"""
DOCX to Markdown converter using raw XML parsing.
Extracts paragraphs, tables, headers/footers — preserves 100% content.
Supports lvlText templates (e.g. "Điều %1.") for correct numbering.
"""

import re
import zipfile
import xml.etree.ElementTree as ET
from io import BytesIO
from typing import Dict, List, Tuple, Optional


WORD_NS = "http://schemas.openxmlformats.org/wordprocessingml/2006/main"
NS = {"w": WORD_NS}


def _get_numbering_map(z: zipfile.ZipFile) -> dict:
    """
    Parse word/numbering.xml to build numId+level -> (format, lvlText, start) mapping.
    Returns: {numId: {ilvl: {"fmt": ..., "lvlText": ..., "start": ...}}}
    """
    mapping = {}
    try:
        numbering_xml = z.read("word/numbering.xml")
    except KeyError:
        return mapping

    root = ET.fromstring(numbering_xml)

    # Build abstractNum lookup
    abstract_nums = {}
    for abstract in root.findall(".//w:abstractNum", NS):
        abstract_id = abstract.get(f"{{{WORD_NS}}}abstractNumId")
        levels = {}
        for lvl in abstract.findall("w:lvl", NS):
            ilvl = lvl.get(f"{{{WORD_NS}}}ilvl")

            num_fmt_el = lvl.find("w:numFmt", NS)
            fmt = num_fmt_el.get(f"{{{WORD_NS}}}val") if num_fmt_el is not None else "bullet"

            lvl_text_el = lvl.find("w:lvlText", NS)
            lvl_text = lvl_text_el.get(f"{{{WORD_NS}}}val") if lvl_text_el is not None else "%1."

            start_el = lvl.find("w:start", NS)
            start_val = int(start_el.get(f"{{{WORD_NS}}}val")) if start_el is not None else 1

            levels[ilvl] = {
                "fmt": fmt,
                "lvlText": lvl_text,
                "start": start_val,
            }
        abstract_nums[abstract_id] = levels

    # Build numId -> abstractNumId mapping
    for num in root.findall(".//w:num", NS):
        num_id = num.get(f"{{{WORD_NS}}}numId")
        abstract_ref = num.find("w:abstractNumId", NS)
        if abstract_ref is not None:
            abstract_id = abstract_ref.get(f"{{{WORD_NS}}}val")
            if abstract_id in abstract_nums:
                mapping[num_id] = abstract_nums[abstract_id]

    return mapping


def _format_number(value: int, fmt: str) -> str:
    """Convert a number to the specified format."""
    if fmt == "decimal":
        return str(value)
    elif fmt == "lowerLetter":
        return chr(ord('a') + value - 1) if 1 <= value <= 26 else str(value)
    elif fmt == "upperLetter":
        return chr(ord('A') + value - 1) if 1 <= value <= 26 else str(value)
    elif fmt == "lowerRoman":
        romans = [(1000,'m'),(900,'cm'),(500,'d'),(400,'cd'),(100,'c'),(90,'xc'),
                  (50,'l'),(40,'xl'),(10,'x'),(9,'ix'),(5,'v'),(4,'iv'),(1,'i')]
        result = ""
        for val, numeral in romans:
            while value >= val:
                result += numeral
                value -= val
        return result
    elif fmt == "upperRoman":
        romans = [(1000,'M'),(900,'CM'),(500,'D'),(400,'CD'),(100,'C'),(90,'XC'),
                  (50,'L'),(40,'XL'),(10,'X'),(9,'IX'),(5,'V'),(4,'IV'),(1,'I')]
        result = ""
        for val, numeral in romans:
            while value >= val:
                result += numeral
                value -= val
        return result
    else:
        return str(value)


def _render_lvl_text(
    lvl_text: str,
    counters: Dict[str, int],
    num_id: str,
    current_ilvl: str,
    numbering_map: dict,
) -> str:
    """
    Render a lvlText template like 'Điều %1.' or '%1.%2.' by substituting
    counter values for each level.
    """
    result = lvl_text

    # Find all %N references in the lvlText
    for match in re.finditer(r"%(\d+)", lvl_text):
        ref_level = str(int(match.group(1)) - 1)  # %1 = ilvl 0, %2 = ilvl 1, etc.
        counter_key = f"{num_id}_{ref_level}"
        counter_val = counters.get(counter_key, 1)

        # Get the format for this referenced level
        ref_fmt = "decimal"
        if num_id in numbering_map and ref_level in numbering_map[num_id]:
            ref_fmt = numbering_map[num_id][ref_level]["fmt"]

        formatted = _format_number(counter_val, ref_fmt)
        result = result.replace(match.group(0), formatted, 1)

    return result


def _extract_run_text(run) -> Tuple[str, bool, bool]:
    """Extract text from a run element with bold/italic info."""
    texts = []
    for t in run.findall(".//w:t", NS):
        if t.text:
            texts.append(t.text)

    text = "".join(texts)

    rpr = run.find("w:rPr", NS)
    bold = False
    italic = False
    if rpr is not None:
        b_el = rpr.find("w:b", NS)
        if b_el is not None:
            val = b_el.get(f"{{{WORD_NS}}}val", "true")
            bold = val.lower() != "false"
        i_el = rpr.find("w:i", NS)
        if i_el is not None:
            val = i_el.get(f"{{{WORD_NS}}}val", "true")
            italic = val.lower() != "false"

    return text, bold, italic


def _merge_runs(runs_data: List[Tuple[str, bool, bool]]) -> str:
    """
    Merge adjacent runs with the same formatting to avoid ****artifacts****.
    """
    if not runs_data:
        return ""

    # Group consecutive runs with same formatting
    groups = []
    for text, bold, italic in runs_data:
        if not text:
            continue
        if groups and groups[-1][1] == bold and groups[-1][2] == italic:
            groups[-1] = (groups[-1][0] + text, bold, italic)
        else:
            groups.append((text, bold, italic))

    # Convert groups to markdown
    parts = []
    for text, bold, italic in groups:
        if bold and italic:
            parts.append(f"***{text}***")
        elif bold:
            parts.append(f"**{text}**")
        elif italic:
            parts.append(f"*{text}*")
        else:
            parts.append(text)

    return "".join(parts)


def _paragraph_to_markdown(
    para,
    numbering_map: dict,
    list_counters: dict,
) -> str:
    """Convert a single paragraph element to markdown."""
    # Extract style
    ppr = para.find("w:pPr", NS)
    style_name = ""
    num_id = None
    ilvl = "0"

    if ppr is not None:
        style_el = ppr.find("w:pStyle", NS)
        if style_el is not None:
            style_name = style_el.get(f"{{{WORD_NS}}}val", "")

        # Check for list numbering
        num_pr = ppr.find("w:numPr", NS)
        if num_pr is not None:
            num_id_el = num_pr.find("w:numId", NS)
            ilvl_el = num_pr.find("w:ilvl", NS)
            if num_id_el is not None:
                num_id = num_id_el.get(f"{{{WORD_NS}}}val")
            if ilvl_el is not None:
                ilvl = ilvl_el.get(f"{{{WORD_NS}}}val", "0")

    # Process all runs — collect then merge
    runs_data = []
    for run in para.findall(".//w:r", NS):
        text, bold, italic = _extract_run_text(run)
        if text:
            runs_data.append((text, bold, italic))

    line = _merge_runs(runs_data).strip()
    if not line:
        return ""

    # Lọc bỏ các dòng chỉ chứa số trang hoặc chữ "Trang / Page" do OCR chèn nhầm
    if line.isdigit() and len(line) <= 3:
        return ""
    if re.fullmatch(r"Trang\s+\d+(/\d+)?|Page\s+\d+(/\d+)?|\d+\s*-.*-.*|^\d+$", line, re.IGNORECASE):
        return ""

    # ── Detect heading level ──────────────────────────────
    # Ưu tiên 1: Đọc outlineLvl trực tiếp từ XML (chính xác nhất, Word tự gắn)
    # Ưu tiên 2: Đọc từ tên Style (Heading1, Heading2...)
    # Ưu tiên 3: Regex nhận diện cấu trúc pháp lý VN (Phần, Chương, Điều, Mục...)
    #            CHỈ áp dụng khi style có chứa "heading" hoặc "title"
    # → Tuyệt đối KHÔNG bịa heading từ text bôi đậm thông thường.

    heading_level = 0

    # ── P1: outlineLvl trong XML (ưu tiên cao nhất) ──
    if ppr is not None:
        outline_el = ppr.find("w:outlineLvl", NS)
        if outline_el is not None:
            try:
                raw_lvl = int(outline_el.get(f"{{{WORD_NS}}}val", "9"))
                if 0 <= raw_lvl <= 5:
                    heading_level = raw_lvl + 1  # outlineLvl 0 = H1, 1 = H2, ...
            except ValueError:
                pass

    # ── P2: Style name (nếu outlineLvl không có) ──
    if heading_level == 0:
        style_lower = style_name.lower()
        if "heading" in style_lower or "title" in style_lower:
            for i in range(1, 7):
                if str(i) in style_name:
                    heading_level = i
                    break
            if heading_level == 0:
                heading_level = 1  # Generic "Heading" hoặc "Title" → H1

    if heading_level > 0:
        return f"{'#' * heading_level} {line}"

    # Apply list formatting with lvlText support
    if num_id is not None and num_id != "0":
        indent = "    " * int(ilvl)
        fmt = "bullet"
        lvl_text = "%1."

        if num_id in numbering_map and ilvl in numbering_map[num_id]:
            lvl_info = numbering_map[num_id][ilvl]
            fmt = lvl_info["fmt"]
            lvl_text = lvl_info["lvlText"]

        if fmt == "bullet":
            return f"{indent}- {line}"

        # Increment counter for current level
        counter_key = f"{num_id}_{ilvl}"
        list_counters[counter_key] = list_counters.get(
            counter_key,
            numbering_map.get(num_id, {}).get(ilvl, {}).get("start", 1) - 1,
        ) + 1

        # ── CRITICAL: Reset sub-level counters when parent level changes ──
        # Khi ĐIỀU 4 (ilvl 0) tăng counter, phải reset ilvl 1,2,3...
        # Nếu không: 4.1 thành 4.10 vì counter ilvl 1 vẫn tiếp tục từ ĐIỀU trước
        current_ilvl_int = int(ilvl)
        if num_id in numbering_map:
            for sub_ilvl in list(numbering_map[num_id].keys()):
                if int(sub_ilvl) > current_ilvl_int:
                    sub_key = f"{num_id}_{sub_ilvl}"
                    if sub_key in list_counters:
                        del list_counters[sub_key]

        # Render the full numbering text
        rendered = _render_lvl_text(
            lvl_text, list_counters, num_id, ilvl, numbering_map
        )

        return f"{indent}{rendered} {line}"

    return line


def _table_to_markdown(table) -> str:
    """Convert a table element to markdown table."""
    rows = table.findall(".//w:tr", NS)
    if not rows:
        return ""

    md_rows = []
    max_cols = 0

    for row in rows:
        cells = row.findall(".//w:tc", NS)
        cell_texts = []
        for cell in cells:
            # Extract all text from cell paragraphs
            cell_parts = []
            for p in cell.findall(".//w:p", NS):
                runs_data = []
                for run in p.findall(".//w:r", NS):
                    text, bold, italic = _extract_run_text(run)
                    if text:
                        runs_data.append((text, bold, italic))
                p_text = _merge_runs(runs_data).strip()
                if p_text:
                    cell_parts.append(p_text)
            cell_texts.append(" ".join(cell_parts))
        md_rows.append(cell_texts)
        max_cols = max(max_cols, len(cell_texts))

    if max_cols == 0:
        return ""

    # Normalize row lengths
    for row in md_rows:
        while len(row) < max_cols:
            row.append("")

    lines = []
    # Header row
    header = md_rows[0]
    lines.append("| " + " | ".join(header) + " |")
    lines.append("| " + " | ".join(["---"] * max_cols) + " |")

    # Data rows
    for row in md_rows[1:]:
        lines.append("| " + " | ".join(row) + " |")

    return "\n".join(lines)


def convert_docx_to_markdown(file_bytes: bytes, filename: str) -> dict:
    """
    Convert DOCX file bytes to markdown content.
    Returns dict matching Docling response format.
    """
    z = zipfile.ZipFile(BytesIO(file_bytes))
    doc_xml = z.read("word/document.xml")
    root = ET.fromstring(doc_xml)

    numbering_map = _get_numbering_map(z)
    list_counters: Dict[str, int] = {}

    # Find the document body
    body = root.find(f"{{{WORD_NS}}}body")
    if body is None:
        body = root

    md_parts = []

    # Process all top-level elements in order (paragraphs and tables)
    for element in body:
        tag = element.tag.replace(f"{{{WORD_NS}}}", "")

        if tag == "p":
            line = _paragraph_to_markdown(element, numbering_map, list_counters)
            if line:
                md_parts.append(line)
            else:
                md_parts.append("")  # preserve blank lines

        elif tag == "tbl":
            table_md = _table_to_markdown(element)
            if table_md:
                md_parts.append("")
                md_parts.append(table_md)
                md_parts.append("")

        elif tag == "sectPr":
            pass  # skip section properties

    # Also extract headers and footers
    header_footer_parts = _extract_headers_footers(z)
    if header_footer_parts:
        md_parts = header_footer_parts + ["", "---", ""] + md_parts

    # Clean up excessive blank lines
    md_content = "\n".join(md_parts)
    while "\n\n\n" in md_content:
        md_content = md_content.replace("\n\n\n", "\n\n")
    md_content = md_content.strip()

    # ── Post-processing: fix heading levels dựa trên nội dung text ──
    md_content = _postprocess_legal_headings(md_content)

    return {
        "document": {
            "filename": filename,
            "md_content": md_content,
        }
    }


# ───────────────────────────────────────────────────────────────
# Post-processing: regex-based heading detection trên markdown
# ───────────────────────────────────────────────────────────────

_RE_DIEU = re.compile(
    r"^(?:#+\s+)?(?:\*{0,3})((?:ĐIỀU|Điều)\s+\d+[\.\:]?\s*.+?)(?:\*{0,3})\s*$",
    re.MULTILINE | re.UNICODE,
)
_RE_H2_NUM = re.compile(
    r"^(?:#+\s+)?(?:\*{0,3})(\d+\.\d+\.?\s+.+?)(?:\*{0,3})\s*$",
    re.MULTILINE | re.UNICODE,
)
_RE_H3_NUM = re.compile(
    r"^(?:#+\s+)?(?:\*{0,3})(\d+\.\d+\.\d+\.?\s+.+?)(?:\*{0,3})\s*$",
    re.MULTILINE | re.UNICODE,
)
_RE_ALL_CAPS_LINE = re.compile(
    r"^(?:#+\s+)?(?:\*{0,3})([A-ZÀ-Ỹ\s\-–:,\.]+)(?:\*{0,3})\s*$",
    re.MULTILINE | re.UNICODE,
)


def _postprocess_legal_headings(md: str) -> str:
    """
    Post-process markdown: phát hiện cấu trúc pháp lý VN và gán đúng heading.
    Chạy regex trên markdown text đã convert, bất kể DOCX heading style.
    """
    lines = md.split("\n")
    result = []

    for line in lines:
        stripped = line.strip()
        if not stripped:
            result.append(line)
            continue

        # Bỏ heading markers hiện có để đánh giá lại
        clean = re.sub(r"^#+\s+", "", stripped)
        # Bỏ bold/italic markers để check text gốc
        text_only = re.sub(r"\*{1,3}", "", clean).strip()

        if not text_only or len(text_only) > 150:
            result.append(line)
            continue

        # ── Priority 1: "ĐIỀU X." pattern → H1 ──
        if re.match(r"^(?:ĐIỀU|Điều)\s+\d+[\.\:\s]", text_only, re.UNICODE | re.IGNORECASE):
            result.append(f"# {clean}")
            continue

        # ── Priority 2: ALL CAPS short line → H1 ──
        letters = [c for c in text_only if c.isalpha()]
        if letters and all(c == c.upper() for c in letters) and len(text_only) < 100:
            # Không phải dòng bảng hoặc list
            if not text_only.startswith("|") and not text_only.startswith("-"):
                result.append(f"# {clean}")
                continue

        # ── Priority 3: "X.Y.Z." pattern → H3 ──
        if re.match(r"^\d+\.\d+\.\d+\.?\s", text_only):
            result.append(f"### {clean}")
            continue

        # ── Priority 4: "X.Y." pattern → H2 ──
        if re.match(r"^\d+\.\d+\.?\s", text_only):
            result.append(f"## {clean}")
            continue

        result.append(line)

    return "\n".join(result)


def _extract_headers_footers(z: zipfile.ZipFile) -> List[str]:
    """Extract text from document headers and footers."""
    parts = []
    for name in z.namelist():
        if name.startswith("word/header") or name.startswith("word/footer"):
            try:
                xml_data = z.read(name)
                root = ET.fromstring(xml_data)
                for p in root.findall(f".//{{{WORD_NS}}}p"):
                    texts = []
                    for t in p.findall(f".//{{{WORD_NS}}}t"):
                        if t.text:
                            texts.append(t.text)
                    line = "".join(texts).strip()
                    if line and line not in parts:
                        parts.append(line)
            except Exception:
                pass
    return parts
