"""
AI Room Designer — Stage 3 generation orchestrator: backend + QC v2 gate +
AUTOMATED RETRY (estimate phase 13, completed here).

Flow for one job (job dir prepared by pipeline.prepare / the app server):
  attempt 0   -> backend.edit(before, locked structure prompt)
  score it with the QC v2 architectural-mask gate
  FAIL/PARTIAL -> regenerate with the FRAMING LOCK appended
                  (the fix that rescued Budget Refresh in Stage 2),
                  up to config.max_retries() retries (default 2)
  end         -> the BEST attempt (gate rank, then arch, then layout) is
                 finished through pipeline.finish (composites + qc + plan);
                 if no attempt passed, the best is still surfaced and the
                 job is FLAGGED in generation.json + qc.json (gate_final).

Usage:
  python3 pipeline/generate.py <job-dir> [--provider gemini|mock|standin]
"""
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
import config  # noqa: E402
from backends import get_backend, BackendError  # noqa: E402
from modes import build_prompt  # noqa: E402
from qc_score import score, verdict  # noqa: E402
import pipeline as pl  # noqa: E402
from PIL import Image  # noqa: E402

FRAMING_LOCK = (
    " CRITICAL FRAMING LOCK: keep the EXACT same camera position, framing, "
    "crop and field of view as the input photo. Do not zoom, widen, shift, "
    "pan or reframe the shot in any way."
)


def _rank(v):
    return 2 if v.startswith("PASS") else 1 if v.startswith("PARTIAL") else 0


def generate(job_dir, provider=None, backend=None):
    job = Path(job_dir)
    meta = json.loads((job / "job.json").read_text())
    mode = meta["mode"]
    before = job / "before.jpg"
    backend = backend or get_backend(provider)
    attempts_dir = job / "attempts"
    attempts_dir.mkdir(exist_ok=True)

    base_prompt = build_prompt(mode)
    attempts, best = [], None
    for i in range(config.max_retries() + 1):
        prompt = base_prompt if i == 0 else base_prompt + FRAMING_LOCK
        raw = attempts_dir / f"attempt{i}.png"
        try:
            backend.edit(str(before), prompt, str(raw))
        except BackendError as e:
            attempts.append({"attempt": i, "error": str(e)})
            break
        # Normalise onto the before frame, exactly like pipeline.finish does.
        norm = attempts_dir / f"attempt{i}_matched.png"
        Image.open(raw).convert("RGB").resize(Image.open(before).size).save(norm)
        s = score(str(before), str(norm))
        v = verdict(s)
        rec = {"attempt": i, "image": str(norm), "scores": s, "verdict": v,
               "framing_lock": i > 0,
               "usage": getattr(backend, "last_usage", None)}
        attempts.append(rec)
        key = (_rank(v), s["arch_edge_corr"], s["layout_corr"])
        if best is None or key > best[0]:
            best = (key, rec)
        if v.startswith("PASS"):
            break

    gen = {"provider": getattr(backend, "name", provider or config.provider()),
           "model": getattr(backend, "model", None),
           "attempts": attempts,
           "retries_used": max(0, len([a for a in attempts if 'scores' in a]) - 1)}
    if best is None:
        gen.update({"status": "error", "flagged": True,
                    "flag": attempts[-1].get("error", "no attempt produced") if attempts else "no attempts"})
        (job / "generation.json").write_text(json.dumps(gen, indent=2))
        return gen

    best_rec = best[1]
    passed = best_rec["verdict"].startswith("PASS")
    gen.update({"status": "done" if passed else "done-flagged",
                "flagged": not passed,
                "flag": None if passed else
                        "no attempt passed the QC gate; best attempt surfaced for review",
                "best_attempt": best_rec["attempt"],
                "final_verdict": best_rec["verdict"]})

    # Finish the best attempt: composites + qc.json + plan.json (existing code).
    pl.finish(str(job), best_rec["image"])
    qc = json.loads((job / "qc.json").read_text())
    qc["generation"] = {k: gen[k] for k in ("provider", "model", "retries_used", "flagged")}
    qc["gate_final"] = ("PASS" if passed else "FLAGGED — surfaced best attempt") + \
                       f" (attempt {best_rec['attempt']}, {gen['retries_used']} retries)"
    (job / "qc.json").write_text(json.dumps(qc, indent=2))
    (job / "generation.json").write_text(json.dumps(gen, indent=2))
    return gen


if __name__ == "__main__":
    prov = None
    if "--provider" in sys.argv:
        prov = sys.argv[sys.argv.index("--provider") + 1]
    out = generate(sys.argv[1], provider=prov)
    print(json.dumps(out, indent=2))