"""Load RJ REIT filings into the playground RAG through the upload API.

Input: a manifest CSV with the columns s3_uri, folder, filing_type, year,
period, title — one row per document. Every row is submitted as an s3-sourced
upload job with its metadata; the server hashes the bytes, rejects exact
duplicates, and runs the fusion pipeline (two at a time). The script then
watches every job to a terminal state and prints a summary. Duplicate
rejections are reported as SKIPPED rather than treated as failures, because
re-running a manifest over already-loaded documents is the expected way to
top a quarter up.

The playground server must be running (default http://127.0.0.1:8101).
"""

from __future__ import annotations

import argparse
import csv
import json
import sys
import time

import requests

REQUIRED = ("s3_uri", "folder", "filing_type", "year", "period", "title")


def submit(base: str, row: dict[str, str]) -> tuple[str, dict[str, str]] | None:
    """Submit one document; returns (job_id, row) or None on refusal."""
    params = {
        "filename": row["s3_uri"].rsplit("/", 1)[-1],
        "pipeline": "fusion",
        "s3_uri": row["s3_uri"],
        "title": row["title"],
        "folder": row["folder"],
        "filing_type": row["filing_type"],
        "year": row["year"],
        "period": row["period"],
    }
    r = requests.post(f"{base}/api/upload", params=params, timeout=60)
    if r.status_code == 409:
        detail = r.json().get("detail", {})
        kind = detail.get("kind") if isinstance(detail, dict) else None
        message = detail.get("message") if isinstance(detail, dict) else str(detail)
        print(f"SKIPPED {row['s3_uri']}: {kind or '409'} — {message}")
        return None
    r.raise_for_status()
    return r.json()["job_id"], row


def main() -> int:
    p = argparse.ArgumentParser(description="Load filings from a manifest CSV into the playground RAG.")
    p.add_argument("manifest", help=f"CSV with columns {', '.join(REQUIRED)}")
    p.add_argument("--base", default="http://127.0.0.1:8101", help="playground server URL")
    p.add_argument("--poll", type=int, default=90, help="seconds between job polls")
    p.add_argument("--dry-run", action="store_true", help="print what would be submitted and exit")
    args = p.parse_args()

    with open(args.manifest, newline="") as f:
        rows = list(csv.DictReader(f))
    missing = [c for c in REQUIRED if rows and c not in rows[0]]
    if missing or not rows:
        print(f"manifest must have columns {REQUIRED} and at least one row; missing: {missing}")
        return 1

    if args.dry_run:
        for row in rows:
            print(
                f"would submit {row['folder']:5s} {row['filing_type']:5s} "
                f"{row['period']} {row['year']}  {row['s3_uri']}  '{row['title']}'"
            )
        print(f"{len(rows)} documents in manifest (dry run, nothing submitted)")
        return 0

    jobs: dict[str, dict] = {}
    for row in rows:
        submitted = submit(args.base, row)
        if submitted:
            job_id, r = submitted
            jobs[job_id] = r
            print(f"submitted {job_id}  {r['folder']:5s} {r['filing_type']:5s} {r['s3_uri']}")
    if not jobs:
        print("nothing to load")
        return 0

    last: dict[str, str] = {}
    while True:
        states: dict[str, dict] = {}
        for job_id in jobs:
            try:
                states[job_id] = requests.get(f"{args.base}/api/jobs/{job_id}", timeout=30).json()
            except Exception as exc:
                states[job_id] = {"stage": f"unreachable: {exc}"}
            stage = states[job_id].get("stage", "?")
            if last.get(job_id) != stage:
                print(f"{time.strftime('%H:%M:%S')}  {jobs[job_id]['folder']:5s} -> {stage}", flush=True)
                last[job_id] = stage
        if all(s.get("stage") in ("done", "failed") for s in states.values()):
            break
        time.sleep(args.poll)

    failed = 0
    print("\n==== LOAD SUMMARY ====")
    for job_id, row in jobs.items():
        s = states[job_id]
        if s.get("stage") == "done":
            print(f"DONE   {row['folder']:5s} {row['s3_uri']}  doc_id={s.get('doc_id')}")
        else:
            failed += 1
            print(f"FAILED {row['folder']:5s} {row['s3_uri']}")
            print(f"       error: {s.get('error')}")
            if s.get("error_info"):
                print(f"       info: {json.dumps(s['error_info'])}")
    print(f"\n{len(jobs) - failed} done, {failed} failed of {len(jobs)} submitted")
    return 1 if failed else 0


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