"""Seed the sweep answer cache from a previously filled workbook.

A filled workbook already holds everything a cache row needs: the prompt is
rebuilt from Metric and Quarter exactly as the sweep builds it, the answer
sits in the 10Q and Deck columns, and provenance in their appended
provenance columns. This script reads one filled copy and inserts one open
row per answered cell into ade_playground.sweep_answers, so a sweep run
after seeding reuses the prior run's answers instead of re-asking.

Rules:
  - A tab seeds only when its ticker has the target quarter's 10-Q/99-2
    pair in the library; the pair is what resolves each answer to the
    doc_key the cache is keyed by. Tabs without a pair are reported.
  - Empty answer cells are skipped: the cell was never written, so there
    is no measurement to record.
  - ERROR cells are skipped: an error row records a failed measurement,
    not an answer, and the sweep never reuses errors anyway.
  - Existing open rows are never touched. Seeding inserts only where no
    open row exists, so re-running the seed cannot clobber newer answers.
"""

from __future__ import annotations

import argparse
import sys
import tempfile
from pathlib import Path

import openpyxl
import requests
from psycopg.types.json import Jsonb

from quber.playground import db as playground_db
from quber.playground.agent import DEFAULT_MODEL as ANSWER_MODEL

sys.path.insert(0, str(Path(__file__).parent))
from run_sweep import IN_SCOPE, answer_status, doc_key_for, fetch  # noqa: E402


def main() -> int:
    p = argparse.ArgumentParser(description="Seed the sweep answer cache from a filled workbook.")
    p.add_argument("workbook", help="filled workbook: s3:// URI or local path")
    p.add_argument("--year", type=int, default=2026, help="target filing year")
    p.add_argument("--period", default="Q2", help="target filing period, e.g. Q2")
    p.add_argument("--base", default="http://127.0.0.1:8101", help="playground server URL")
    args = p.parse_args()

    src = fetch(args.workbook, Path(tempfile.mkdtemp(prefix="rj-reit-seed-")))
    docs = requests.get(f"{args.base}/api/documents", timeout=30).json()

    def target_doc(ticker: str, filing_type: str) -> dict | None:
        hits = [
            d
            for d in docs
            if d["folder"] == ticker
            and d["filing_type"] == filing_type
            and d["year"] == args.year
            and d["period"] == args.period
        ]
        return hits[0] if len(hits) == 1 else None

    wb = openpyxl.load_workbook(src, data_only=False)
    inserted = existing = skipped_empty = skipped_error = 0
    with playground_db.connect() as conn:
        for name in wb.sheetnames:
            if name == "Summary":
                continue
            ten_q, deck = target_doc(name, "10-Q"), target_doc(name, "99-2")
            if ten_q is None or deck is None:
                print(f"SKIP {name}: no {args.period} {args.year} 10-Q/99-2 pair in the library", flush=True)
                continue
            ws = wb[name]
            col = {h: i + 1 for i, h in enumerate(c.value for c in ws[1])}
            keys = {"10Q": doc_key_for(conn, ten_q["id"]), "Deck": doc_key_for(conn, deck["id"])}
            tab_inserted = 0
            for r in range(2, ws.max_row + 1):
                found = ws.cell(row=r, column=col["Found"]).value
                metric = ws.cell(row=r, column=col["Metric"]).value
                quarter = ws.cell(row=r, column=col["Quarter"]).value
                if found not in IN_SCOPE or not metric:
                    continue
                question = f"Find {metric}. Return {quarter} value"
                for label in ("10Q", "Deck"):
                    cell = ws.cell(row=r, column=col[label]).value
                    if cell is None or str(cell).strip() == "":
                        skipped_empty += 1
                        continue
                    value = str(cell)
                    status = answer_status(value)
                    if status == "error":
                        skipped_error += 1
                        continue
                    response = {label: value}
                    prov_col = col.get(f"{label} Provenance")
                    prov = ws.cell(row=r, column=prov_col).value if prov_col else None
                    if prov:
                        response[f"{label} Provenance"] = str(prov)
                    cur = conn.execute(
                        "INSERT INTO ade_playground.sweep_answers"
                        " (doc_key, question, model, status, response)"
                        " VALUES (%s, %s, %s, %s, %s)"
                        " ON CONFLICT (doc_key, question, model) WHERE valid_to IS NULL DO NOTHING",
                        (keys[label], question, ANSWER_MODEL, status, Jsonb(response)),
                    )
                    if cur.rowcount:
                        inserted += 1
                        tab_inserted += 1
                    else:
                        existing += 1
            print(f"{name}: {tab_inserted} rows seeded", flush=True)

    print(
        f"\nSEED DONE: {inserted} inserted, {existing} already open,"
        f" {skipped_empty} empty cells skipped, {skipped_error} error cells skipped",
        flush=True,
    )
    return 0


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