from copy import copy
from typing import Generator

from docx.shared import Length, Parented, Pt, Twips

from . import Paragraph
from .caption import Caption, CaptionInfo
from .page_break import PageBreak
from .renderable import Renderable
from .requires_numbering import RequiresNumbering
from ..docx_elements import *
from ..docx_elements import _twips
from ..layout_tracker import LayoutState
from ..profiles import DEFAULT_TABLE_CONTINUATION, TABLE_CONTINUATION_MODES
from ..rendered_info import RenderedInfo

CELL_OFFSET = Pt(9) - Twips(108 * 2)
# Slack only for modes that fragment by our height estimate (legacy/caption).
ROW_HEIGHT_SLACK = Pt(4)

# Modes that do NOT cut the table into fragments — Word owns page breaks
# (word: post-split via COM after save).
_WORD_PAGED_MODES = frozenset({"off", "soft", "word"})


class Table(Renderable, RequiresNumbering):
    def __init__(self, parent: Parented, rows: int, cols: int, caption_info: CaptionInfo):
        super().__init__("Таблица")
        self._parent = parent
        self._caption_info = caption_info
        self._cols = cols
        self._continuation_mode = DEFAULT_TABLE_CONTINUATION

        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"]))

        self._cell_margin_lr = left_margin + right_margin
        self._number = "?"
        self.landscape = bool(caption_info and caption_info.landscape)
        if caption_info and caption_info.unique_name:
            self.unique_name = caption_info.unique_name

        self._rows: list[list[list[Paragraph]]] = [[[] for i in range(cols)] for j in range(rows)]
        self._merge: list[list[tuple[str, str]]] = [
            [("none", "none") for _ in range(cols)] for _ in range(rows)
        ]

    def set_continuation_mode(self, mode: str) -> None:
        if mode not in TABLE_CONTINUATION_MODES:
            raise ValueError(
                f"table continuation must be one of {TABLE_CONTINUATION_MODES}, got {mode!r}"
            )
        self._continuation_mode = mode

    def _table_width(self, layout_state: LayoutState) -> Length:
        return Twips(_twips(layout_state.max_width) + _twips(self._cell_margin_lr))

    def add_paragraph_to_cell(self, row: int, col: int) -> Paragraph:
        paragraph = Paragraph(self._parent)
        try:
            paragraph.style = "Table Text"
        except KeyError:
            pass
        paragraph.first_line_indent = 0
        paragraph._docx_paragraph.paragraph_format.space_before = 0
        paragraph._docx_paragraph.paragraph_format.space_after = 0
        paragraph._docx_paragraph.paragraph_format.line_spacing = 1
        self._rows[row][col].append(paragraph)
        return paragraph

    def set_cell_merge(self, row: int, col: int, merge_v: str = "none", merge_h: str = "none") -> None:
        self._merge[row][col] = (merge_v or "none", merge_h or "none")

    def set_number(self, number: str):
        self._number = number

    def _col_width(self, layout_state: LayoutState) -> Length:
        return self._table_width(layout_state) / self._cols

    def _grid_span(self, merge_row: list[tuple[str, str]], col: int) -> int:
        _, merge_h = merge_row[col]
        if merge_h == "continue":
            return 0
        span = 1
        while col + span < self._cols and merge_row[col + span][1] == "continue":
            span += 1
        return span

    def _build_docx_row(self, docx_table, row_idx: int, apply_merge: bool, layout_state: LayoutState):
        row = self._rows[row_idx]
        merge_row = self._merge[row_idx]
        docx_row = create_table_row(docx_table, header=(row_idx == 0))
        row_height = 0
        col_w = self._col_width(layout_state)

        col = 0
        while col < self._cols:
            merge_v, merge_h = merge_row[col]
            if apply_merge and merge_h == "continue":
                col += 1
                continue

            span = self._grid_span(merge_row, col) if apply_merge else 1
            if span < 1:
                span = 1

            docx_cell = create_table_cell(docx_row, col_w * span)
            v_merge = merge_v if apply_merge and merge_v in ("restart", "continue") else None
            grid_span = span if apply_merge and span > 1 else None
            if v_merge or grid_span:
                apply_cell_merge(docx_cell, v_merge=v_merge, grid_span=grid_span)

            cell_height = 0
            if not (apply_merge and merge_v == "continue"):
                for paragraph in row[col]:
                    # Method guide §3: do not leave caption + column headers alone
                    # at page bottom — glue header row to the following data row.
                    if row_idx == 0:
                        paragraph._docx_paragraph.paragraph_format.keep_with_next = True
                    cell_layout_state = LayoutState(
                        layout_state.max_height, layout_state.max_width
                    )
                    cell_layout_state.max_width = col_w * span - CELL_OFFSET
                    for paragraph_rendered_info in paragraph.render(None, cell_layout_state):
                        docx_cell._element.append(paragraph_rendered_info.docx_element._element)
                        cell_height += paragraph_rendered_info.height
            else:
                from ..util import create_element
                docx_cell._element.append(create_element("w:p"))

            row_height = max(cell_height, row_height)
            docx_row._element.append(docx_cell._element)
            col += span

        row_height += Pt(0.5) + (
            ROW_HEIGHT_SLACK if self._continuation_mode not in _WORD_PAGED_MODES else Pt(0)
        )
        return docx_row, row_height

    def _peek_row_height(self, row_idx: int, layout_state: LayoutState) -> Length:
        """Estimate row height without attaching cell paragraphs to a table."""
        from .paragraph_sizer import ParagraphSizer

        row = self._rows[row_idx]
        merge_row = self._merge[row_idx]
        row_height = 0
        col_w = self._col_width(layout_state)
        col = 0
        while col < self._cols:
            merge_v, merge_h = merge_row[col]
            if merge_h == "continue":
                col += 1
                continue
            span = self._grid_span(merge_row, col)
            if span < 1:
                span = 1
            cell_height = 0
            if merge_v != "continue":
                max_w = col_w * span - CELL_OFFSET
                for paragraph in row[col]:
                    cell_height += ParagraphSizer(
                        paragraph._docx_paragraph, None, max_w
                    ).calculate_height().full
            row_height = max(cell_height, row_height)
            col += span
        row_height += Pt(0.5) + (
            ROW_HEIGHT_SLACK if self._continuation_mode not in _WORD_PAGED_MODES else Pt(0)
        )
        return row_height

    def _start_block_height(
        self, previous_rendered: RenderedInfo | None, layout_state: LayoutState
    ) -> Length:
        """Caption + header + first data row (method: no orphan title/header)."""
        from docx.text.paragraph import Paragraph as DocxParagraph
        from .paragraph_sizer import ParagraphSizer

        # Dry-run caption sizing (temporary Caption; not yielded into the body).
        cap = Caption(
            self._parent, "Таблица", self._caption_info, self._number, True
        )
        prev_p = (
            previous_rendered.docx_element
            if previous_rendered and isinstance(previous_rendered.docx_element, DocxParagraph)
            else None
        )
        total = ParagraphSizer(
            cap._docx_paragraph, prev_p, layout_state.max_width
        ).calculate_height().full
        if self._rows:
            total += self._peek_row_height(0, layout_state)
        if len(self._rows) > 1:
            total += self._peek_row_height(1, layout_state)
        return total

    def _should_split_fragment(
        self, row_height: Length, layout_state: LayoutState, rows_in_fragment: int
    ) -> bool:
        # off/soft: never fragment — height estimate ≠ Word; mid-page «Продолжение» bug
        if self._continuation_mode in _WORD_PAGED_MODES:
            return False
        if rows_in_fragment == 0:
            return False
        return row_height > layout_state.remaining_page_height

    def _make_continuation_paragraph(self) -> Paragraph:
        continuation_paragraph = Paragraph(self._parent)
        continuation_paragraph.add_run(f"Продолжение Таблицы {self._number}")
        continuation_paragraph.style = "Название таблицы"
        continuation_paragraph.first_line_indent = 0
        return continuation_paragraph

    def _emit_page_break_and_optional_caption(
        self, layout_state: LayoutState
    ) -> Generator[RenderedInfo, None, None]:
        """For legacy/caption only: break page and insert «Продолжение Таблицы N»."""
        page_break_info = next(PageBreak(self._parent).render(None, layout_state))
        layout_state.add_height(page_break_info.height)
        yield page_break_info

        # Never page_break_before on «Название таблицы»: keep_with_next +
        # pageBreakBefore makes a blank page that debug does not count.
        continuation_paragraph = self._make_continuation_paragraph()
        continuation_paragraph._docx_paragraph.paragraph_format.keep_with_next = True
        continuation_paragraph._docx_paragraph.paragraph_format.page_break_before = False
        info = next(continuation_paragraph.render(None, copy(layout_state)))
        layout_state.add_height(info.height)
        yield info

    def render(self, previous_rendered: RenderedInfo, layout_state: LayoutState) \
            -> Generator[RenderedInfo, None, None]:
        # Method guide: never leave only «Таблица N» or «Таблица N» + column
        # headers at the bottom of a page. If caption+header+first data row do
        # not fit, push the whole block to the next page (layout leftover;
        # Word also gets keep_with_next on caption/header).
        leftover = Pt(0)
        start_h = self._start_block_height(previous_rendered, layout_state)
        if (
            layout_state.current_page_height > 0
            and start_h > layout_state.remaining_page_height
        ):
            leftover = layout_state.remaining_page_height
            layout_state.add_height(leftover)
            previous_rendered = None

        caption_rendered_infos = list(
            Caption(self._parent, "Таблица", self._caption_info, self._number, True)
            .render(previous_rendered, copy(layout_state))
        )
        if leftover and caption_rendered_infos:
            first = caption_rendered_infos[0]
            caption_rendered_infos[0] = RenderedInfo(
                first.docx_element, first.height + leftover
            )
        layout_state.add_height(sum(info.height for info in caption_rendered_infos))
        yield from caption_rendered_infos

        docx_table = create_table(
            self._parent, 0, self._cols, self._table_width(layout_state)
        )

        table_height = Pt(0.5)
        apply_merge = True
        rows_in_fragment = 0

        for row_idx in range(len(self._rows)):
            docx_row, row_height = self._build_docx_row(
                docx_table, row_idx, apply_merge, layout_state
            )

            if self._should_split_fragment(row_height, layout_state, rows_in_fragment):
                # Never cut after the header alone (method: no title+column headers
                # without a data row). Prefer slight overflow over an orphan fragment.
                if rows_in_fragment >= 2:
                    yield RenderedInfo(docx_table, table_height)
                    yield from self._emit_page_break_and_optional_caption(layout_state)

                    docx_table = create_table(
                        self._parent, 0, self._cols, self._table_width(layout_state)
                    )
                    apply_merge = False
                    rows_in_fragment = 0
                    table_height = Pt(0.5)

                    docx_row, row_height = self._build_docx_row(
                        docx_table, row_idx, apply_merge, layout_state
                    )

            docx_table._element.append(docx_row._element)
            layout_state.add_height(row_height)
            table_height += row_height
            rows_in_fragment += 1

        yield RenderedInfo(docx_table, table_height)