# pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportGeneralTypeIssues=false
from __future__ import annotations

import io
import zipfile

from app.services.document_parser.formats.docx.toc import detect_doc_tocs, detect_sdt_toc
from docx import Document
from docx.oxml.table import CT_Tbl
from docx.oxml.text.paragraph import CT_P
from docx.table import Table
from docx.text.paragraph import Paragraph
from loguru import logger
from lxml import etree


def iter_block_items(doc_data):
    doc_stream = io.BytesIO(doc_data)
    doc = Document(doc_stream)

    # python-docx mapping
    p_tbl_map = []
    for child in doc.element.body:
        if isinstance(child, CT_P):
            p_tbl_map.append(("p", child))
        elif isinstance(child, CT_Tbl):
            p_tbl_map.append(("tbl", child))

    with zipfile.ZipFile(io.BytesIO(doc_data), "r") as docx:
        xml = docx.read("word/document.xml")
        rels = etree.fromstring(docx.read("word/_rels/document.xml.rels"))
        rel_map = {
            r.get("Id"): r.get("Target") for r in rels.findall(".//{*}Relationship")
        }
        ns = {
            "a": "http://schemas.openxmlformats.org/drawingml/2006/main",
            "r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships",
            "w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main",
            "v": "urn:schemas-microsoft-com:vml",
            "o": "urn:schemas-microsoft-com:office:office",
        }
        r_ns = "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}"

        root = etree.fromstring(xml)
        body = root.find(".//w:body", namespaces=ns)

        ele_num = 1
        map_index = 0  # point to p_tbl_map
        toc_field_active = False

        for elem in body.iterchildren():
            if not isinstance(elem.tag, str):
                continue

            tag = etree.QName(elem.tag).localname

            # --- SDT (Structured Document Tag) container ---
            # TOC generated by MS Word is usually in sdt
            if tag == "sdt":
                sdt_toc_info = detect_sdt_toc(elem, ns)
                is_toc_sdt = sdt_toc_info["is_toc_sdt"]

                sdt_content = elem.find(".//w:sdtContent", namespaces=ns)
                if sdt_content is not None:
                    for p_elem in sdt_content.findall(".//w:p", namespaces=ns):
                        texts = p_elem.xpath(".//w:t/text()", namespaces=ns)
                        text = "".join(texts).strip()

                        if is_toc_sdt:
                            label = "TOC-AREA"
                            toc_info = detect_doc_tocs(p_elem, ns)
                        else:
                            toc_info = detect_doc_tocs(p_elem, ns)
                            if toc_info["is_style"] or toc_info["is_field_start"]:
                                label = "TOC-AREA"
                            else:
                                label = "PTXT"

                        if text:
                            meta = None
                            if "TOC" in label:
                                meta = {
                                    "toc_level": toc_info.get("toc_level"),
                                    "toc_outline_level": toc_info.get("outline_level"),
                                    "toc_left_indent": toc_info.get("left_indent"),
                                    "toc_style_name": toc_info.get("style_name"),
                                    "toc_source": "sdt",
                                }
                            yield ele_num, text, label, meta
                            ele_num += 1
                continue

            # --- text paras ---
            if tag == "p":
                texts = elem.xpath(".//w:t/text()", namespaces=ns)
                text = "".join(texts).strip()

                if map_index < len(p_tbl_map) and p_tbl_map[map_index][0] == "p":
                    p_obj = Paragraph(p_tbl_map[map_index][1], doc)
                else:
                    p_obj = None

                toc_info = detect_doc_tocs(elem, ns)
                if toc_info["is_field_start"]:
                    toc_field_active = True

                if toc_info["is_style"] or toc_field_active:
                    label = "TOC-AREA"
                else:
                    label = "PTXT"

                if text or p_obj is not None:
                    meta = None
                    if "TOC" in label:
                        meta = {
                            "toc_level": toc_info.get("toc_level"),
                            "toc_outline_level": toc_info.get("outline_level"),
                            "toc_left_indent": toc_info.get("left_indent"),
                            "toc_style_name": toc_info.get("style_name"),
                            "toc_source": "paragraph",
                        }
                    yield ele_num, p_obj or text, label, meta
                    ele_num += 1

                # images (DrawingML: <a:blip>)
                seen_rids = set()
                blips = elem.xpath(".//a:blip", namespaces=ns)
                for b in blips:
                    rid = b.get(f"{r_ns}embed")
                    if not rid or rid in seen_rids:
                        continue
                    target = rel_map.get(rid)
                    if not target or not target.startswith("media/"):
                        continue
                    seen_rids.add(rid)
                    data = docx.read("word/" + target)
                    yield (
                        ele_num,
                        None,
                        "IMAGE",
                        {
                            "image_name": target.split("/")[-1],
                            "from": "paragraph",
                            "size": len(data),
                            "data": data,
                        },
                    )
                    ele_num += 1

                # TODO: Re-evaluate VML group extraction strategy. 
                # Complex VML composite images (<v:group>) are currently skipped because extracting 
                # <v:imagedata> piece-by-piece loses textual overlay and positioning. 
                # Future plan: Use LibreOffice headless conversion to render the entire document 
                # and map the perfectly rendered images back to the layout via text anchors.
                #
                # Temporary: detect VML-only paragraphs and inject a placeholder so the
                # paragraph isn't silently swallowed, leaving its parent section empty.
                if not text and not seen_rids:
                    # No text and no DrawingML images — check for VML content
                    vml_groups = elem.xpath(".//v:group", namespaces=ns)
                    vml_images_check = elem.xpath(".//v:imagedata", namespaces=ns)
                    if vml_groups or vml_images_check:
                        vml_placeholder = "[VML graphic \u2014 extraction not yet supported]"
                        yield ele_num, vml_placeholder, "PTXT", None
                        ele_num += 1
                        logger.debug(
                            f"Injected VML placeholder for paragraph with "
                            f"{len(vml_groups)} v:group, {len(vml_images_check)} v:imagedata"
                        )
                """
                # images (VML: <v:imagedata>) — convert to PNG
                from PIL import Image as PILImage

                vml_images = elem.xpath(".//v:imagedata", namespaces=ns)
                for v in vml_images:
                    rid = v.get(f"{r_ns}id")
                    if not rid or rid in seen_rids:
                        continue
                    target = rel_map.get(rid)
                    if not target or not target.startswith("media/"):
                        continue
                    seen_rids.add(rid)
                    raw_data = docx.read("word/" + target)
                    # Convert to PNG for uniform downstream handling
                    try:
                        pil_img = PILImage.open(io.BytesIO(raw_data))
                        png_buf = io.BytesIO()
                        pil_img.save(png_buf, format="PNG")
                        png_data = png_buf.getvalue()
                    except Exception as e:
                        logger.warning(f"Failed to convert VML image to PNG: {e}")
                        continue
                    orig_name = target.split("/")[-1]
                    png_name = os.path.splitext(orig_name)[0] + ".png"
                    yield (
                        ele_num,
                        None,
                        "IMAGE",
                        {
                            "image_name": png_name,
                            "from": "paragraph_vml",
                            "size": len(png_data),
                            "data": png_data,
                        },
                    )
                    ele_num += 1
                """
                map_index += 1

                if toc_info["is_field_end"]:
                    toc_field_active = False

            # --- tables ---
            elif tag == "tbl":
                if map_index < len(p_tbl_map) and p_tbl_map[map_index][0] == "tbl":
                    tbl = Table(p_tbl_map[map_index][1], doc)
                else:
                    tbl = Table(elem, doc)

                # Extract images from each cell, keyed by (row_idx, col_idx)
                cell_images = {}  # {(row_idx, col_idx): [{'image_name', 'data', 'size'}]}
                for row_idx, tr in enumerate(elem.findall(".//w:tr", namespaces=ns)):
                    for col_idx, tc in enumerate(tr.findall(".//w:tc", namespaces=ns)):
                        cell_seen_rids = set()
                        imgs_in_cell = []
                        # DrawingML images in cell
                        blips = tc.xpath(".//a:blip", namespaces=ns)
                        for b in blips:
                            rid = b.get(f"{r_ns}embed")
                            if not rid or rid in cell_seen_rids:
                                continue
                            target = rel_map.get(rid)
                            if not target or not target.startswith("media/"):
                                continue
                            cell_seen_rids.add(rid)
                            data = docx.read("word/" + target)
                            if (
                                len(data) < 10 * 1024
                            ):  # Skip small images (<10KB, likely icons)
                                continue
                            imgs_in_cell.append(
                                {
                                    "image_name": target.split("/")[-1],
                                    "data": data,
                                    "size": len(data),
                                }
                            )
                        # TODO: VML in tables is temporarily skipped to avoid extracting 
                        # fragmented textless background images. (Same as paragraph VML logic)
                        """
                        # VML images in cell — convert to PNG
                        from PIL import Image as PILImage

                        vml_in_cell = tc.xpath(".//v:imagedata", namespaces=ns)
                        for v in vml_in_cell:
                            rid = v.get(f"{r_ns}id")
                            if not rid or rid in cell_seen_rids:
                                continue
                            target = rel_map.get(rid)
                            if not target or not target.startswith("media/"):
                                continue
                            cell_seen_rids.add(rid)
                            raw_data = docx.read("word/" + target)
                            try:
                                pil_img = PILImage.open(io.BytesIO(raw_data))
                                png_buf = io.BytesIO()
                                pil_img.save(png_buf, format="PNG")
                                png_data = png_buf.getvalue()
                            except Exception as e:
                                logger.warning(
                                    f"Failed to convert VML cell image to PNG: {e}"
                                )
                                continue
                            if len(png_data) < 10 * 1024:
                                continue
                            orig_name = target.split("/")[-1]
                            png_name = os.path.splitext(orig_name)[0] + ".png"
                            imgs_in_cell.append(
                                {
                                    "image_name": png_name,
                                    "data": png_data,
                                    "size": len(png_data),
                                }
                            )
                        """
                        if imgs_in_cell:
                            cell_images[(row_idx, col_idx)] = imgs_in_cell

                yield ele_num, tbl, "TABLE", cell_images if cell_images else None
                ele_num += 1
                map_index += 1
            else:
                continue

        # --- handle p_tbl_map at the end ---
        while map_index < len(p_tbl_map):
            tag, node = p_tbl_map[map_index]
            if tag == "p":
                toc_info = detect_doc_tocs(node, ns)
                label = "TOC-AREA" if toc_info["is_style"] else "PTXT"
                meta = None
                if "TOC" in label:
                    meta = {
                        "toc_level": toc_info.get("toc_level"),
                        "toc_outline_level": toc_info.get("outline_level"),
                        "toc_left_indent": toc_info.get("left_indent"),
                        "toc_style_name": toc_info.get("style_name"),
                        "toc_source": "tail-map",
                    }
                yield ele_num, Paragraph(node, doc), label, meta
            elif tag == "tbl":
                yield ele_num, Table(node, doc), "TABLE", None
            ele_num += 1
            map_index += 1

