"""Stage 1-4 extraction + frontend layout mirror (ParagraphEditor.tsx)."""
from __future__ import annotations

import hashlib
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "gateway"))
import pdfengine  # type: ignore


def line_advances(line) -> tuple[dict[int, list[float]], float]:
    runs = line.runs
    seq: list[tuple[int, int, float]] = []
    aligned: dict[int, bool] = {}
    for ri, r in enumerate(runs):
        gs = r.glyphs
        ok = bool(r.text) and len(gs) == len(r.text)
        aligned[ri] = ok
        if ok:
            for ci, g in enumerate(gs):
                seq.append((ri, ci, g.origin_x))
    per_run: dict[int, list[float]] = {}
    for ri, r in enumerate(runs):
        if aligned.get(ri):
            per_run[ri] = [0.0] * len(r.text)
    anchor_x = seq[0][2] if seq else line.x
    for k, (ri, ci, ox) in enumerate(seq):
        gs = runs[ri].glyphs
        g = gs[ci]
        char_w = g.bbox_w if g.bbox_w > 0 else (runs[ri].font_size or 12) * 0.5
        per_run[ri][ci] = (seq[k + 1][2] - ox) if k + 1 < len(seq) else char_w
    for ri in list(per_run.keys()):
        if any(a <= 0 for a in per_run[ri]):
            del per_run[ri]
    return per_run, anchor_x


def para_eff_size(lines) -> float:
    nominal = 0.0
    for line in lines:
        for r in line.runs:
            sz = max(r.font_size or 0, r.h or 0)
            if sz > nominal:
                nominal = sz
    return nominal or 12.0


def median(xs: list[float]) -> float:
    if not xs:
        return 0.0
    s = sorted(xs)
    return s[len(s) // 2]


def compute_layout(para) -> dict:
    lines = para.lines
    eff_size = para_eff_size(lines)
    column_left = float("inf")
    first_baseline_y = float("-inf")
    baselines: list[float] = []
    right_edges: list[float] = []
    object_indices: list[int] = []
    seed_runs: list[dict] = []
    orig_lines: list[dict] = []

    for li, line in enumerate(lines):
        if line.baseline_y is not None:
            baselines.append(line.baseline_y)
            first_baseline_y = max(first_baseline_y, line.baseline_y)
        column_left = min(column_left, line.x)
        right_edges.append(line.x + line.w)
        line_runs = line.runs
        per_run, anchor_x = line_advances(line)
        line_frags: list[dict] = []
        for ri, r in enumerate(line_runs):
            for oi in (r.object_indices or []):
                object_indices.append(oi)
            orig = r.text or ""
            text = orig
            if li > 0 and ri == 0 and seed_runs:
                prev = seed_runs[-1]["text"]
                if prev and not prev.endswith((" ", "\t")) and not text.startswith((" ", "\t")):
                    text = " " + text
            adv = per_run.get(ri)
            fc = getattr(r, "fill_color", None) or getattr(r, "color", None) or "#000000"
            safe_color = fc if fc and fc != "#ffffff00" else "#000000"
            r_size = max(r.font_size or 0, r.h or 0) or eff_size
            seed_runs.append({
                "text": text, "fid": r.internal_font_id or "", "size": r_size,
                "color": safe_color, "fontName": r.font_name or "", "advances": adv,
            })
            if orig:
                line_frags.append({
                    "text": orig, "fid": r.internal_font_id or "", "size": r_size,
                    "color": safe_color, "advances": adv,
                })
        if line_frags:
            orig_lines.append({"frags": line_frags, "x": anchor_x, "baselineY": line.baseline_y or 0})

    max_right = max(right_edges) if right_edges else column_left + 250
    column_right = max(max_right, column_left + 250) if column_left != float("inf") else 250
    deltas = [baselines[i] - baselines[i + 1] for i in range(len(baselines) - 1)]
    dom_size = next((r["size"] for r in seed_runs if r["text"].strip()), 12.0)
    leading = abs(median(deltas)) if deltas else dom_size * 1.2
    col_w = column_right - column_left
    align = "left"
    if len(lines) >= 2:
        reaching = sum(1 for i in range(len(right_edges) - 1) if right_edges[i] >= column_right - col_w * 0.04)
        if reaching >= (len(lines) - 1) * 0.7:
            align = "justify"

    return {
        "columnLeft": column_left, "columnRight": column_right,
        "firstBaselineY": first_baseline_y, "leading": leading,
        "oldLineCount": len(lines), "align": align,
        "objectIndices": object_indices, "seedRuns": seed_runs, "origLines": orig_lines,
        "domSize": dom_size,
    }


def extract_flat_runs(seed_runs: list[dict], dominant_fid: str, dom_size: float, dom_color: str) -> list[dict]:
    """Mirror extractFlatRuns when DOM matches seedRuns exactly (edit mode entry)."""
    out: list[dict] = []
    for r in seed_runs:
        if not r["text"]:
            continue
        frag = {
            "text": r["text"], "internalFontId": r["fid"] or dominant_fid,
            "fontSize": r["size"], "color": r["color"],
        }
        adv = r.get("advances")
        if adv and len(adv) == len(r["text"]):
            frag["advances"] = adv
        out.append(frag)
    return out


def build_reflow_data(layout: dict, runs: list[dict], orig_lines: list[dict] | None, para_id: str) -> dict:
    line_position = {}
    if layout["origLines"]:
        line_position = {
            "lineX": [l["x"] for l in layout["origLines"]],
            "lineBaselineY": [l["baselineY"] for l in layout["origLines"]],
        }
    lines_data = {}
    if orig_lines:
        lines_data = {
            "lines": [
                [
                    {
                        "text": f["text"], "internalFontId": f["fid"],
                        "fontSize": f["size"], "color": f["color"],
                        **({"advances": f["advances"]} if f.get("advances") else {}),
                    }
                    for f in l["frags"]
                ]
                for l in orig_lines
            ],
        }
    dominant_fid = next((r["fid"] for r in layout["seedRuns"] if r["text"].strip() and r["fid"]), "")
    dom_run = next((r for r in layout["seedRuns"] if r["text"].strip() and r["fid"] == dominant_fid), None)
    dom_size = (dom_run or layout["seedRuns"][0])["size"] if layout["seedRuns"] else 12.0
    dom_color = (dom_run or layout["seedRuns"][0])["color"] if layout["seedRuns"] else "#000000"
    out_runs = runs if runs else [{"text": " ", "internalFontId": dominant_fid, "fontSize": dom_size, "color": "#000000"}]
    return {
        "objectIndices": layout["objectIndices"],
        "runs": out_runs,
        **line_position,
        **lines_data,
        "columnLeft": layout["columnLeft"],
        "columnRight": layout["columnRight"],
        "pushColumnLeft": layout["columnLeft"],
        "firstBaselineY": layout["firstBaselineY"],
        "leading": layout["leading"],
        "oldLineCount": layout["oldLineCount"],
        "align": layout["align"],
        "paraId": para_id,
    }


def paragraph_snapshot_from_extraction(para, layout: dict) -> dict:
    """Original PDF paragraph metrics from extraction + derived advances."""
    lines = para.lines
    all_text = "".join(r.text for ln in lines for r in ln.runs)
    runs_detail = []
    for ln in lines:
        per_run, anchor_x = line_advances(ln)
        for ri, r in enumerate(ln.runs):
            adv = per_run.get(ri)
            runs_detail.append({
                "text": r.text,
                "fontId": r.internal_font_id,
                "fontName": r.font_name,
                "fontSize": max(r.font_size or 0, r.h or 0),
                "color": getattr(r, "fill_color", None) or getattr(r, "color", None) or "#000000",
                "glyphCount": len(r.glyphs),
                "advanceCount": len(adv) if adv else 0,
                "advances": adv,
                "bbox": [r.x, r.y, r.w, r.h],
                "lineWidth": sum(adv) if adv else r.w,
            })
    line_widths = [ln.w for ln in lines]
    orig_line_texts = ["".join(r.text for r in ln.runs) for ln in lines]
    return {
        "stage": "ORIGINAL_PDF_EXTRACTION",
        "text": all_text,
        "fontId": runs_detail[0]["fontId"] if runs_detail else "",
        "fontName": runs_detail[0]["fontName"] if runs_detail else "",
        "fontSize": layout["domSize"],
        "color": runs_detail[0]["color"] if runs_detail else "#000000",
        "glyphCount": sum(len(r.glyphs) for ln in lines for r in ln.runs),
        "advanceCount": sum(rd["advanceCount"] for rd in runs_detail),
        "advances": [a for rd in runs_detail for a in (rd["advances"] or [])],
        "boundingBox": [para.x, para.y, para.w, para.h],
        "paragraphWidth": layout["columnRight"] - layout["columnLeft"],
        "lineWidths": line_widths,
        "baseline": layout["firstBaselineY"],
        "ascent": None,
        "descent": None,
        "lineHeight": layout["leading"],
        "origLines": orig_line_texts,
        "runs": runs_detail,
        "origLinesProvided": False,
    }


def snapshot_from_payload(stage: str, layout: dict, runs: list[dict], orig_lines: list[dict] | None) -> dict:
    all_text = "".join(r["text"] for r in runs)
    dom = runs[0] if runs else {}
    adv_flat = []
    for r in runs:
        if r.get("advances"):
            adv_flat.extend(r["advances"])
    return {
        "stage": stage,
        "text": all_text,
        "fontId": dom.get("internalFontId", ""),
        "fontName": layout["seedRuns"][0]["fontName"] if layout["seedRuns"] else "",
        "fontSize": dom.get("fontSize", layout["domSize"]),
        "color": dom.get("color", "#000000"),
        "glyphCount": len(all_text),
        "advanceCount": len(adv_flat),
        "advances": adv_flat,
        "boundingBox": None,
        "paragraphWidth": layout["columnRight"] - layout["columnLeft"],
        "lineWidths": None,
        "baseline": layout["firstBaselineY"],
        "ascent": None,
        "descent": None,
        "lineHeight": layout["leading"],
        "origLines": ["".join(f["text"] for f in l["frags"]) for l in orig_lines] if orig_lines else None,
        "runs": [
            {
                "text": r["text"],
                "fontId": r.get("internalFontId"),
                "fontSize": r.get("fontSize"),
                "color": r.get("color"),
                "advanceCount": len(r.get("advances") or []),
                "advances": r.get("advances"),
                "advanceSource": "frontend_origin_x_derived" if r.get("advances") else "missing",
            }
            for r in runs
        ],
        "origLinesProvided": orig_lines is not None,
    }


def snapshot_from_wasm_layout(stage: str, layout_json: dict, seed_font_name: str) -> dict:
    lines = layout_json.get("lines") or []
    all_text = "".join(ln.get("text", "") for ln in lines)
    adv_flat = []
    line_widths = []
    for ln in lines:
        adv = ln.get("adv") or []
        adv_flat.extend(adv)
        line_widths.append(sum(adv))
    dom_fs = lines[0].get("fontSize") if lines else None
    return {
        "stage": stage,
        "text": all_text,
        "fontId": "",
        "fontName": seed_font_name,
        "fontSize": dom_fs,
        "color": None,
        "glyphCount": len(all_text),
        "advanceCount": len(adv_flat),
        "advances": adv_flat,
        "boundingBox": None,
        "paragraphWidth": (layout_json.get("columnRight") or 0) - (layout_json.get("columnLeft") or 0),
        "lineWidths": line_widths,
        "baseline": lines[0].get("baselineY") if lines else None,
        "ascent": None,
        "descent": None,
        "lineHeight": None,
        "origLines": [ln.get("text", "") for ln in lines],
        "runs": [
            {
                "text": ln.get("text"),
                "lineX0": ln.get("x0"),
                "fontSize": ln.get("fontSize"),
                "advanceCount": len(ln.get("adv") or []),
                "advances": ln.get("adv"),
                "advanceSum": sum(ln.get("adv") or []),
                "advanceSource": "wasm_reflow_output",
            }
            for ln in lines
        ],
        "origLinesProvided": None,
    }


def para_signature(snap: dict) -> str:
    parts = [
        snap.get("text") or "",
        str(snap.get("fontId") or ""),
        str(snap.get("fontSize") or ""),
        json.dumps(snap.get("advances") or []),
        json.dumps(snap.get("lineWidths") or []),
        str(snap.get("baseline") or ""),
        json.dumps(snap.get("boundingBox") or []),
    ]
    return hashlib.sha256("".join(parts).encode()).hexdigest()


def main():
    pdf_path = ROOT / "corpus" / "basic" / "hello_world.pdf"
    if len(sys.argv) > 1:
        pdf_path = Path(sys.argv[1])
    doc = pdfengine.PdfDocument.load_from_file(str(pdf_path), "")
    model = doc.get_page(0).extract_document_model()
    para = model.paragraphs[0]
    layout = compute_layout(para)
    fonts = doc.get_fonts(0, 0)

    dominant_fid = next((r["fid"] for r in layout["seedRuns"] if r["text"].strip() and r["fid"]), "")
    dom_run = next((r for r in layout["seedRuns"] if r["text"].strip()), layout["seedRuns"][0])
    dom_size = dom_run["size"]
    dom_color = dom_run["color"]

    para_id = f"forensic-{hashlib.md5(str(pdf_path).encode()).hexdigest()[:8]}"

    # Stage snapshots
    original = paragraph_snapshot_from_extraction(para, layout)

    flat_entry = extract_flat_runs(layout["seedRuns"], dominant_fid, dom_size, dom_color)
    payload_entry = build_reflow_data(layout, flat_entry, layout["origLines"], para_id)
    edit_entry = snapshot_from_payload("EDIT_MODE_ENTRY_PAYLOAD", layout, flat_entry, layout["origLines"])

    # After one character
    flat_one = []
    for i, r in enumerate(flat_entry):
        if i == 0:
            flat_one.append({**r, "text": r["text"] + "x"})
        else:
            flat_one.append(r)
    payload_one = build_reflow_data(layout, flat_one, None, para_id)
    edit_one = snapshot_from_payload("AFTER_ONE_CHAR_PAYLOAD", layout, flat_one, None)

    out = {
        "pdf": str(pdf_path),
        "fonts": [
            {
                "internalFontId": f.internal_font_id,
                "fontName": f.font_name,
                "ascent": f.ascent,
                "descent": f.descent,
                "capHeight": getattr(f, "cap_height", None),
                "unitsPerEm": getattr(f, "units_per_em", None),
            }
            for f in fonts
        ],
        "layout": layout,
        "original": original,
        "editEntry": edit_entry,
        "editOne": edit_one,
        "payloads": {
            "editEntry": {
                "version": "1.0",
                "operations": [{
                    "id": "forensic-entry", "type": "reflow_paragraph", "pageIndex": 0,
                    "data": payload_entry,
                }],
            },
            "editOne": {
                "version": "1.0",
                "operations": [{
                    "id": "forensic-one", "type": "reflow_paragraph", "pageIndex": 0,
                    "data": payload_one,
                }],
            },
        },
        "signatures": {
            "original": para_signature(original),
            "editEntryPayload": para_signature(edit_entry),
            "editOnePayload": para_signature(edit_one),
        },
    }
    out_path = Path(__file__).resolve().parent / "forensic_extract_out.json"
    out_path.write_text(json.dumps(out, indent=2), encoding="utf-8")
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()