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

Runs on the ALREADY-TAGGED output of remediate_pdf.py -- never re-tags --
so figure indices stay stable across the phases (see the plan's "why a
standalone post-tagging stage" section).

Phase 1 (--build-manifest, script): render every /Figure struct element to
a PNG and write a JSON manifest carrying every pdf_tree_edit.list_figures()
context field (current alt, is_placeholder, caption_candidate, width,
height, surrounding_text, reading_order_index, repeat_count, in_link,
in_heading, mcids/page) plus the PNG path. Every entry's "disposition" is
null -- this script NEVER concludes decorative/content; it only surfaces
evidence, including size, which is a datum here, never a decider (a 14x14
image can be a status icon; see AC6 in the test suite).

Phase 2 (agent, no script involved): Claude (or another session AI) opens
each PNG alongside its context fields and edits the manifest in place,
writing "disposition" ("describe" | "decorative"), "alt" (required,
non-empty, non-placeholder when "describe"), and a one-clause "reason".
Mirrors audit_image_alt.py / apply_image_alt.py's manifest-for-agent shape
(the Canvas-HTML sibling of this PDF workflow).

Phase 3 (--apply, script): backs up the PDF bytes first (abort if the
backup itself fails), then applies every annotated figure -- "describe" ->
pdf_tree_edit.set_alts, "decorative" -> pdf_tree_edit.demote_figures_to_artifact
-- saves in place, and re-verifies: veraPDF must still be compliant AND
list_figures() must show zero remaining placeholder alts. Any failure
restores the pre-apply backup and reports it -- never leaves a half-written
file claimed done. A figure left at disposition: null, or a "describe" whose
alt is empty/whitespace/placeholder, is refused outright (reported under
"needs_manual"), never guessed.

Usage:
    python figure_alt.py --pdf tagged.pdf --build-manifest [--outdir DIR]
    python figure_alt.py --pdf tagged.pdf --apply --manifest DIR/figure-alt-manifest.json
"""

from __future__ import annotations

import argparse
import json
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 fitz  # PyMuPDF
import pikepdf

import pdf_tree_edit
import verapdf_gate

MANIFEST_NAME = "figure-alt-manifest.json"
RENDER_DPI = 150
CROP_PAD_PT = 8


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


def _render_figure_png(fitz_doc, page_1based, objid, out_path):
    """Best-effort render of the figure's placed image, padded and cropped
    to its page rect. Falls back to rendering the whole page when the
    image's placement rect can't be resolved (e.g. a vector-wrapped figure,
    or an unresolvable page) -- a PNG is still produced either way (AC1)."""
    page = None
    if page_1based is not None and 1 <= page_1based <= fitz_doc.page_count:
        page = fitz_doc[page_1based - 1]
    elif fitz_doc.page_count:
        page = fitz_doc[0]
    if page is None:
        return False

    rect = None
    if objid is not None:
        try:
            rects = page.get_image_rects(objid)
            if rects:
                rect = rects[0]
        except Exception:  # noqa: BLE001
            rect = None

    if rect is not None:
        padded = fitz.Rect(rect.x0 - CROP_PAD_PT, rect.y0 - CROP_PAD_PT,
                            rect.x1 + CROP_PAD_PT, rect.y1 + CROP_PAD_PT)
        padded &= page.rect
        pix = page.get_pixmap(clip=padded, dpi=RENDER_DPI)
    else:
        pix = page.get_pixmap(dpi=RENDER_DPI)
    pix.save(str(out_path))
    return True


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

    pdf = pikepdf.open(str(pdf_path))
    try:
        records = pdf_tree_edit.list_figures(pdf, path=str(pdf_path))
        image_info = pdf_tree_edit.figures_image_info(pdf)
    finally:
        pdf.close()

    fitz_doc = fitz.open(str(pdf_path))
    try:
        manifest = []
        for rec, info in zip(records, image_info):
            png_path = figures_dir / f"fig{rec['index']}.png"
            objid = info.get("objid") if info else None
            _render_figure_png(fitz_doc, rec.get("page"), objid, png_path)
            entry = dict(rec)
            entry["png"] = str(png_path)
            # Load-bearing: the script NEVER concludes decorative/content —
            # every field above is evidence for the AI, disposition is null
            # for every figure regardless of size/position/repeat_count.
            entry["disposition"] = None
            entry["alt"] = None
            entry["reason"] = None
            manifest.append(entry)
    finally:
        fitz_doc.close()

    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 _partition_dispositions(manifest):
    """Split manifest entries into (describe_alt_by_index, decorative_indices,
    needs_manual) per the honesty rules -- null disposition or an empty/
    placeholder "describe" alt is refused, never guessed."""
    describe: dict = {}
    decorative: list = []
    needs_manual: list = []

    for entry in manifest:
        idx = entry.get("index")
        disposition = entry.get("disposition")
        if disposition is None:
            needs_manual.append({"index": idx, "reason": "disposition not annotated"})
            continue
        if disposition == "describe":
            alt = entry.get("alt")
            if not alt or not alt.strip() or pdf_tree_edit.is_placeholder_alt(alt):
                needs_manual.append({
                    "index": idx,
                    "reason": "describe with empty/whitespace/placeholder alt — rejected",
                })
                continue
            describe[idx] = alt
        elif disposition == "decorative":
            decorative.append(idx)
        else:
            needs_manual.append({
                "index": idx, "reason": f"unknown disposition {disposition!r}",
            })
    return describe, decorative, needs_manual


def apply_manifest(pdf_path, manifest_path):
    """Phase 3. Returns a result dict; never raises on a refused figure 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"))
    describe, decorative, needs_manual = _partition_dispositions(manifest)

    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:
        applied_alts = pdf_tree_edit.set_alts(pdf, describe) if describe else {}
        demoted = pdf_tree_edit.demote_figures_to_artifact(pdf, decorative) if decorative else 0
        pdf.save(str(pdf_path))
    finally:
        pdf.close()

    vera = verapdf_gate.run_verapdf(str(pdf_path))
    review = pikepdf.open(str(pdf_path))
    try:
        figures_after = pdf_tree_edit.list_figures(review, path=str(pdf_path))
    finally:
        review.close()
    placeholders_after = [f for f in figures_after if f["is_placeholder"]]

    ok = bool(vera.get("compliant")) and not placeholders_after
    result = {
        "ok": ok,
        "pdf": str(pdf_path),
        "backup": str(backup_path),
        "applied": {"described": applied_alts, "demoted_count": demoted},
        "needs_manual": needs_manual,
        "verapdf_compliant": vera.get("compliant"),
        "placeholders_remaining": len(placeholders_after),
    }
    if not ok:
        shutil.copy2(backup_path, pdf_path)
        result["restored"] = True
        result["error"] = (
            "post-apply verify failed (veraPDF not compliant and/or placeholder "
            "alt text remains) — restored the pre-apply backup; nothing was "
            "kept. Re-annotate the reported figures and re-run --apply."
        )
    return result


def main():
    parser = argparse.ArgumentParser(
        description="Two-phase PDF figure-alt 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/PNG output dir for "
                         "--build-manifest (default: <pdf_stem>.figure-alt/ beside the PDF)")
    parser.add_argument("--manifest", help="Path to the annotated manifest JSON (--apply)")
    parser.add_argument("--confirm", help="Optional label recorded for the run; this is a "
                         "local file edit (not a Canvas push), so no gating verb is required")
    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), "figures": len(manifest)}
        if args.json:
            print(json.dumps(result, indent=2))
        else:
            print(f"Manifest: {manifest_path} ({len(manifest)} figure(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()
