#!/usr/bin/env python3
"""Two-phase PDF heading-promotion CLI (file-remediation post-tagging stage,
docs/plan-heading-promote.md).

Runs on an already-tagged PDF (E3 output) that has zero (or too few) /Hn
struct elements -- veraPDF checks structural VALIDITY, never whether
headings exist, so a document can pass 106/106 while every element is
tagged /P (confirmed live: chapter10-remediated.pdf, 854 /P, 0 /H). The
Ally checker correctly flags "This PDF does not have headings" on exactly
this kind of file.

Doctrine mirrors figure_alt.py: SCRIPT clusters, AI (or a human) assigns
levels, SCRIPT applies + verifies. This module never decides a cluster is
a heading or picks its level -- font size/boldness/color is evidence, not a
verdict (a large word mid-sentence isn't a heading; a colored label might
not be one either).

Phase 1 (--build-manifest, script): pdf_tree_edit.heading_candidates()
clusters every non-heading leaf struct element's text by font signature,
excludes the document's body-text signature, and returns one record per
occurrence with "level": null. Written to heading-manifest.json alongside a
font_size_rank summary (distinct cluster signatures, largest first) so the
agent can see relative sizing at a glance.

Phase 2 (agent, no script involved): Claude (or another session AI) reads
the manifest and font_size_rank summary and assigns each occurrence a
"level" (1-6) or leaves it null ("not a heading") -- batching by cluster_id
is expected (a whole cluster usually gets one level), but each occurrence
is annotated individually so an outlier inside a cluster can be overridden.

Phase 3 (--apply, script): backs up the PDF first (abort if the backup
itself fails), promotes every occurrence with a non-null level via
pdf_tree_edit.promote_headings (which also normalizes level jumps), saves,
and re-verifies: veraPDF must still be compliant AND at least one heading
must be present. Any failure restores the pre-apply backup and reports it
-- never leaves a half-written file claimed done, and never claims success
from an all-null manifest.

Usage:
    python heading_promote.py --pdf tagged.pdf --build-manifest [--outdir DIR]
    python heading_promote.py --pdf tagged.pdf --apply --manifest DIR/heading-manifest.json
"""

from __future__ import annotations

import argparse
import json
import re
import shutil
import sys
from datetime import datetime, timezone
from pathlib import Path

SCRIPTS_DIR = Path(__file__).resolve().parent
if str(SCRIPTS_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPTS_DIR))

import pikepdf

import pdf_tree_edit
import verapdf_gate

MANIFEST_NAME = "heading-manifest.json"
_HEADING_RE = re.compile(r"^/H\d+$")


def _default_outdir(pdf_path):
    p = Path(pdf_path)
    return p.with_name(f"{p.stem}.headings")


def _font_size_rank(candidates):
    """[{cluster_id, size, bold, rgb, font, count}, ...] -- one row per
    distinct cluster_id, sorted largest-size first (the same order
    heading_candidates() already assigns cluster_id in, so this is a
    de-duplicated summary, not a re-sort)."""
    seen = {}
    for c in candidates:
        cid = c["cluster_id"]
        if cid not in seen:
            seen[cid] = {
                "cluster_id": cid, "size": c["size"], "bold": c["bold"],
                "rgb": c["rgb"], "font": c["font"], "count": c["cluster_count"],
            }
    return [seen[cid] for cid in sorted(seen)]


def build_manifest(pdf_path, outdir=None):
    """Phase 1. Returns (manifest_path, manifest_dict)."""
    pdf_path = Path(pdf_path)
    outdir = Path(outdir) if outdir else _default_outdir(pdf_path)
    outdir.mkdir(parents=True, exist_ok=True)

    pdf = pikepdf.open(str(pdf_path))
    try:
        candidates = pdf_tree_edit.heading_candidates(pdf, str(pdf_path))
    finally:
        pdf.close()

    manifest = {
        "pdf": str(pdf_path),
        "candidates": candidates,
        "font_size_rank": _font_size_rank(candidates),
    }
    manifest_path = outdir / MANIFEST_NAME
    manifest_path.write_text(
        json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")
    return manifest_path, manifest


def _backup_pdf(pdf_path, outdir):
    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    backup_dir = Path(outdir) / "backups"
    backup_dir.mkdir(parents=True, exist_ok=True)
    backup_path = backup_dir / f"{Path(pdf_path).stem}.pre-apply-{ts}.pdf"
    shutil.copy2(pdf_path, backup_path)
    return backup_path


def _count_headings(pdf):
    return sum(1 for e in pdf_tree_edit.walk_elems(pdf, None)
               if _HEADING_RE.match(str(e.get("/S", ""))))


def apply_manifest(pdf_path, manifest_path):
    """Phase 3. Returns a result dict; never raises on a refused/unresolved
    occurrence or a failed verify -- both are reported, not exceptions."""
    pdf_path = Path(pdf_path)
    manifest_path = Path(manifest_path)
    if not pdf_path.exists():
        return {"ok": False, "error": f"pdf not found: {pdf_path}"}
    if not manifest_path.exists():
        return {"ok": False, "error": f"manifest not found: {manifest_path}"}

    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    candidates = manifest.get("candidates", [])
    assignments = [
        {"page": c.get("page"), "mcids": c.get("mcids"), "level": c.get("level")}
        for c in candidates if c.get("level") is not None
    ]
    skipped_null = sum(1 for c in candidates if c.get("level") is None)

    try:
        backup_path = _backup_pdf(pdf_path, manifest_path.parent)
    except OSError as exc:
        return {"ok": False, "error": f"backup failed — aborting, no changes made: {exc}"}

    pdf = pikepdf.open(str(pdf_path), allow_overwriting_input=True)
    try:
        promotion = pdf_tree_edit.promote_headings(pdf, assignments)
        pdf.save(str(pdf_path))
    finally:
        pdf.close()

    vera = verapdf_gate.run_verapdf(str(pdf_path))
    review = pikepdf.open(str(pdf_path))
    try:
        heading_count = _count_headings(review)
    finally:
        review.close()

    ok = bool(vera.get("compliant")) and heading_count > 0
    result = {
        "ok": ok,
        "pdf": str(pdf_path),
        "backup": str(backup_path),
        "promoted": promotion["promoted"],
        "unresolved": promotion["unresolved"],
        "skipped_null_level": skipped_null,
        "verapdf_compliant": vera.get("compliant"),
        "heading_count": heading_count,
    }
    if not ok:
        shutil.copy2(backup_path, pdf_path)
        result["restored"] = True
        result["error"] = (
            "post-apply verify failed (veraPDF not compliant and/or zero "
            "headings resulted) — restored the pre-apply backup; nothing "
            "was kept. Assign levels to the reported candidates and re-run "
            "--apply."
        )
    return result


def main():
    parser = argparse.ArgumentParser(
        description="Two-phase PDF heading-promotion CLI (build-manifest / apply)")
    parser.add_argument("--pdf", required=True)
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--build-manifest", action="store_true")
    group.add_argument("--apply", action="store_true")
    parser.add_argument("--outdir", help="Manifest output dir for "
                         "--build-manifest (default: <pdf_stem>.headings/ beside the PDF)")
    parser.add_argument("--manifest", help="Path to the annotated manifest JSON (--apply)")
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    if args.build_manifest:
        manifest_path, manifest = build_manifest(args.pdf, args.outdir)
        result = {
            "ok": True, "manifest": str(manifest_path),
            "candidates": len(manifest["candidates"]),
            "clusters": len(manifest["font_size_rank"]),
        }
        if args.json:
            print(json.dumps(result, indent=2))
        else:
            print(f"Manifest: {manifest_path} "
                  f"({result['candidates']} candidate(s), {result['clusters']} cluster(s))")
        return

    if not args.manifest:
        parser.error("--apply requires --manifest FILE")
    result = apply_manifest(args.pdf, args.manifest)
    print(json.dumps(result, indent=2, default=str))
    sys.exit(0 if result.get("ok") else 1)


if __name__ == "__main__":
    main()
