"""`design-canvas` command line.

    design-canvas prepare [--train 200 --validation 60 --test 100]   build the task pack
    design-canvas serve [--port 8000] [--workers N]                OpenEnv server + editor (N processes, one URL)
    design-canvas mcp                                              stdio MCP server for Claude & co.
    design-canvas export-sft --split train --out sft.jsonl         expert trajectories
    design-canvas export-hf --cua artifacts/cua/cua.jsonl --out artifacts/hf/cua   cua-v1 dataset (Hub-ready)
    design-canvas smoke                                            offline end-to-end check

    design-canvas corpus build --out DIR          index all of FineEnvs/crello-bucket
    design-canvas corpus publish --out DIR        upload the index next to the data
    design-canvas corpus stats                    what the served corpus holds
    design-canvas evalset build --out FILE        freeze the stratified evaluation set
    design-canvas evalset verify FILE             check a frozen set against the corpus
    design-canvas evalset export FILE --out DIR   write it as an offline task pack
    design-canvas evalset subset FILE --name N --size K [--within lite]   nest a smaller subset in the set
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from pathlib import Path


def _pack_arg(p: argparse.ArgumentParser):
    p.add_argument("--pack", default=os.environ.get("DESIGN_CANVAS_PACK"),
                   help="task pack directory (default ~/.cache/design_canvas/pack)")


def cmd_prepare(args) -> int:
    from huggingface_hub import hf_hub_download

    from .core.pack import CRELLO_REPO, CRELLO_REVISION, DEFAULT_PACK, build_pack

    out = Path(args.pack or DEFAULT_PACK).expanduser()

    def dl(name):
        return hf_hub_download(CRELLO_REPO, name, repo_type="dataset", revision=CRELLO_REVISION)

    shards = {"train": "data/train-{:05d}-of-00031.parquet",
              "validation": "data/validation-{:05d}-of-00003.parquet",
              "test": "data/test-{:05d}-of-00004.parquet"}
    per = {"train": args.train, "validation": args.validation, "test": args.test}
    sources = {}
    for split, pattern in shards.items():
        if per[split] <= 0:
            continue
        # Each shard holds ~600 templates; about 60% pass the filters.
        need = max(1, min(args.max_shards, 1 + per[split] // 350))
        sources[split] = [dl(pattern.format(i)) for i in range(need)]
        print(f"{split}: {len(sources[split])} shard(s)")
    t = time.time()
    manifest = build_pack(out, sources, per, dl("resources/fonts.pickle"),
                          max_elements=args.max_elements)
    print(json.dumps(manifest, indent=2))
    print(f"pack ready at {out} in {time.time() - t:.0f}s")
    return 0


def cmd_serve(args) -> int:
    if args.pack:
        os.environ["DESIGN_CANVAS_PACK"] = str(Path(args.pack).expanduser())
    os.environ["PORT"] = str(args.port)
    from .server.app import main

    print(f"DesignGym on http://{args.host}:{args.port}  "
          f"(editor /editor/, playground /web/, MCP /mcp)")
    main(host=args.host, port=args.port, workers=args.workers)
    return 0


def cmd_mcp(args) -> int:
    if args.pack:
        os.environ["DESIGN_CANVAS_PACK"] = str(Path(args.pack).expanduser())
    from .server.stdio_mcp import run

    run(editor_port=args.editor_port)
    return 0


def cmd_export_sft(args) -> int:
    from .core.pack import load_pack
    from .train_data import export

    pack = load_pack(args.pack)
    n = export(pack, args.split, Path(args.out), surfaces=args.surfaces.split(","),
               modes=args.modes.split(","), limit=args.limit)
    print(f"wrote {n} trajectories to {args.out}")
    return 0


def cmd_export_cua(args) -> int:
    from .core.pack import load_pack
    from .cua import export

    summary = export(load_pack(args.pack), args.split, Path(args.out), limit=args.limit, start=args.start,
                     modes=tuple(args.modes.split(",")), seed=args.seed, workers=args.workers,
                     server=args.server, max_elements=args.max_elements, image_format=args.image_format)
    print(json.dumps(summary, indent=2))
    return 0


def cmd_export_hf(args) -> int:
    from .cua_format import Window, export_hf

    summary = export_hf(Path(args.cua), Path(args.out), image_format=args.image_format, min_reward=args.min_reward,
                        window=Window(args.max_images, args.keep_images), shard_mb=args.shard_mb,
                        with_windows=args.windows)
    print(json.dumps(summary, indent=2))
    return 0


def cmd_smoke(args) -> int:
    from .smoke import run

    return run(args.pack)


def _corpus(args):
    from .core.corpus import CorpusPack

    manifest = args.corpus or os.environ.get("DESIGN_CANVAS_CORPUS")
    if not manifest:
        raise SystemExit("set --corpus or DESIGN_CANVAS_CORPUS to a corpus manifest")
    return CorpusPack(manifest)


def cmd_corpus(args) -> int:
    from .core import corpus

    if args.action == "build":
        source = corpus.BucketSource(args.bucket, args.source_root)
        m = corpus.build_index(args.out, source, max_elements=args.max_elements, workers=args.workers)
        print(json.dumps({k: m[k] for k in ("snapshot_id", "splits", "rejected", "difficulty")}, indent=2))
    elif args.action == "publish":
        print(corpus.publish_index(args.out, args.bucket))
    elif args.action == "warm":
        pack = _corpus(args)
        out = pack.warm(args.splits.split(","), workers=min(args.workers, 8),
                        progress=lambda g, n, d: print(f"  {g}/{n} row groups, {d} designs", flush=True))
        out["failed"] = out["failed"][:20]
        print(json.dumps(out, indent=2))
        return 1 if out["failed"] else 0
    else:
        pack = _corpus(args)
        print(json.dumps({**pack.stats(), "splits": pack.manifest["splits"],
                          "rejected": pack.manifest["rejected"]}, indent=2))
    return 0


def cmd_evalset(args) -> int:
    from .core import evalset as ev

    pack = _corpus(args)
    if args.action == "build":
        record = ev.build(pack, split=args.split, size=args.size, lite=args.lite, seed=args.seed,
                          validate=not args.no_validate)
        ev.save(record, args.out)
        print(json.dumps({"evalset_id": record["evalset_id"], **record["summary"],
                          "lite": ev.summarize(record, "lite"), "failures": record["failures"]}, indent=2))
        return 1 if record["failures"] else 0
    record = ev.load(args.file, pack.snapshot_id)
    ev.check_against(record, pack)
    if args.action == "calibrate":
        from .core.difficulty import calibrate

        rewards: dict[str, list[float]] = {}
        for trace in args.trace:
            for line in Path(trace).read_text().splitlines():
                r = json.loads(line)
                if r.get("evalset_task") and not r.get("error"):
                    rewards.setdefault(r["evalset_task"], []).append(r["reward"] or 0.0)
        ids = {e["task_id"] for e in record["tasks"]}
        rows = [r for r in pack.rows(record["split"]) if r["task_id"] in ids]
        print(json.dumps(calibrate(rows, rewards, pack.manifest["difficulty"]["weights"]), indent=2))
        return 0
    if args.action == "export":
        print(json.dumps(ev.export_pack(pack, record, args.out), indent=2))
    elif args.action == "subset":
        ev.add_subset(record, args.name, args.size, args.within)
        ev.save(record, args.file)
        print(json.dumps({"evalset_id": record["evalset_id"], "subset": args.name,
                          **ev.summarize(record, args.name), "task_ids": record["subsets"][args.name]}, indent=2))
    else:
        print(json.dumps({"evalset_id": record["evalset_id"], "ok": True, **record["summary"]}, indent=2))
    return 0


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(prog="design-canvas", description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = parser.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("prepare", help="build the task pack from Crello")
    _pack_arg(p)
    p.add_argument("--train", type=int, default=200)
    p.add_argument("--validation", type=int, default=60)
    p.add_argument("--test", type=int, default=100)
    p.add_argument("--max-elements", type=int, default=20)
    p.add_argument("--max-shards", type=int, default=4)
    p.set_defaults(fn=cmd_prepare)

    p = sub.add_parser("serve", help="run the OpenEnv server with the editor")
    _pack_arg(p)
    p.add_argument("--host", default="127.0.0.1")
    p.add_argument("--port", type=int, default=int(os.environ.get("PORT", "8007")))
    p.add_argument("--workers", type=int, default=int(os.environ.get("DESIGN_CANVAS_WORKERS", "1")),
                   help="server processes behind one URL (a session-aware router); 1 = a single process")
    p.set_defaults(fn=cmd_serve)

    p = sub.add_parser("mcp", help="stdio MCP server (e.g. claude mcp add design-canvas -- ...)")
    _pack_arg(p)
    p.add_argument("--editor-port", type=int, default=8007,
                   help="also serve the live editor on this port (0 to disable)")
    p.set_defaults(fn=cmd_mcp)

    p = sub.add_parser("export-sft", help="write expert trajectories as chat JSONL")
    _pack_arg(p)
    p.add_argument("--split", default="train")
    p.add_argument("--out", default="sft.jsonl")
    p.add_argument("--surfaces", default="tools,html")
    p.add_argument("--modes", default="reference,description")
    p.add_argument("--limit", type=int, default=None)
    p.set_defaults(fn=cmd_export_sft)

    p = sub.add_parser("export-cua", help="record verified computer-use trajectories (screenshots + actions)")
    _pack_arg(p)
    p.add_argument("--split", default="train")
    p.add_argument("--out", default="artifacts/cua")
    p.add_argument("--limit", type=int, default=20)
    p.add_argument("--start", type=int, default=0)
    p.add_argument("--modes", default="reference,description")
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--workers", type=int, default=2)
    p.add_argument("--max-elements", type=int, default=12)
    p.add_argument("--image-format", default="webp", choices=["webp", "png", "jpeg"])
    p.add_argument("--server", default=None, help="a running server with DESIGN_CANVAS_DEBUG=1; default in-process")
    p.set_defaults(fn=cmd_export_cua)

    p = sub.add_parser("export-hf", help="write recorded trajectories as a cua-v1 Hugging Face dataset folder")
    p.add_argument("--cua", required=True, help="the recorder's cua.jsonl (images next to it)")
    p.add_argument("--out", required=True)
    p.add_argument("--image-format", default="keep", choices=["keep", "png", "jpeg"],
                   help="keep = the recorder's format (WebP), smallest")
    p.add_argument("--min-reward", type=float, default=0.98)
    p.add_argument("--max-images", type=int, default=8)
    p.add_argument("--keep-images", type=int, default=3)
    p.add_argument("--shard-mb", type=int, default=500)
    p.add_argument("--windows", action="store_true", help="also store the SFT windows view (duplicates screenshots)")
    p.set_defaults(fn=cmd_export_hf)

    p = sub.add_parser("smoke", help="offline end-to-end check")
    _pack_arg(p)
    p.set_defaults(fn=cmd_smoke)

    p = sub.add_parser("corpus", help="index, publish or inspect the bucket-backed corpus")
    p.add_argument("action", choices=["build", "publish", "stats", "warm"])
    p.add_argument("--out", default="artifacts/corpus-index", help="index directory (build, publish)")
    p.add_argument("--bucket", default="FineEnvs/crello-bucket")
    p.add_argument("--source-root", default=None, help="local copy or mount of the bucket")
    p.add_argument("--max-elements", type=int, default=30)
    p.add_argument("--workers", type=int, default=8)
    p.add_argument("--corpus", default=None, help="corpus manifest (stats, warm)")
    p.add_argument("--splits", default="train,validation,test", help="splits to download ahead of time (warm)")
    p.set_defaults(fn=cmd_corpus)

    p = sub.add_parser("evalset", help="build, verify or export the frozen evaluation set")
    p.add_argument("action", choices=["build", "verify", "export", "calibrate", "subset"])
    p.add_argument("file", nargs="?", help="frozen set (verify, export, subset)")
    p.add_argument("--name", help="subset name (subset)")
    p.add_argument("--within", default="lite", help="subset to nest inside (subset)")
    p.add_argument("--out", help="output file (build) or pack directory (export)")
    p.add_argument("--corpus", default=None, help="corpus manifest (default DESIGN_CANVAS_CORPUS)")
    p.add_argument("--split", default="test")
    p.add_argument("--size", type=int, default=300, help="set size (build) or subset size (subset)")
    p.add_argument("--lite", type=int, default=60)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--no-validate", action="store_true")
    p.add_argument("--trace", action="append", default=[], help="rollout trace.jsonl (calibrate)")
    p.set_defaults(fn=cmd_evalset)

    args = parser.parse_args(argv)
    return args.fn(args)


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