""" Selective read implementation of the binary DOC file format """

from __future__ import annotations

import zlib
from bisect import bisect
from dataclasses import dataclass
from datetime import datetime
from functools import cached_property, lru_cache
from io import SEEK_CUR, BytesIO, IOBase
from typing import TYPE_CHECKING, Generator, Literal
if TYPE_CHECKING:
    from typing_extensions import Self
    from _typeshed import StrPath

from winstg import Storage, Stream, open_storage
from logutil import errmsg


def _prop(s: str) -> int:
    gr, pr = (int(i, base=16) for i in s.split(':'))
    return pr + (gr << 10)

class DOCFile:
    PROP_PARA_HEADING_LEVEL = _prop('1:0a')
    PROP_PARA_IN_TABLE = _prop('1:16')
    PROP_PARA_END_OF_ROW = _prop('1:17')
    PROP_PARA_INDIRECT = _prop('1:46')
    PROP_CHAR_EMBED = _prop('2:03')

    _MAGIC = 0xA5EC
    _LW_OFFS = 64
    _LW_CCTEXT_IDX = 3
    _LCB_OFFS = 154
    _LCB_STSH_IDX = 2
    _LCB_CHPX_IDX = 24
    _LCB_PAPX_IDX = 26
    _LCB_DOP_IDX = 62
    _LCB_CLX_IDX = 66
    _LCB_FT_IDX = 174

    _GUID_SUMMARYINFORMATION = "{F29F85E0-4FF9-1068-AB91-08002B27B3D9}"
    _SI_TEMPLATE_KEY = 0x07
    _SI_CTIME_KEY = 0x0C
    _SI_MTIME_KEY = 0x0D

    _stg: Storage
    _prop_cache: dict[int, bytes]
    _prop_cache_left: int
    _prop_cache_right: int
    _cur_fkp_left: int
    _cur_fkp_right: int


    def __init__(self, path: StrPath) -> None:
        self._stg = open_storage(path)
        self._prop_cache = {}
        self._prop_cache_key = -1
        self._cur_fkp_left = -1
        self._cur_fkp_right = -1

    def read(self, count: int) -> str:
        """ Read from the document text stream. """
        if self._text_compressed:
            data = self._textstream.read(count)
            text = str(data, encoding='ansi')
        else:
            data = self._textstream.read(count * 2)
            text = str(data, encoding='utf-16')
        return text

    def get_pprop_at(self, pos: int, propid: int) -> int | None:
        """ Read a formatting property at the specified position. """
        if pos != self._prop_cache_key:
            self._prop_cache_key = pos
            self._prop_cache = {prop.key: prop.value for prop in self._read_props(pos, 'papx')}
        data = self._prop_cache.get(propid)
        r = None if data is None \
            else int.from_bytes(data, 'little')
        return r
    
    def get_cprop_at(self, pos: int, propid: int) -> int | None:
        """ Read a character property at the specified position. """
        for prop in self._read_props(pos, 'chpx'):
            if prop.key == propid:
                return int.from_bytes(prop.value, 'little')
        return None
    
    def get_picture(self, key: int) -> Picture:
        pic = self._open_datastream(key)
        return Picture.unpack_from_stream(pic)


    @cached_property
    def _doc_info(self) -> _DocInfo:
        with self._open_docstream() as fib:
            magic = int.from_bytes(fib.read(2), 'little')
            if magic != self._MAGIC: raise RuntimeError("Bad magic number")
            fib.seek(9, SEEK_CUR)
            tbl_idx = (int.from_bytes(fib.read(1), 'little') & 0x02) >> 1
        with self._open_docstream(self._LW_OFFS + 4 * self._LW_CCTEXT_IDX) as lw:
            cc_text = int.from_bytes(lw.read(4), 'little')
        return _DocInfo(tbl_idx, cc_text)

    @property
    def len_text(self) -> int:
        return self._doc_info.cc_text

    @cached_property
    def properties(self) -> DOCProperties:
        propset_stm = self._stg['\005SummaryInformation']
        assert isinstance(propset_stm, Stream)

        propset = propset_stm.open_propset(self._GUID_SUMMARYINFORMATION)
        ctime = propset[self._SI_CTIME_KEY]
        mtime = propset[self._SI_MTIME_KEY]
        try:
            template = propset[self._SI_TEMPLATE_KEY]
        except IndexError:
            template = ""

        assert isinstance(ctime, datetime)
        assert isinstance(mtime, datetime)
        assert isinstance(template, str)
        return DOCProperties(ctime, mtime, template)

    @cached_property
    def _tabstream_key(self) -> str:
        idx = self._doc_info.tbl_idx
        return f'{idx}Table'
    
    @cached_property
    def _docstream(self) -> Stream:
        stm = self._stg['WordDocument']
        assert isinstance(stm, Stream)
        return stm
    
    @cached_property
    def _datastream(self) -> Stream:
        stm = self._stg['Data']
        assert isinstance(stm, Stream)
        return stm
    
    @cached_property
    def _tabstream(self) -> Stream:
        stm = self._stg[self._tabstream_key]
        assert isinstance(stm, Stream)
        return stm
    
    @cached_property
    def _textstream(self) -> Stream:
        if self._text_compressed:
            offs = self._range_table.main_range_offset // 2
        else:
            offs = self._range_table.main_range_offset
        return self._open_docstream(offs)

    @cached_property
    def _text_compressed(self) -> bool:
        return self._range_table.main_range_compressed
    
    @cached_property
    def _range_table(self) -> _RangeTable:
        with self._open_docstream(self._LCB_OFFS + 4 * self._LCB_CLX_IDX) as lcb:
            clx_offs = int.from_bytes(lcb.read(4), 'little')
        with self._open_tabstream(clx_offs) as clx:
            while True:
                dc = int.from_bytes(clx.read(1), 'little')
                if dc == 1:
                    pc_sz = int.from_bytes(clx.read(2), 'little')
                    clx.seek(pc_sz, SEEK_CUR)
                elif dc == 2:
                    break
                else:
                    raise RuntimeError(f"Failed to parse document CLX: bad discriminant {dc}")
            rg_sz = int.from_bytes(clx.read(4), 'little')
            rg_count = (rg_sz - 4) // 12
            if rg_count > 1: raise RuntimeError("Unsupported multirange document")
            rg_start = int.from_bytes(clx.read(4), 'little')
            rg_end = int.from_bytes(clx.read(4), 'little')

            clx.seek(2, SEEK_CUR)
            fc = int.from_bytes(clx.read(4), 'little')
            offs = fc & ~0xC000_0000
            is_compressed = bool(fc & 0x4000_0000)
        return _RangeTable(rg_start, rg_end, offs, is_compressed)
    
    def _read_bte(self, which: Literal['papx', 'chpx']) -> Generator[tuple[int, int, bytes], None, None]:
        if   which == 'papx': idx = self._LCB_PAPX_IDX
        elif which == 'chpx': idx = self._LCB_CHPX_IDX
        else: raise ValueError()

        with self._open_docstream(self._LCB_OFFS + 4 * idx) as lcb:
            bte_offs = int.from_bytes(lcb.read(4), 'little')
            bte_sz = int.from_bytes(lcb.read(4), 'little')
        metaspan_count = (bte_sz - 4) // 8
        with self._open_tabstream(bte_offs) as fkpd, self._open_docstream() as docstream:
            metaspan_bounds = [int.from_bytes(fkpd.read(4), 'little') \
                for _ in range(metaspan_count + 1)]
            for i in range(metaspan_count):
                fkp_offs = int.from_bytes(fkpd.read(4), 'little') * 512
                docstream.seek(fkp_offs)
                fkp = docstream.read(512)
                yield (metaspan_bounds[i], metaspan_bounds[i+1], fkp)

    @cached_property
    def _fkp_directory_papx(self) -> list[tuple[int, int, bytes]]:
        return list(self._read_bte('papx'))
    
    @cached_property
    def _fkp_directory_chpx(self) -> list[tuple[int, int, bytes]]:
        return list(self._read_bte('chpx'))

    @cached_property
    def _style_table(self) -> _StyleTable:
        with self._open_docstream(self._LCB_OFFS + 4 * self._LCB_STSH_IDX) as lcb:
            stshOffs = int.from_bytes(lcb.read(4), 'little')
        with self._open_tabstream(stshOffs) as stsh:
            return _StyleTable.unpack_from_stream(stsh)
        
    def _select_papx_fkp(self, choffs: int) -> bytes:
        if not (self._cur_fkp_left <= choffs < self._cur_fkp_right):
            i = bisect([t[1] for t in self._fkp_directory_papx], choffs)
            if i in range(len(self._fkp_directory_papx)):
                self._cur_fkp_left, self._cur_fkp_right, self._cur_fkp = self._fkp_directory_papx[i]
            else: raise IndexError(f"No papx metaspan for ch offs {choffs}")
        return self._cur_fkp
    
    def _select_chpx_fkp(self, choffs: int) -> bytes:
        i = bisect([t[1] for t in self._fkp_directory_chpx], choffs)
        if i in range(len(self._fkp_directory_papx)):
            return self._fkp_directory_chpx[i][2]
        else: raise IndexError(f"No chpx metaspan for ch offs {choffs}")

    def _read_props(self, pos: int, which: Literal['papx', 'chpx']) -> Generator[Property, None, None]:
        choffs = self._range_table.main_range_offset \
            // (2 if self._range_table.main_range_compressed else 1) \
            + (pos - self._range_table.main_range_start) \
            * (1 if self._range_table.main_range_compressed else 2)
        
        if   which == 'papx': fkp = self._select_papx_fkp(choffs)
        elif which == 'chpx': fkp = self._select_chpx_fkp(choffs)
        else: raise ValueError()

        span_count = fkp[-1]
        for i in range(span_count):
            span_left = int.from_bytes(fkp[4*i:4*i+4], 'little')
            if span_left > choffs: break
        else: i = span_count
        i -= 1
        if i < 0: raise IndexError(f"Failed to find span containing ch offs {choffs}")

        if which == 'papx':
            papx_fkp_offs = 4 * (span_count + 1) + 13 * i
            papx_offs = fkp[papx_fkp_offs] * 2
            papx_sz = fkp[papx_offs]
            if papx_sz == 0:
                papx_sz = fkp[papx_offs+1] * 2
                props = fkp[papx_offs+2:papx_offs+2+papx_sz]
            else:
                papx_sz = 2 * papx_sz - 1
                props = fkp[papx_offs+1:papx_offs+1+papx_sz]
            istyle = int.from_bytes(props[0:2], 'little')
            props = props[2:]
        elif which == 'chpx':
            chpx_offs_offs = 4 * (span_count + 1) + i
            if chpx_offs_offs:
                chpx_offs = fkp[chpx_offs_offs] * 2
                chpx_sz = fkp[chpx_offs]
                props = fkp[chpx_offs+1:chpx_offs+1+chpx_sz]
            else:
                props = None
            istyle = None

        if istyle is not None:
            yield from self._style_table.all_props(istyle)

        if props is not None:
            props_stm = BytesIO(props)
            while True:
                try:
                    prop = Property.unpack_from_stream(props_stm)
                except EOFError:
                    break
                else:
                    if prop.key == self.PROP_PARA_INDIRECT:
                        yield from self._read_extra_props(int.from_bytes(prop.value, 'little'))
                    else:
                        yield prop

    def _read_extra_props(self, offs: int) -> Generator[Property, None, None]:
        with self._open_datastream(offs) as prc:
            prc_sz = int.from_bytes(prc.read(2), 'little')
            prc_base = prc.tell()
            while True:
                yield Property.unpack_from_stream(prc)
                if prc.tell() - prc_base >= prc_sz: break

    def _open_docstream(self, offs: int = 0) -> Stream:
        stm = self._docstream.clone()
        stm.seek(offs)
        return stm
    
    def _open_tabstream(self, offs: int = 0) -> Stream:
        stm = self._tabstream.clone()
        stm.seek(offs)
        return stm
    
    def _open_datastream(self, offs: int = 0) -> Stream:
        stm = self._datastream.clone()
        stm.seek(offs)
        return stm
    
@dataclass
class _RangeTable:
    main_range_start: int
    main_range_end: int
    main_range_offset: int
    main_range_compressed: bool

@dataclass
class _DocInfo:
    tbl_idx: int
    cc_text: int

@dataclass
class DOCProperties:
    ctime: datetime
    mtime: datetime
    template_name: str

class _StyleTable:
    _table: dict[int, Style]
    _resolved: dict[int, list[Property]]

    def __init__(self, table: dict[int, Style]) -> None:
        self._table = table
        self._resolved = {}

    @classmethod
    def unpack_from_stream(cls, stsh: IOBase) -> Self:
        styles = {}

        stshi_sz = int.from_bytes(stsh.read(2), 'little')
        style_count = int.from_bytes(stsh.read(2), 'little')
        stdf_sz = int.from_bytes(stsh.read(2), 'little')
        stsh.seek(stshi_sz - 4, SEEK_CUR)

        for i in range(style_count):
            std_sz = int.from_bytes(stsh.read(2), 'little')
            if std_sz == 0: continue

            std_base = stsh.tell()
            stsh.seek(std_base)
            istyle = int.from_bytes(stsh.read(2), 'little') & 0x0fff
            iparent_stk = int.from_bytes(stsh.read(2), 'little')
            iparent: int | None = (iparent_stk & 0xfff0) >> 4
            if iparent == 0x0fff: iparent = None
            stk = iparent_stk & 0x000f
            stsh.seek(6, SEEK_CUR)
            if (stdf_sz > 10): stsh.seek(stdf_sz - 10, SEEK_CUR)
            name_sz = int.from_bytes(stsh.read(2), 'little')
            name_bytes = stsh.read(name_sz * 2)
            name = str(name_bytes, encoding='utf-16')
            stsh.seek(2, SEEK_CUR)

            props: list[Property] = []

            if stk == 3:
                tapx_sz = int.from_bytes(stsh.read(2), 'little')
                stsh.seek(2, SEEK_CUR)
                tapx = BytesIO(stsh.read(tapx_sz - 2))
                while True:
                    try:
                        props.append(Property.unpack_from_stream(tapx))
                    except EOFError:
                        break
                if tapx_sz % 2: stsh.seek(1, SEEK_CUR)

            if stk in [1, 4]:
                papx_sz = int.from_bytes(stsh.read(2), 'little')
                aistyle = int.from_bytes(stsh.read(2), 'little')
                if aistyle != 0x0FFE: istyle = aistyle
                papx = BytesIO(stsh.read(papx_sz - 2))
                while True:
                    try:
                        props.append(Property.unpack_from_stream(papx))
                    except EOFError:
                        break
                if papx_sz % 2: stsh.seek(1, SEEK_CUR)
                
            if stk in [1, 2, 3]:
                chpx_sz = int.from_bytes(stsh.read(2), 'little')
                chpx = BytesIO(stsh.read(chpx_sz))
                while True:
                    try:
                        props.append(Property.unpack_from_stream(chpx))
                    except EOFError:
                        break
                if chpx_sz % 2: stsh.seek(1, SEEK_CUR)

            if stk not in [1, 2, 3, 4]:
                raise RuntimeError()

            stsh.seek(std_base + std_sz)
            styles[istyle] = Style(props, name, iparent)

        return cls(styles)

    def own_props(self, istyle: int) -> list[Property]:
        return self._table[istyle].properties
    
    @lru_cache(maxsize=None)
    def all_props(self, istyle: int) -> list[Property]:
        def _gen() -> Generator[Property, None, None]:
            iparent  = self._table[istyle].iparent
            if iparent is not None:
                yield from self.all_props(iparent)
            yield from self.own_props(istyle)
        
        
        if istyle not in self._table:
            errmsg(f"DOC: missing definition for style ID {istyle}")
            return []

        return list(_gen())

@dataclass
class Style:
    properties: list[Property]
    name: str
    iparent: int | None

@dataclass
class Property:
    key: int
    value: bytes

    @classmethod
    def unpack_from_stream(cls, sprm: IOBase) -> Self:
        op = int.from_bytes(sprm.read(2), 'little')
        if op == 0: raise EOFError()
        op_pr =  (op & 0x01FF) >> 0
        op_prf = (op & 0x0200) >> 9
        op_gr =  (op & 0x1C00) >> 10
        op_ar =  (op & 0xE000) >> 13

        if   op_ar in [0, 1]:    arg_sz = 1
        elif op_ar in [2, 4, 5]: arg_sz = 2
        elif op_ar == 3:         arg_sz = 4
        elif op_ar == 7:         arg_sz = 3
        elif op_ar == 6:         arg_sz = 0
        else: raise RuntimeError()

        if arg_sz == 0 and op_pr == 0x08 and op_gr == 0x05:
            arg_sz = int.from_bytes(sprm.read(2), 'little') - 1
        elif arg_sz == 0 and op_pr == 0x15 and op_gr == 0x01:
            raise NotImplementedError("Tab list")
        elif arg_sz == 0:
            arg_sz = int.from_bytes(sprm.read(1), 'little')
            if arg_sz == 0: raise RuntimeError()
        value = sprm.read(arg_sz)

        key = op & 0x1DFF
        return cls(key, value)

@dataclass
class Picture:
    data: bytes

    @classmethod
    def unpack_from_stream(cls, pic: IOBase) -> Self:
        pic_base = pic.tell()
        pic_sz = int.from_bytes(pic.read(4), 'little')
        pic.seek(2, SEEK_CUR)
        pic_type = int.from_bytes(pic.read(2), 'little')
        pic.seek(60, SEEK_CUR)
        if pic_type != 0x64:
            raise RuntimeError('Unsupported picture type')
        
        while pic.tell() - pic_base < pic_sz:
            rec_vi = int.from_bytes(pic.read(2), 'little')
            rec_type = int.from_bytes(pic.read(2), 'little')
            rec_sz = int.from_bytes(pic.read(4), 'little')
            if rec_type == 0xF007:
                rec_btype = (rec_vi & 0xfff0) >> 4
                if rec_btype != 0x002: break

                pic.seek(20, SEEK_CUR)
                inrec_sz = int.from_bytes(pic.read(4), 'little')
                pic.seek(9, SEEK_CUR)
                name_sz = int.from_bytes(pic.read(1), 'little')
                pic.seek(2, SEEK_CUR)
                name_data = pic.read(name_sz)
                name = str(name_data, encoding='utf-16')

                inrec_vi = int.from_bytes(pic.read(2), 'little')
                inrec_type = int.from_bytes(pic.read(2), 'little')
                inrec_sz = int.from_bytes(pic.read(4), 'little')
                
                inrec_btype = (inrec_vi & 0xfff0) >> 4
                if inrec_btype == 0x3D4:
                    pic.seek(16, SEEK_CUR)
                elif inrec_btype == 0x3D5:
                    pic.seek(16, SEEK_CUR)
                else:
                    raise RuntimeError()
                
                emf_orig_sz = int.from_bytes(pic.read(4), 'little')
                pic.seek(24, SEEK_CUR)
                emf_sz = int.from_bytes(pic.read(4), 'little')
                emf_compression_type = int.from_bytes(pic.read(1), 'little')
                pic.seek(1, SEEK_CUR)

                data = pic.read(emf_sz)
                if emf_compression_type == 0x00:
                    data = zlib.decompress(data)
                elif emf_compression_type == 0xFE:
                    pass
                else:
                    raise RuntimeError()
                
                return Picture(data)
            else:
                pic.seek(rec_sz, SEEK_CUR)
        raise RuntimeError('Unsupported picture type')
