"""Pre-Docling header sentinel injection -> Docling -> Markdown chunking.

Standalone CLI/dev tool. Delegates the sentinel + ATX conversion logic to
src.extraction_v2.xlsx_header_injection so the production airflow flow and
this debug script share one implementation.

Inputs an .xlsx (typically the .preprocessed.xlsx already produced by
embed_mermaid_and_parse.py so mermaid is already embedded), runs the same
pre-Docling header injection used by ExcelParserService, then runs Docling
and the LlamaIndex hierarchical chunker.

Outputs in CWD:
    <base>.headered.xlsx
    <base>.headered.md
    <base>.headered.chunks.json
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import re
import shutil
import sys
import tempfile
import xml.etree.ElementTree as ET
import zipfile
from pathlib import Path
from typing import Dict, List, Tuple

# Repo-root-relative imports so the script can run via `uv run python scripts/...`
REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT))

from src.extraction_v2.xlsx_header_injection import (
    NS,
    _load_shared_strings,
    _read_cell_text,
    convert_sentinels_to_atx,
    inject_header_sentinels,
)

logger = logging.getLogger("inject_headers_and_parse")


# ---------------------------------------------------------------------------
# Step 2a: Collect embedded mermaid blocks from the xlsx (so we can restore
# their newlines after Docling collapses them into a single table-cell line).
# Only used when running Docling stand-alone (the airflow flow uses
# `unwrap_mermaid_tables` instead).
# ---------------------------------------------------------------------------

def collect_mermaid_blocks(xlsx_path: str) -> List[str]:
    workdir = tempfile.mkdtemp(prefix="mermaid_collect_")
    try:
        with zipfile.ZipFile(xlsx_path) as z:
            z.extractall(workdir)

        ss_path = os.path.join(workdir, "xl/sharedStrings.xml")
        shared = _load_shared_strings(ET.parse(ss_path).getroot()) if os.path.exists(ss_path) else []

        out: List[str] = []
        ws_dir = os.path.join(workdir, "xl/worksheets")
        for fname in sorted(os.listdir(ws_dir)):
            if not fname.endswith(".xml"):
                continue
            root = ET.parse(os.path.join(ws_dir, fname)).getroot()
            sd = root.find(f"{{{NS}}}sheetData")
            if sd is None:
                continue
            for row in sd.findall(f"{{{NS}}}row"):
                for c in row.findall(f"{{{NS}}}c"):
                    text = _read_cell_text(c, shared)
                    if text and text.lstrip().startswith("```mermaid"):
                        out.append(text)
        return out
    finally:
        shutil.rmtree(workdir, ignore_errors=True)


def restore_mermaid_blocks(md: str, mermaid_texts: List[str]) -> Tuple[str, int]:
    restored = 0
    for original in mermaid_texts:
        tokens = original.split()
        if not tokens:
            continue
        flex = r"\s+".join(re.escape(t) for t in tokens)
        pattern = r"\|?\s*" + flex + r"\s*\|?"
        replacement = "\n\n" + original + "\n\n"
        new_md, n = re.subn(pattern, replacement, md, count=1)
        if n == 0:
            logger.warning("could not relocate mermaid block (%d chars)", len(original))
            continue
        md = new_md
        restored += 1
    return md, restored


def docling_to_md(xlsx_path: str) -> str:
    from docling.document_converter import DocumentConverter
    return DocumentConverter().convert(xlsx_path).document.export_to_markdown()


# ---------------------------------------------------------------------------
# Step 4: hierarchical chunking (line-based mermaid-aware leaf splitter)
# ---------------------------------------------------------------------------

MERMAID_FENCE_START_RE = re.compile(r"^```mermaid\b")
FENCE_END_RE = re.compile(r"^```\s*$")
HEADER_ONLY_RE = re.compile(r"^\s*#{1,6} [^\n]+\s*$")

CHUNK_TARGET_TOKENS = 512
EMBEDDING_CONTEXT_TOKENS = 8191
MIN_CHUNK_TOKENS = 64


def _split_oversize_mermaid(block_lines: List[str], tokenize, limit: int) -> List[str]:
    opener, closer = block_lines[0], block_lines[-1]
    inner = block_lines[1:-1]
    wrapper_tok = tokenize(opener + "\n\n" + closer)
    parts: List[str] = []
    cur: List[str] = []
    cur_tok = wrapper_tok
    for ln in inner:
        ln_tok = tokenize(ln + "\n")
        if wrapper_tok + ln_tok > limit:
            raise ValueError(f"single mermaid line exceeds embedding context ({ln_tok} tokens)")
        if cur_tok + ln_tok > limit and cur:
            parts.append(opener + "\n" + "\n".join(cur) + "\n" + closer)
            cur = [ln]
            cur_tok = wrapper_tok + ln_tok
        else:
            cur.append(ln)
            cur_tok += ln_tok
    if cur:
        parts.append(opener + "\n" + "\n".join(cur) + "\n" + closer)
    return parts


def chunk_with_atomic_mermaid(text: str, target: int, max_size: int, tokenize) -> List[Tuple[str, bool]]:
    lines = text.split("\n")
    chunks: List[Tuple[str, bool]] = []
    cur: List[str] = []
    cur_tok = 0
    cur_has_mermaid = False

    def flush() -> None:
        nonlocal cur, cur_tok, cur_has_mermaid
        if cur:
            chunks.append(("\n".join(cur), cur_has_mermaid))
        cur, cur_tok, cur_has_mermaid = [], 0, False

    i = 0
    n = len(lines)
    while i < n:
        ln = lines[i]
        if MERMAID_FENCE_START_RE.match(ln):
            j = i + 1
            while j < n and not FENCE_END_RE.match(lines[j]):
                j += 1
            end = j if j < n else n - 1
            block_lines = lines[i:end + 1]
            block_text = "\n".join(block_lines)
            block_tok = tokenize(block_text)

            if block_tok > max_size:
                flush()
                for part in _split_oversize_mermaid(block_lines, tokenize, max_size):
                    chunks.append((part, True))
                i = end + 1
                continue

            if cur_tok + block_tok > max_size and cur:
                flush()
            cur.extend(block_lines)
            cur_tok += block_tok
            cur_has_mermaid = True
            i = end + 1
        else:
            tok_n = tokenize(ln + "\n")
            limit = max_size if cur_has_mermaid else target
            if cur_tok + tok_n > limit and cur:
                flush()
            cur.append(ln)
            cur_tok += tok_n
            i += 1
    flush()

    merged: List[Tuple[str, bool]] = []
    for text, has_m in chunks:
        if merged and tokenize(text) < MIN_CHUNK_TOKENS:
            prev_text, prev_has = merged[-1]
            combined = prev_text + "\n" + text
            if tokenize(combined) <= max_size:
                merged[-1] = (combined, prev_has or has_m)
                continue
        merged.append((text, has_m))
    return merged


def hierarchical_chunk(md: str, base_name: str, out_path: str) -> Dict[str, int]:
    import tiktoken
    from llama_index.core import Document
    from llama_index.core.node_parser import MarkdownNodeParser
    from llama_index.core.schema import NodeRelationship, RelatedNodeInfo, TextNode

    enc = tiktoken.get_encoding("cl100k_base")
    def _tok(s: str) -> int:
        return len(enc.encode(s))

    base_meta = {
        "filename": f"{base_name}.headered.md",
        "document_type": "markdown",
        "page_number": 1,
    }

    sections = MarkdownNodeParser().get_nodes_from_documents(
        [Document(text=md, metadata=dict(base_meta))]
    )

    all_nodes: List = list(sections)
    mermaid_leaf_count = 0
    for section in sections:
        leaf_specs = chunk_with_atomic_mermaid(
            section.text, CHUNK_TARGET_TOKENS, EMBEDDING_CONTEXT_TOKENS, _tok
        )
        if len(leaf_specs) == 1 and leaf_specs[0][0] == section.text:
            continue
        child_rels = section.relationships.get(NodeRelationship.CHILD) or []
        if child_rels and not isinstance(child_rels, list):
            child_rels = [child_rels]
        for leaf_index, (text, has_mermaid) in enumerate(leaf_specs):
            leaf = TextNode(
                text=text,
                metadata={
                    **dict(section.metadata),
                    "leaf_index": leaf_index,
                    "has_mermaid": has_mermaid,
                },
            )
            leaf.relationships[NodeRelationship.PARENT] = RelatedNodeInfo(node_id=section.node_id)
            child_rels.append(RelatedNodeInfo(node_id=leaf.node_id))
            if has_mermaid:
                mermaid_leaf_count += 1
            all_nodes.append(leaf)
        section.relationships[NodeRelationship.CHILD] = child_rels

    nodes = all_nodes

    drop: set[str] = set()
    for n in nodes:
        if HEADER_ONLY_RE.match(n.text):
            drop.add(n.node_id)
    changed = True
    while changed:
        changed = False
        for n in nodes:
            if n.node_id in drop:
                continue
            child_rel = n.relationships.get(NodeRelationship.CHILD)
            if not child_rel:
                continue
            children = child_rel if isinstance(child_rel, list) else [child_rel]
            if children and all(r.node_id in drop for r in children):
                drop.add(n.node_id)
                changed = True
    nodes = [n for n in nodes if n.node_id not in drop]

    node_by_id = {n.node_id: n for n in nodes}
    surviving = set(node_by_id.keys())

    def _depth(n) -> int:
        d = 0
        cur = n
        while True:
            parent_rel = cur.relationships.get(NodeRelationship.PARENT)
            if parent_rel is None:
                return d
            pid = parent_rel.node_id
            if pid not in node_by_id:
                return d + 1
            cur = node_by_id[pid]
            d += 1

    leaf_ids: set[str] = set()
    for n in nodes:
        child_rel = n.relationships.get(NodeRelationship.CHILD)
        children = child_rel if isinstance(child_rel, list) else ([child_rel] if child_rel else [])
        if not any(r.node_id in surviving for r in children):
            leaf_ids.add(n.node_id)

    serialized = []
    for n in nodes:
        parent_rel = n.relationships.get(NodeRelationship.PARENT)
        child_rels = n.relationships.get(NodeRelationship.CHILD) or []
        if not isinstance(child_rels, list):
            child_rels = [child_rels]
        serialized.append({
            "node_id": n.node_id,
            "depth": _depth(n),
            "is_leaf": n.node_id in leaf_ids,
            "parent_id": parent_rel.node_id if parent_rel else None,
            "child_ids": [r.node_id for r in child_rels],
            "tokens": _tok(n.text),
            "text": n.text,
            "metadata": dict(n.metadata),
        })
    with open(out_path, "w", encoding="utf-8") as fh:
        json.dump(serialized, fh, ensure_ascii=False, indent=2)

    from collections import Counter
    depth_hist = dict(Counter(s["depth"] for s in serialized))
    return {
        "total_nodes": len(serialized),
        "leaves": len(leaf_ids),
        "parents": len(serialized) - len(leaf_ids),
        "mermaid_bearing_leaves": mermaid_leaf_count,
        "depth_histogram": depth_hist,
    }


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("xlsx", help="Input .xlsx (typically the .preprocessed.xlsx with mermaid embedded)")
    ap.add_argument("-v", "--verbose", action="store_true")
    args = ap.parse_args()
    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )

    src = args.xlsx
    if not os.path.exists(src):
        logger.error("input not found: %s", src)
        return 2
    if not src.lower().endswith(".xlsx"):
        logger.error("input must be .xlsx (run embed_mermaid_and_parse first if you have .xls)")
        return 2

    base_name = Path(src).stem
    if base_name.endswith(".preprocessed"):
        base_name = base_name[: -len(".preprocessed")]

    cwd = os.getcwd()
    headered_xlsx = os.path.join(cwd, f"{base_name}.headered.xlsx")
    headered_md = os.path.join(cwd, f"{base_name}.headered.md")
    chunks_json = os.path.join(cwd, f"{base_name}.headered.chunks.json")

    shutil.copyfile(src, headered_xlsx)
    n, body, size_map = inject_header_sentinels(headered_xlsx)
    logger.info(
        "injected sentinels into %d header cell(s) (body=%s, size_to_level=%s) -> %s",
        n, body, size_map, headered_xlsx,
    )

    mermaid_blocks = collect_mermaid_blocks(src)
    logger.info("found %d mermaid block(s) embedded in xlsx", len(mermaid_blocks))

    md_raw = docling_to_md(headered_xlsx)
    md_raw, n_restored = restore_mermaid_blocks(md_raw, mermaid_blocks)
    logger.info("restored mermaid newlines for %d/%d block(s)", n_restored, len(mermaid_blocks))
    md, counts = convert_sentinels_to_atx(md_raw)
    logger.info("ATX headers materialized: %s", counts)

    with open(headered_md, "w", encoding="utf-8") as fh:
        fh.write(md)
    logger.info("wrote %s (%d chars)", headered_md, len(md))

    stats = hierarchical_chunk(md, base_name, chunks_json)
    logger.info("chunking stats: %s", stats)
    logger.info("wrote %s", chunks_json)

    print()
    print(f"Headered xlsx:  {headered_xlsx}")
    print(f"Headered md:    {headered_md}")
    print(f"Chunks JSON:    {chunks_json}")
    print(f"ATX counts:     {counts}")
    print(f"Chunks:         {stats}")
    return 0


if __name__ == "__main__":
    sys.exit(main())