#!/usr/bin/env python3
"""Record and copy a codex-ppt slide image generated by a subagent."""

from __future__ import annotations

import argparse
import json
import re
import shutil
from pathlib import Path

from slide_run_state import (
    deck_dir_from_target,
    ensure_file,
    find_slide,
    locked_jobs,
    now_iso,
    rel_to_deck,
    resolve_deck_path,
    set_run_status,
    sha256_file,
    update_jobs_run_status,
)


FORBIDDEN_BACKEND_TERMS = (
    "pillow",
    "html",
    "svg",
    "canvas",
    "python-pptx",
    "pptxgenjs",
    "manual",
    "script render",
)


def _validate_backend(backend: str) -> str:
    value = backend.strip()
    if not value:
        raise SystemExit("--backend-used must not be empty")
    lowered = value.lower()
    for term in FORBIDDEN_BACKEND_TERMS:
        if term in lowered:
            raise SystemExit(f"Forbidden backend for final slide image: {backend}")
    if "image" not in lowered and "built-in" not in lowered:
        raise SystemExit(
            "--backend-used must name the selected image backend, such as "
            "`built-in image tool` or `scripts/image_gen.py`."
        )
    return value


def _backend_family(value: str) -> str:
    lowered = value.lower()
    if "image_gen.py" in lowered or "scripts/image_gen" in lowered:
        return "cli-api-fallback"
    if "cli" in lowered or "api" in lowered or "fallback" in lowered:
        return "cli-api-fallback"
    if "built-in" in lowered or "builtin" in lowered or "image_generate" in lowered:
        return "built-in-image-tool"
    if "image_gen" in lowered:
        return "built-in-image-tool"
    return ""


def _normalized_backend_label(value: str) -> str:
    return re.sub(r"[^a-z0-9]+", "", value.lower())


def _method_backend_labels(method: object) -> list[str]:
    if not isinstance(method, dict):
        return []
    labels: list[str] = []
    for key in ("backend_used", "selected_backend", "backend", "tool_name"):
        value = method.get(key)
        if isinstance(value, str) and value.strip():
            labels.append(value.strip())
    return labels


def _load_prompt_job(deck_dir: Path, slide: dict) -> dict:
    job_ref = slide.get("job")
    if not isinstance(job_ref, str) or not job_ref.strip():
        return {}
    path = resolve_deck_path(deck_dir, job_ref)
    if not path.exists():
        return {}
    data = json.loads(path.read_text(encoding="utf-8"))
    return data if isinstance(data, dict) else {}


def _expected_backend_labels(jobs: dict, slide: dict, prompt_job: dict) -> list[str]:
    labels: list[str] = []
    for value in (jobs.get("selected_backend"), prompt_job.get("expected_backend")):
        if isinstance(value, str) and value.strip():
            labels.append(value.strip())
    labels.extend(_method_backend_labels(jobs.get("sample_generation_method")))
    labels.extend(_method_backend_labels(prompt_job.get("sample_generation_method")))
    unique_labels: list[str] = []
    for label in labels:
        if label not in unique_labels:
            unique_labels.append(label)
    return unique_labels


def _matched_expected_backend(backend_used: str, expected_labels: list[str]) -> str | None:
    if not expected_labels:
        return None
    used_normalized = _normalized_backend_label(backend_used)
    used_family = _backend_family(backend_used)
    for label in expected_labels:
        if _normalized_backend_label(label) == used_normalized:
            return label
        expected_family = _backend_family(label)
        if expected_family and expected_family == used_family:
            return label
    raise SystemExit(
        "Backend mismatch: worker reported "
        f"`{backend_used}`, but slide job expects one of: {', '.join(expected_labels)}"
    )


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("deck", help="Deck directory or slide_jobs.json")
    parser.add_argument("--slide", required=True, help="slide_01 or 1")
    parser.add_argument("--agent-id", required=True)
    parser.add_argument("--backend-used", required=True)
    parser.add_argument("--selected-source", required=True, help="Generated image selected by the worker.")
    parser.add_argument("--qa-note", required=True)
    args = parser.parse_args()

    deck_dir = deck_dir_from_target(args.deck)
    with locked_jobs(deck_dir) as jobs:
        slide = find_slide(jobs, args.slide)
        if slide.get("status") != "dispatched":
            raise SystemExit(f"{slide['slide_id']} must be dispatched before result recording; got {slide.get('status')}")
        dispatch = slide.get("dispatch") or {}
        if dispatch.get("agent_id") != args.agent_id:
            raise SystemExit(
                f"Agent id mismatch for {slide['slide_id']}: dispatch={dispatch.get('agent_id')} result={args.agent_id}"
            )

        backend_used = _validate_backend(args.backend_used)
        prompt_job = _load_prompt_job(deck_dir, slide)
        expected_backend_labels = _expected_backend_labels(jobs, slide, prompt_job)
        matched_expected_backend = _matched_expected_backend(backend_used, expected_backend_labels)
        source = ensure_file(Path(args.selected_source).expanduser().resolve(), "selected source image")
        out_ref = slide.get("out") or f"origin_image/{slide['slide_id']}.png"
        target = resolve_deck_path(deck_dir, out_ref)
        try:
            target.relative_to(deck_dir)
        except ValueError as exc:
            raise SystemExit(f"Final slide image must live inside deck dir: {target}") from exc
        target.parent.mkdir(parents=True, exist_ok=True)
        if source != target:
            shutil.copy2(source, target)

        slide["result"] = {
            "agent_id": args.agent_id,
            "backend_used": backend_used,
            "selected_source": str(source),
            "selected_source_sha256": sha256_file(source),
            "final_image": rel_to_deck(deck_dir, target),
            "final_image_sha256": sha256_file(target),
            "expected_backend": matched_expected_backend,
            "expected_backend_labels": expected_backend_labels,
            "sample_generation_method_matched": bool(matched_expected_backend),
            "qa_note": args.qa_note,
            "recorded_at": now_iso(),
        }
        slide["status"] = "recorded"
        update_jobs_run_status(jobs)
        run_status = jobs.get("run_status")
        slide_id = slide["slide_id"]
    if run_status == "slides_recorded":
        set_run_status(deck_dir, "slides_recorded", "all slide images recorded")

    print(json.dumps({"slide_id": slide_id, "status": "recorded"}, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
