"""
DOCX parser — uses python-docx for structure-aware extraction.

Extracts heading styles, tables (dual-format), inline images, and paragraphs.
"""

from __future__ import annotations

import re
from collections.abc import Iterator

from .base import BaseParser, FileRecord, RawElement


def _heading_level_from_style(style_name: str) -> int | None:
    """Extract heading level from DOCX style name like 'Heading 1', 'Heading 2'."""
    m = re.match(r"[Hh]eading\s*(\d)", style_name)
    if m:
        return min(int(m.group(1)), 6)
    return None


def _extract_merge_map(table) -> list[dict]:
    """Extract merge topology from a python-docx table. Returns list of MergedCell dicts."""
    from docx.oxml.ns import qn

    merge_map: list[dict] = []
    seen: set = set()
    for row_i, row in enumerate(table.rows):
        for col_i, cell in enumerate(row.cells):
            tc = cell._tc
            cid = id(tc)
            if cid in seen:
                continue
            seen.add(cid)
            vm = tc.find(qn("w:vMerge"))
            if vm is not None and vm.get(qn("w:val")) != "restart":
                continue  # vertical continuation cell — skip
            tcPr = tc.find(qn("w:tcPr"))
            gs_el = tcPr.find(qn("w:gridSpan")) if tcPr is not None else None
            colspan = int(gs_el.get(qn("w:val")) or 1) if gs_el is not None else 1
            rowspan = 1
            r = row_i + 1
            while r < len(table.rows):
                ntc = table.rows[r].cells[col_i]._tc
                nvm = ntc.find(qn("w:vMerge"))
                if nvm is not None and nvm.get(qn("w:val")) != "restart":
                    rowspan += 1
                    r += 1
                else:
                    break
            merge_map.append(
                {
                    "row": row_i,
                    "col": col_i,
                    "rowspan": rowspan,
                    "colspan": colspan,
                    "value": cell.text.strip(),
                }
            )
    return merge_map


def _build_html_table(merge_map: list[dict], col_count: int) -> str:
    """Build an HTML table string with proper rowspan/colspan from merge_map."""
    # Group cells by row
    rows: dict[int, list[dict]] = {}
    for cell in merge_map:
        rows.setdefault(cell["row"], []).append(cell)

    lines = ["<table>"]
    for row_i in sorted(rows):
        lines.append("  <tr>")
        for cell in sorted(rows[row_i], key=lambda c: c["col"]):
            rs = cell["rowspan"]
            cs = cell["colspan"]
            attrs = ""
            if rs > 1:
                attrs += f' rowspan="{rs}"'
            if cs > 1:
                attrs += f' colspan="{cs}"'
            val = cell["value"] or ""
            lines.append(f"    <td{attrs}>{val}</td>")
        lines.append("  </tr>")
    lines.append("</table>")
    return "\n".join(lines)


def _table_to_raw(table) -> RawElement:  # type: ignore[no-untyped-def]
    """Convert a python-docx Table to a RawElement with merge-aware dual-format table data."""
    try:
        merge_map = _extract_merge_map(table)
        merge_error = False
    except Exception:
        merge_map = []
        merge_error = True

    if merge_error:
        # Fallback: flat extraction
        rows_data: list[list[str]] = []
        for row in table.rows:
            rows_data.append([cell.text.strip() for cell in row.cells])
        headers = rows_data[0] if rows_data else []
        data_rows = rows_data[1:] if len(rows_data) > 1 else []
        col_count = len(headers) if headers else (len(rows_data[0]) if rows_data else 0)
        if headers:
            sep = "|" + "|".join("---" for _ in headers) + "|"
            md_lines = ["|" + "|".join(headers) + "|", sep] + [
                "|" + "|".join(row) + "|" for row in data_rows
            ]
        else:
            md_lines = ["|" + "|".join(row) + "|" for row in data_rows]
        markdown = "\n".join(md_lines)
        return RawElement(
            element_type="table:simple",
            text=markdown,
            table_data={
                "markdown": markdown,
                "structured": {"headers": headers, "rows": data_rows},
                "has_header_row": bool(headers),
                "row_count": len(data_rows),
                "header_row_count": 1 if headers else 0,
                "col_count": col_count,
                "merge_map": [],
                "html": None,
                "table_extraction_method": "docx-native",
            },
            extra={
                "extraction_method": "python-docx:table",
                "confidence": 0.95,
                "table_merge_error": True,
            },
        )

    is_complex = any(c["rowspan"] > 1 or c["colspan"] > 1 for c in merge_map)
    element_type = "table:complex" if is_complex else "table:simple"

    # Build flat headers/rows from merge_map values
    row_groups: dict[int, list[str]] = {}
    for cell in merge_map:
        row_groups.setdefault(cell["row"], []).append(cell["value"])

    sorted_rows = [row_groups[i] for i in sorted(row_groups)]
    headers = sorted_rows[0] if sorted_rows else []
    data_rows = sorted_rows[1:] if len(sorted_rows) > 1 else []
    col_count = len(headers) if headers else (len(data_rows[0]) if data_rows else 0)

    # Build GFM markdown
    if headers:
        sep = "| " + " | ".join("---" for _ in headers) + " |"
        md_lines = ["| " + " | ".join(headers) + " |", sep]
        for row in data_rows:
            padded = list(row) + [""] * max(0, len(headers) - len(row))
            md_lines.append("| " + " | ".join(padded) + " |")
    else:
        md_lines = []
        for row in data_rows:
            md_lines.append("| " + " | ".join(row) + " |")
    markdown = "\n".join(md_lines)

    html = _build_html_table(merge_map, col_count) if is_complex else None

    return RawElement(
        element_type=element_type,
        text=markdown,
        table_data={
            "markdown": markdown,
            "structured": {"headers": headers, "rows": data_rows},
            "has_header_row": bool(headers),
            "row_count": len(data_rows),
            "header_row_count": 1 if headers else 0,
            "col_count": col_count,
            "merge_map": merge_map if is_complex else [],
            "html": html,
            "table_extraction_method": "docx-native",
        },
        extra={"extraction_method": "python-docx:table", "confidence": 0.95},
    )


def _extract_docx_figures(
    doc: object,
    figures_dir: object,
) -> Iterator[RawElement]:
    """Extract embedded images from a DOCX document and yield figure RawElements."""
    import logging

    _log = logging.getLogger(__name__)

    # Collect all inline shapes / drawings via relationship parts
    rels = doc.part.rels  # type: ignore[attr-defined]
    img_idx = 0
    seen_rids: set = set()

    # Walk all paragraphs and runs to find drawing/inline image relationships
    body = doc.element.body  # type: ignore[attr-defined]
    for para in body.iter():
        tag = para.tag.split("}")[-1] if "}" in para.tag else para.tag
        if tag not in ("drawing", "pict"):
            continue
        # Find blipFill / blip references
        for blip in para.iter():
            btag = blip.tag.split("}")[-1] if "}" in blip.tag else blip.tag
            if btag != "blip":
                continue
            # r:embed or r:link attribute
            r_embed = blip.get(
                "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}embed"
            )
            if not r_embed or r_embed in seen_rids:
                continue
            seen_rids.add(r_embed)
            rel = rels.get(r_embed)
            if rel is None:
                continue
            try:
                img_bytes = rel.target_part.blob
                img_idx += 1
                # Determine format from content type
                ct = rel.target_part.content_type or ""
                if "png" in ct:
                    ext = "png"
                elif "jpeg" in ct or "jpg" in ct:
                    ext = "jpg"
                elif "gif" in ct:
                    ext = "gif"
                else:
                    ext = "png"
                fname = f"fig_{img_idx:04d}.{ext}"
                asset_path = figures_dir / fname  # type: ignore[operator]
                asset_path.write_bytes(img_bytes)  # type: ignore[union-attr]
                yield RawElement(
                    element_type="media:figure",
                    text=f"[Figure: {fname}]",
                    figure_data={
                        "asset_path": str(asset_path),
                        "asset_format": ext,
                    },
                    extra={"extraction_method": "python-docx:image", "confidence": 0.85},
                )
            except Exception as exc:
                _log.warning("DOCX figure extraction failed rid=%s: %s", r_embed, exc)


class DOCXParser(BaseParser):
    handles = [("application/vnd.openxmlformats-officedocument.wordprocessingml.document", None)]

    def parse(self, file_record: FileRecord) -> Iterator[RawElement]:
        yield from self.parse_with_figures(file_record, figures_dir=None)

    def parse_with_figures(
        self, file_record: FileRecord, figures_dir: object = None
    ) -> Iterator[RawElement]:
        try:
            from docx import Document  # type: ignore[import-untyped]
        except ImportError as e:
            raise ImportError("python-docx required: pip install python-docx") from e

        doc = Document(str(file_record.path))

        # Collect all block-level elements in document order
        # python-docx doesn't preserve paragraph/table order in doc.paragraphs + doc.tables
        # We must iterate doc.element.body children

        body = doc.element.body
        para_map = {p._element: p for p in doc.paragraphs}
        table_map = {t._element: t for t in doc.tables}

        for child in body:
            tag = child.tag.split("}")[-1] if "}" in child.tag else child.tag

            if tag == "p":
                para = para_map.get(child)
                if para is None:
                    continue
                text = para.text.strip()
                if not text:
                    continue

                style_name = para.style.name if para.style else ""
                level = _heading_level_from_style(style_name)

                if level is not None:
                    yield RawElement(
                        element_type="structural:section_header"
                        if level > 1
                        else "structural:title",
                        text=text,
                        heading_level=level,
                        detection_tier=1,  # native DOCX heading styles
                        extra={
                            "extraction_method": "python-docx:heading",
                            "confidence": 0.98,
                            "style_name": style_name,
                        },
                    )
                else:
                    yield RawElement(
                        element_type="text:narrative",
                        text=text,
                        extra={
                            "extraction_method": "python-docx:paragraph",
                            "confidence": 0.95,
                            "style_name": style_name,
                        },
                    )

            elif tag == "tbl":
                table = table_map.get(child)
                if table is not None:
                    yield _table_to_raw(table)

        if figures_dir is not None:
            yield from _extract_docx_figures(doc, figures_dir)
