from typing import Generator

from docx.enum.table import WD_CELL_VERTICAL_ALIGNMENT
from docx.oxml import CT_Tbl
from docx.shared import Pt, Twips
from docx.table import Table
from docx.text.paragraph import Paragraph as DocxParagraph

from .requires_numbering import RequiresNumbering
from ..layout_tracker import LayoutState
from ..renderable import Renderable
from ..rendered_info import RenderedInfo
from ..util import create_element
from ..latex_math import latex_to_omml


_HEIGHT = Pt(50)


class Equation(Renderable, RequiresNumbering):
    def __init__(self, parent, latex_formula: str, unique_name: str | None = None,
                 numbered: bool = False):
        super().__init__("Формула")
        self.unique_name = unique_name
        self._numbered = numbered
        word_math = latex_to_omml(latex_formula)

        sect = parent.part.document.sections[0]

        left_margin = Twips(int(
            parent.part.styles["Normal Table"]._element.xpath("w:tblPr/w:tblCellMar/w:left")[0].attrib[
                "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}w"]))
        right_margin = Twips(int(
            parent.part.styles["Normal Table"]._element.xpath("w:tblPr/w:tblCellMar/w:right")[0].attrib[
                "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}w"]))

        table_width = sect.page_width - sect.right_margin - sect.left_margin + left_margin + right_margin

        self._parent = parent
        self._table = table = Table(CT_Tbl.new_tbl(1, 2, table_width), parent)

        left_cell = table.cell(0, 0)
        right_cell = table.cell(0, 1)
        right_cell.width = Pt(40)
        left_cell.width = table_width - right_cell.width

        table.rows[0].height = _HEIGHT

        left_paragraph = left_cell.paragraphs[0]
        left_paragraph.style = "Formula Content"
        left_paragraph._p.append(word_math)
        left_cell.vertical_alignment = WD_CELL_VERTICAL_ALIGNMENT.CENTER

        self._right_paragraph = right_cell.paragraphs[0]
        self._right_paragraph.style = "Formula Numbering"
        self._number_text = "?"
        right_cell.vertical_alignment = WD_CELL_VERTICAL_ALIGNMENT.CENTER

        self._spacer_before = create_element("w:p")
        self._spacer_after = create_element("w:p")

    def enable_numbering(self):
        self._numbered = True

    @property
    def needs_numbering(self) -> bool:
        return self._numbered

    def set_number(self, number: str):
        self._number_text = str(number)
        self._numbered = True
        # Rebuild right cell content
        p = self._right_paragraph._p
        for child in list(p):
            if child.tag.endswith("}r"):
                p.remove(child)
        p.append(create_element("w:r", "("))
        p.append(create_element("w:r", self._number_text))
        p.append(create_element("w:r", ")"))

    def render(self, previous_rendered: RenderedInfo, layout_state: LayoutState) -> Generator[
            "RenderedInfo | Renderable", None, None]:
        spacer_h = Pt(21)  # ~ one empty line at 14pt * 1.5
        height = _HEIGHT + spacer_h * 2

        if height > layout_state.remaining_page_height:
            height += layout_state.remaining_page_height

        yield RenderedInfo(DocxParagraph(self._spacer_before, self._parent), spacer_h)
        yield RenderedInfo(self._table, _HEIGHT)
        yield RenderedInfo(DocxParagraph(self._spacer_after, self._parent), spacer_h)