"""`uvw` command line: environment, weight hashing and frozen-config generation."""

from __future__ import annotations

import argparse
import json
import sys
from dataclasses import asdict
from pathlib import Path

from uvw import RUNS


def _cmd_env(args: argparse.Namespace) -> int:
    from uvw.provenance import environment

    print(json.dumps(environment(), indent=2))
    return 0


def _cmd_hash(args: argparse.Namespace) -> int:
    from uvw.models import load_pipeline, load_registry, weight_hashes

    variant = load_registry()[args.variant]
    print(json.dumps({"variant": variant.name, **weight_hashes(load_pipeline(variant))}, indent=2))
    return 0


def _cmd_generate(args: argparse.Namespace) -> int:
    """Generate for a JSONL file of {prompt_id, prompt, seed}; save images and a records file."""
    from uvw.generate import GenItem, generate, load_inference_config
    from uvw.models import load_pipeline, load_registry, weight_hashes
    from uvw.provenance import RunRecord, config_hash

    cfg = load_inference_config()
    if args.steps:
        cfg["num_inference_steps"] = args.steps
    variant = load_registry()[args.variant]
    rows = [json.loads(line) for line in Path(args.prompts).read_text().splitlines() if line.strip()]
    if args.limit:
        rows = rows[: args.limit]

    record = RunRecord(job="generate", variants=[variant.name], configs={"inference": config_hash(cfg)})
    out = Path(args.out or RUNS / record.run_id)
    (out / "images").mkdir(parents=True, exist_ok=True)
    pipe = load_pipeline(variant)
    record.weights = {f"{variant.name}/{k}": v for k, v in weight_hashes(pipe).items()}

    items = [GenItem(r["prompt_id"], r.get("text") or r["prompt"], int(r["seed"])) for r in rows]
    with open(out / "generations.jsonl", "w") as f:
        for g in generate(pipe, items, cfg, record.cost):
            name = f"{g.item.prompt_id}_s{g.item.seed}.png"
            if not args.no_images:
                g.image.save(out / "images" / name)
            f.write(json.dumps({**asdict(g.item), "image": name, "image_sha256": g.image_sha256,
                                "config_hash": g.config_hash, "variant": variant.name}) + "\n")
    record.outputs = [str(out / "generations.jsonl")]
    record.finish()
    print(record.write(out))
    return 0


def _cmd_splits(args: argparse.Namespace) -> int:
    """Build sources → template-level splits → leakage report → data/splits/."""
    from uvw import ROOT
    from uvw.data import leakage
    from uvw.data.sources import build_all
    from uvw.data.splits import assign, write_splits

    out = Path(args.out or ROOT / "data" / "splits")
    out.mkdir(parents=True, exist_ok=True)
    df, excluded = build_all()
    emb = leakage.clip_text_embeddings(df.text.tolist())
    df["group_id"] = leakage.group_templates(df, emb, threshold=args.threshold).values
    df["_row"] = range(len(df))
    df = assign(df)
    emb, df = emb[df._row.to_numpy()], df.drop(columns="_row")
    overlap = leakage.template_overlap(df)
    exact = leakage.exact_duplicates(df)
    near = leakage.near_duplicates(df, emb, threshold=args.threshold)
    near.to_csv(out / "near_duplicates.csv", index=False)
    report = {"rows": len(df), "excluded": len(excluded), "template_overlap": len(overlap),
              "exact_cross_split_duplicates": len(exact), "near_cross_split_duplicates": len(near),
              "threshold": args.threshold}
    if len(exact) or len(near):
        raise SystemExit(f"Cross-split duplicates remain: {report}")
    if len(overlap):
        print(overlap.head())
        raise SystemExit("Template overlap across splits: refusing to write.")
    manifest = write_splits(df, out)
    (out / "excluded.jsonl").write_text("".join(json.dumps(e) + "\n" for e in excluded))
    exact.to_csv(out / "exact_duplicates.csv", index=False)
    (out / "leakage_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))
    for name, info in sorted(manifest["files"].items()):
        print(f"{name:28s} rows={info['rows']:6d} templates={info['templates']:5d} {info['by_role']}")
    return 0


def _cmd_evaluate(args: argparse.Namespace) -> int:
    from uvw.evaluate import evaluate

    for concept in args.concepts:
        out = evaluate(args.variant, concept, args.split, n_seeds=args.seeds, limit=args.limit,
                       out=Path(args.out) / f"{args.variant}-{concept}-{args.split}" if args.out else None,
                       keep_images=args.keep_images, edits=Path(args.edits) if args.edits else None)
        print(out)
    return 0


def _cmd_edit_agent(args: argparse.Namespace) -> int:
    """Run the editor agent locally (ChatGPT plan via atbot) against a remote tool server."""
    from uvw import RUNS
    from uvw.harness.agent import run_editor
    from uvw.harness.server import RemoteSession

    session = RemoteSession(args.url, RUNS / "edits-local" / args.run_id)
    summary = run_editor(session, editor_model=args.editor, model_name=args.model_name)
    print(json.dumps(summary, indent=2))
    return 0


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(prog="uvw")
    sub = parser.add_subparsers(dest="command", required=True)
    sub.add_parser("env", help="Print the environment record")
    p = sub.add_parser("hash", help="Hash a variant's weights")
    p.add_argument("variant")
    p = sub.add_parser("generate", help="Generate images with the frozen inference config")
    p.add_argument("variant")
    p.add_argument("prompts", help="JSONL with prompt_id, text (or prompt), seed")
    p.add_argument("--out")
    p.add_argument("--limit", type=int)
    p.add_argument("--steps", type=int, help="Override steps (smoke tests only)")
    p.add_argument("--no-images", action="store_true", help="Keep hashes and scores only")
    p = sub.add_parser("splits", help="Build frozen dev/val/test prompt splits")
    p.add_argument("--out")
    p.add_argument("--threshold", type=float, default=0.95)
    p = sub.add_parser("evaluate", help="Generate + score a variant on concept splits")
    p.add_argument("variant")
    p.add_argument("concepts", nargs="+", choices=["nudity", "vangogh", "church", "coco"])
    p.add_argument("--split", default="dev", choices=["dev", "val"])  # test: evaluator only (003)
    p.add_argument("--seeds", type=int, default=1)
    p.add_argument("--limit", type=int, help="Max prompts per role (pilots)")
    p.add_argument("--out")
    p.add_argument("--keep-images", default="none", choices=["none", "benign", "all"])
    p.add_argument("--edits", help="edited_tensors.safetensors from a committed candidate")
    p = sub.add_parser("edit-agent", help="Run the editor agent against a tool server")
    p.add_argument("--url", default="http://127.0.0.1:8765")
    p.add_argument("--run-id", required=True)
    p.add_argument("--editor", default="gpt-6-astra")
    p.add_argument("--model-name", default="Stable Diffusion v1.4")
    args = parser.parse_args(argv)
    commands = {"edit-agent": _cmd_edit_agent, "evaluate": _cmd_evaluate, "env": _cmd_env, "hash": _cmd_hash, "generate": _cmd_generate, "splits": _cmd_splits}
    return commands[args.command](args)


if __name__ == "__main__":
    sys.exit(main())