#!/usr/bin/env python3
"""Step 05 — parse all batch results, build per-image and property-level tables for every tier and image budget K,
and the baseline property tables B0–B4 (per tier). Writes:
  data/processed/images_<tier>.parquet            per-image extraction rows (ok + errors)
  data/processed/text_<tier>.parquet              per-listing text extraction rows
  data/processed/props_<tier>_K<k>.parquet        aggregated visual attributes (K = 1,2,4,8,10,12 where available)
  data/processed/levels_<tier>.parquet            long table: level ∈ {B0..B4} × uid with cost-model inputs
  results/raw/aggregate_summary.json
Also handles the supplementary "sold" sample (suffix _sold) when present.
"""
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.aggregate import parse_results, aggregate_images, build_levels  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]
RAW = ROOT / "results" / "raw" / "batches"
OUT = ROOT / "data" / "processed"
KS = {"haiku": [1, 2, 4, 8, 10, 12], "sonnet": [1, 2, 4, 8, 10], "opus": [1, 2, 4, 8, 10]}


def main(suffix: str = "") -> None:
    frame_file = "sample_sold.parquet" if suffix else "sample.parquet"
    smp = pd.read_parquet(OUT / frame_file)
    summary = {}
    texts = {}
    for tier in ["haiku", "sonnet", "opus"]:
        f = RAW / f"txt_{tier}{suffix}.results.jsonl"
        if f.exists():
            t = parse_results(f"txt_{tier}{suffix}")
            t.to_parquet(OUT / f"text_{tier}{suffix}.parquet", index=False)
            texts[tier] = t
            summary[f"text_{tier}"] = {"n": int(len(t)), "ok": int(t.ok.sum())}
    for tier in ["haiku", "sonnet", "opus"]:
        f = RAW / f"img_{tier}{suffix}.results.jsonl"
        if not f.exists():
            continue
        img = parse_results(f"img_{tier}{suffix}")
        img.to_parquet(OUT / f"images_{tier}{suffix}.parquet", index=False)
        summary[f"img_{tier}"] = {"n": int(len(img)), "ok": int(img.ok.sum()), "listings": int(img.uid.nunique()),
                                   "mean_in_tokens": float(img.input_tokens.mean()), "mean_out_tokens": float(img.output_tokens.mean())}
        props_full = None
        for k in KS[tier]:
            p = aggregate_images(img, K=k)
            p.to_parquet(OUT / f"props_{tier}{suffix}_K{k}.parquet", index=False)
            props_full = p
        txt = texts.get(tier) if tier in texts else texts.get("sonnet")
        levels = build_levels(smp, props_full, txt)
        long = pd.concat([df.assign(level=lvl).reset_index() for lvl, df in levels.items()], ignore_index=True)
        long.to_parquet(OUT / f"levels_{tier}{suffix}.parquet", index=False)
        summary[f"levels_{tier}"] = {lvl: int(len(df)) for lvl, df in levels.items()}
    json.dump(summary, open(ROOT / "results" / "raw" / f"aggregate_summary{suffix}.json", "w"), indent=2)
    print(json.dumps(summary, indent=1))


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else "")