#!/usr/bin/env python3
"""Submit / monitor / collect the extraction batches.

  .venv/bin/python scripts/04_run_batches.py submit [img_haiku img_sonnet img_opus txt_haiku txt_sonnet txt_opus]
  .venv/bin/python scripts/04_run_batches.py status
  .venv/bin/python scripts/04_run_batches.py collect

Design (configs/data.yaml, configs/models.yaml): Haiku 5.5 sees up to 12 images per listing (image-count sensitivity),
Sonnet 5.5 and Opus 5.5 the first 10 (cost control; comparisons between tiers use the common first 10).
custom_id = "<tier>-<uid with ':'→'__'>-<image_index or 'T'>" (letters, digits, '_', '-' only).
"""
import json
import sys
from pathlib import Path

import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from src.extraction import client as C  # noqa: E402
from src.extraction.textprep import build_text  # noqa: E402
from src.feature_schema.schema import IMAGE_SCHEMA, SCHEMA_VERSION, TEXT_SCHEMA  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]
RAW = ROOT / "results" / "raw" / "batches"
IMG_CAP = {"haiku": 12, "sonnet": 10, "opus": 10}


def cid(tier: str, uid: str, idx) -> str:
    return f"{tier}-{uid.replace(':', '__')}-{idx}"


def build(name: str, smp: pd.DataFrame) -> list[tuple[str, dict]]:
    kind, tier = name.split("_")
    rows = smp[smp[f"tier_{tier}"]]
    reqs = []
    if kind == "img":
        for r in rows.itertuples():
            for i, url in enumerate(list(r.image_urls_used)[:IMG_CAP[tier]]):
                reqs.append((cid(tier, r.uid, i), C.image_request(tier, url, IMAGE_SCHEMA)))
    else:
        stats = []
        for r in rows.itertuples():
            text, st = build_text(r._asdict())
            st["uid"] = r.uid
            stats.append(st)
            reqs.append((cid(tier, r.uid, "T"), C.text_request(tier, text, TEXT_SCHEMA)))
        pd.DataFrame(stats).to_csv(RAW / f"{name}{SUFFIX}.text_stats.csv", index=False)
    return reqs


import os
SAMPLE_FILE = os.environ.get("WP20_SAMPLE", "sample.parquet")
SUFFIX = os.environ.get("WP20_BATCH_SUFFIX", "")


def submit(names):
    smp = pd.read_parquet(ROOT / "data" / "processed" / SAMPLE_FILE)
    meta = {"schema_version": SCHEMA_VERSION, "prompt_version": C.PROMPTS["prompt_version"], "models": C.MODELS["tiers"],
            "image_caps": IMG_CAP, "n_listings": int(len(smp))}
    json.dump(meta, open(RAW / "run_meta.json", "w"), indent=2, default=str)
    for name in names:
        out_name = name + SUFFIX
        if (RAW / f"{out_name}.batch_ids.txt").exists():
            print(f"{out_name}: already submitted, skipping")
            continue
        reqs = build(name, smp)
        ids = C.batch_submit(out_name, reqs)
        print(f"{out_name}: {len(reqs)} requests → batches {ids}")


def all_batches():
    out = []
    for f in sorted(RAW.glob("*.batch_ids.txt")):
        name = f.name.replace(".batch_ids.txt", "")
        for line in open(f):
            bid = line.split("\t")[0].strip()
            if bid:
                out.append((name, bid))
    return out


def status():
    for name, bid in all_batches():
        b = C.batch_status(bid)
        rc = b.request_counts
        print(f"{name:12s} {bid} {b.processing_status:10s} ok={rc.succeeded} err={rc.errored} proc={rc.processing} exp={rc.expired} can={rc.canceled}")


def collect():
    for name, bid in all_batches():
        b = C.batch_status(bid)
        if b.processing_status == "ended":
            n = C.batch_collect(name, bid)
            print(f"{name}: collected {n} new results from {bid}")
        else:
            print(f"{name}: {bid} still {b.processing_status}")


if __name__ == "__main__":
    cmd = sys.argv[1]
    if cmd == "submit":
        submit(sys.argv[2:] or ["img_haiku", "img_sonnet", "img_opus", "txt_haiku", "txt_sonnet", "txt_opus"])
    elif cmd == "status":
        status()
    elif cmd == "collect":
        collect()