"""P1/P4: download 1-minute bars for a symbol set through /exports (resumable, quota-checked) and build dense grids.
Usage: download_pilot.py --config configs/pilot.yaml [--symbols-from pilot|feasibility|all] [--workers 3]"""
from __future__ import annotations

import argparse
import json
import sys
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

import polars as pl

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from chronos_x.config import load_config  # noqa: E402
from chronos_x.data.client import FinDataClient  # noqa: E402
from chronos_x.data.exports import download_universe, fetch_instrument  # noqa: E402
from chronos_x.data.features import build_grid, save_grid  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]


def symbols_for(man: dict, which: str) -> list[dict]:
    inst = {r["symbol"]: r for r in man["instruments"]}
    if which == "pilot":
        return [inst[s] for s in man["pilot_symbols"]]
    if which == "feasibility":
        return [inst[s] for s in man["feasibility_symbols"]]
    return list(inst.values())


def build_one(args) -> dict:
    sym, ac, cache, start, end, version = args
    t0 = time.time()
    raw = pl.read_parquet(Path(cache) / "bars" / f"{sym}.parquet")
    inst = json.loads((Path(cache) / "instruments" / f"{sym}.json").read_text())
    arrays, meta = build_grid(raw, inst.get("corporate_actions", []), start, end, sym, ac, version)
    save_grid(arrays, meta, Path(cache) / "grids" / f"{start}_{end}" / sym)
    return {"symbol": sym, "seconds": round(time.time() - t0, 1), **meta.counts, "qa": meta.qa}


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="configs/pilot.yaml")
    ap.add_argument("--symbols-from", default="pilot", choices=["pilot", "feasibility", "all"])
    ap.add_argument("--workers", type=int, default=3)
    ap.add_argument("--skip-download", action="store_true")
    a = ap.parse_args()
    cfg = load_config(a.config)
    man = json.loads((ROOT / cfg["universe"]["manifest"]).read_text())
    rows = symbols_for(man, a.symbols_from)
    cache = Path(cfg["paths"]["cache"])
    start, end = cfg["data"]["start"], cfg["data"]["end"]
    c = FinDataClient(log_path=ROOT / "results" / "logs" / f"provenance_download_{cfg['name']}.jsonl")
    out_dir = ROOT / "results" / "metrics"; out_dir.mkdir(parents=True, exist_ok=True)
    if not a.skip_download:
        t0 = time.time()
        ds_for = {r["symbol"]: cfg["data"]["dataset_etf"] if r["asset_class"] == "etf" else cfg["data"]["dataset_stock"]
                  for r in rows}
        res = download_universe(c, [r["symbol"] for r in rows], start, end, cache, ds_for, workers=a.workers)
        summary = {"config": cfg["name"], "symbols_from": a.symbols_from, "n": len(res), "start": start, "end": end,
                   "downloaded": sum(r["status"] == "downloaded" for r in res), "cached": sum(r["status"] == "cached" for r in res),
                   "failed": [r for r in res if r["status"] == "failed"], "rows": sum(r.get("row_count", 0) for r in res),
                   "dataset_versions": sorted({r.get("dataset_version") for r in res if r.get("dataset_version")}),
                   "wall_s": round(time.time() - t0, 1)}
        (out_dir / f"download_{cfg['name']}_{a.symbols_from}.json").write_text(json.dumps(summary, indent=1))
        print(json.dumps({k: v for k, v in summary.items() if k != "failed"}), "failed:", len(summary["failed"]))
    # grids
    t0 = time.time()
    jobs = []
    for r in rows:
        p = cache / "bars" / f"{r['symbol']}.parquet"
        if p.exists():
            fetch_instrument(c, r["symbol"], cache)
            jobs.append((r["symbol"], r["asset_class"], str(cache), start, end, None))
    results = []
    with ProcessPoolExecutor(max_workers=max(1, a.workers * 2)) as ex:
        for f in as_completed([ex.submit(build_one, j) for j in jobs]):
            results.append(f.result())
    results.sort(key=lambda r: r["symbol"])
    cov = [r["coverage"] for r in results]
    rep = {"config": cfg["name"], "symbols_from": a.symbols_from, "n_grids": len(results), "wall_s": round(time.time() - t0, 1),
           "coverage_mean": sum(cov) / max(1, len(cov)), "coverage_min": min(cov) if cov else None,
           "rows_regular": sum(r["rows_regular"] for r in results), "rows_extended": sum(r["rows_extended_hours"] for r in results),
           "rows_dropped_integrity": sum(r["rows_dropped_integrity"] for r in results),
           "duplicates_in_raw": sum(r["qa"].get("duplicates", 0) for r in results), "per_symbol": results}
    (out_dir / f"grids_{cfg['name']}_{a.symbols_from}.json").write_text(json.dumps(rep, indent=1))
    print(json.dumps({k: v for k, v in rep.items() if k != "per_symbol"}))


if __name__ == "__main__":
    main()