""" DOCX file configuration provider for doc_file_processing """

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from functools import cached_property
from typing import Any, TYPE_CHECKING, Generator, Iterator, Optional, Union
from xml.dom import minidom
from zipfile import ZipFile
if TYPE_CHECKING: from _typeshed import StrPath

from document import *


class DOCXReader(Reader):
    """ Configuration reader implementation for DOCX files"""

    _doc: _DOCX

    def __init__(self, path: StrPath) -> None:
        self._doc = _DOCX(path)

    def parse(self) -> Document:
        return _Parser(self._doc).parse_document()
    
    def write_image(self, key: Any, dst: StrPath) -> None:
        self._doc.extract_picture(key, dst)
        
    def timestamp(self) -> datetime:
        return self._doc.ctime()

    def is_config(self) -> bool:
        return self._doc.template_name().startswith('ABB_ProjectDoc')


class _Parser:
    _doc: _DOCX
    _stream: Iterator[RawContent]
    _next_buf: Optional[RawContent]

    def __init__(self, doc: _DOCX) -> None:
        self._doc = doc
        self._stream = doc.read_document()
        self._next_buf = None

    @property
    def _next(self) -> Optional[RawContent]:
        if self._next_buf is None:
            self._next_buf = next(self._stream, None)
        return self._next_buf
    
    def _advance(self) -> Optional[RawContent]:
        """ Advance the parser and return the token advanced past """
        c = self._next
        self._next_buf = None
        return c
    
    def parse_document(self) -> Document:
        """ Parse the document. """
        content = list(self.parse_content())
        assert self._next is None
        return Document(content)
    
    def parse_content(self, heading_limit: int = 0) -> Generator[Content, None, None]:
        """ Parse a stream of content (tables, paragraphs, figures, and/or subchapters).
        
        `heading_limit` sets the current heading level -- deeper headings will initiate a
        subchapter in the content stream, while higher level headings will end the content
        stream. """
        while True:
            if isinstance(self._next, RawText):
                text = self.parse_plain_text().strip()
                if text: yield TextBlock(text)
            elif isinstance(self._next, Marker) and self._next == Marker.TABLE_BEGIN:
                data = self.parse_table_body()
                index, caption = self.parse_caption()
                yield Table(index, caption, data[0], data[1:])
            elif isinstance(self._next, HeadingBegin) and (lvl := self._next.level) >= heading_limit:
                self._advance()
                title = self.parse_plain_text().strip()
                assert self._advance() == Marker.HEADING_END
                content = list(self.parse_content(heading_limit=lvl+1))
                yield Chapter(title, lvl, content)
            elif isinstance(self._next, RawFigure):
                key = self._next.key
                self._advance()
                index, caption = self.parse_caption()
                yield Figure(index, caption, key)
            else:
                return
                
    def parse_plain_text(self) -> str:
        """ Parse and join a run of text. """
        text = ''
        while isinstance(self._next, RawText):
            text += self._next.text
            self._advance()
        return text
    
    def parse_caption(self) -> tuple[Optional[int], str]:
        """ Parse a figure/table caption, returning index and name. """
        caption_text = ''
        if isinstance(self._next, RawText) and self._next.text == '\n': self._advance()
        while isinstance(self._next, RawText) and not self._next.text.endswith('\n'):
            caption_text += self._next.text
            self._advance()
        
        if match := CAPTION_TEXT_RE.match(caption_text):
            return int(match['index']), match['name']
        else:
            return None, caption_text
        
    def parse_table_body(self) -> tuple[tuple[str, ...], ...]:
        """ Parse the contents of a table. """
        def _row() -> Generator[str, None, None]:
            while self._next == Marker.CELL_BEGIN:
                self._advance()
                yield self.parse_plain_text().strip('\n')
                assert self._advance() == Marker.CELL_END

        def _rows() -> Generator[tuple[str, ...], None, None]:
            while self._next == Marker.ROW_BEGIN:
                self._advance()
                yield tuple(_row())
                assert self._advance() == Marker.ROW_END

        assert self._advance() == Marker.TABLE_BEGIN
        body = tuple(_rows())
        assert self._advance() == Marker.TABLE_END

        return body
        

def _xml_lookup(root: minidom.Document | minidom.Element, *path: str) -> Optional[minidom.Node]:
    """ Look up a descendent node by a path consisting of tag names """
    node: minidom.Node = root
    for term in path:
        if not (isinstance(node, minidom.Element) or isinstance(node, minidom.Document)): return None
        nodes = node.getElementsByTagName(term)
        if nodes: node = nodes[0]
        else: return None
    return node

@dataclass
class RawFigure:
    key: str

@dataclass
class RawText:
    text: str

@dataclass
class HeadingBegin:
    level: int

class Marker(Enum):
    TABLE_BEGIN = 0
    TABLE_END = 1
    ROW_BEGIN = 2
    ROW_END = 3
    CELL_BEGIN = 4
    CELL_END = 5
    HEADING_END = 7

RawContent = Union[RawText, RawFigure, HeadingBegin, Marker]

class _DOCX:
    _arc: ZipFile
    _data: minidom.Document

    def __init__(self, path: StrPath) -> None:
        self._arc = ZipFile(path)
        with self._arc.open('word/document.xml') as data_xml:
            self._data = minidom.parse(data_xml)

    def _style_outline_level(self, istyle: str) -> Optional[int]:
        """ Get outline level of style identified by `istyle` """
        return self._style_table.get(istyle, None)

    def extract_picture(self, key: str, dst: StrPath) -> None:
        """ Write embedded media file identified by `key` to path `dst`. """
        path = 'word/' + self._rel_table[key]
        with self._arc.open(path) as data, open(dst, 'wb') as outfile:
            outfile.write(data.read())

    def read_document(self) -> Generator[RawContent, None, None]:
        """ Read and tokenize the document. """
        root = _xml_lookup(self._data, 'w:document', 'w:body')
        assert root is not None
        return self.read_content(root)

    def read_content(self, root: minidom.Node) -> Generator[RawContent, None, None]:
        """ Parse and yield content tokens from the descendents of node `root` in the document. """
        for child in root.childNodes:
            child: minidom.Node
            if isinstance(child, minidom.Element):
                if child.nodeName in ['w:p', 'w:r', 'w:t', 'w:fldSimple']:
                    is_heading = False
                    if child.nodeName == 'w:p':
                        if isinstance(style := _xml_lookup(child, 'w:pPr', 'w:pStyle'), minidom.Element):
                            lvl = self._style_outline_level(style.getAttribute('w:val'))
                            if lvl is not None:
                                is_heading = True
                                yield HeadingBegin(lvl)

                    yield from self.read_content(child)

                    if child.nodeName == 'w:p':
                        yield RawText('\n')
                    if is_heading:
                        yield Marker.HEADING_END
                elif child.nodeName == 'w:tbl':
                    yield from self.read_table(child)
                elif child.nodeName == 'w:drawing':
                    yield self.read_figure(child)
            elif isinstance(child, minidom.Text):
                yield RawText(child.data)

    def read_figure(self, node: minidom.Element) -> RawFigure:
        """ Get information for an embedded figure at `node`. """
        blip = _xml_lookup(node, 'wp:inline', 'a:graphic', 'a:graphicData', 'pic:pic', 'pic:blipFill', 'a:blip')
        assert isinstance(blip, minidom.Element)
        key = blip.getAttribute('r:embed')
        return RawFigure(key)
    
    def read_table(self, node: minidom.Node) -> Generator[RawContent, None, None]:
        """ Parse and yield content tokens for a table at `node`. """
        yield Marker.TABLE_BEGIN
        for row in node.childNodes:
            if isinstance(row, minidom.Element) and row.nodeName == 'w:tr':
                yield Marker.ROW_BEGIN
                for cell in row.childNodes:
                    if isinstance(cell, minidom.Element) and cell.nodeName == 'w:tc':
                        yield Marker.CELL_BEGIN
                        yield from self.read_content(cell)
                        yield Marker.CELL_END
                yield Marker.ROW_END
        yield Marker.TABLE_END

    def ctime(self) -> datetime:
        """ Get the document creation time. """
        with self._arc.open('docProps/core.xml') as props_xml:
            props = minidom.parse(props_xml)
            created = _xml_lookup(props, 'cp:coreProperties', 'dcterms:created')
            assert created is not None
            text = created.childNodes[0]
            assert isinstance(text, minidom.Text)
            return datetime.strptime(text.data, '%Y-%m-%dT%H:%M:%SZ')
        
    def template_name(self) -> str:
        """ Get the document template name. """
        with self._arc.open('docProps/app.xml') as props_xml:
            props = minidom.parse(props_xml)
            node = _xml_lookup(props, 'Properties', 'Template')
            assert node is not None
            text = node.childNodes[0]
            assert isinstance(text, minidom.Text)
            return text.data

    @cached_property
    def _style_table(self) -> dict[str, int]:
        """ Document style definitions.
        
        The only information stored for each style is its outline level, as an int. """

        def _gen(root: minidom.Node) -> Generator[tuple[str, int], None, None]:
            for node in root.childNodes:
                if isinstance(node, minidom.Element) and node.nodeName == 'w:style':
                    istyle = node.getAttribute('w:styleId')
                    lvl = _xml_lookup(node, 'w:pPr', 'w:outlineLvl')
                    if isinstance(lvl, minidom.Element):
                        yield istyle, int(lvl.getAttribute('w:val'))
        
        with self._arc.open('word/styles.xml') as styles_xml:
            styles = minidom.parse(styles_xml)
        root = _xml_lookup(styles, 'w:styles')
        assert root is not None
        return dict(_gen(root))
    
    @cached_property
    def _rel_table(self) -> dict[str, str]:
        """ Document relationship definitions. """
        def _gen(root: minidom.Node) -> Generator[tuple[str, str], None, None]:
            for node in root.childNodes:
                if isinstance(node, minidom.Element) and node.nodeName == 'Relationship':
                    yield node.getAttribute('Id'), node.getAttribute('Target')

        with self._arc.open('word/_rels/document.xml.rels') as rels_xml:
            rels = minidom.parse(rels_xml)
        root = _xml_lookup(rels, 'Relationships')
        assert root is not None
        return dict(_gen(root))
