"""Fill the RJ REIT metric sweep workbook from the playground RAG.

For each ticker tab whose folder holds both a 10-Q and a 99-2 for the target
quarter, applies each in-scope row's prompt to both documents through the
playground batch API and writes the answers into the 10Q and Deck columns.

The rules, exactly as the workbook defines them:
  - A row is in scope when its Found column is "xbrl" or
    "docling (non-xbrl tool)"; not_reported rows are left untouched.
  - The prompt is the sheet's own formula computed here (openpyxl does not
    evaluate formulas): "Find {Metric}. Return {Quarter} value".
  - A scalar answer records its printed value; an unanswerable records the
    #not found# sentinel; a failed question records ERROR: <message>. An
    in-scope row with an empty Metric cell is reported and skipped.

Tabs without a loaded document pair are reported as skipped, never silently
dropped. Only the 10Q and Deck cells are written; the workbook is saved as a
new file after each ticker (checkpoint) and uploaded to S3 at the end when
--output is an s3:// URI. One batch runs at a time so total provider
concurrency stays at the batch runner's own bound.

Answers are cached in ade_playground.sweep_answers, keyed by the document's
content-derived doc_key, the exact prompt text, and the answering model. A
question with an open cache row is reused — value and provenance copied from
the stored {field header: value} response — and only misses reach the batch
API. Error rows are never reused. There is no bypass flag: to re-measure,
close the rows (SET valid_to = now()) and the next run asks fresh.
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
import tempfile
import time
from pathlib import Path

import openpyxl
import psycopg
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

IN_SCOPE = {"xbrl", "docling (non-xbrl tool)"}
NOT_IN_DOCUMENT = "#not found#"


def table_number(ref: dict) -> str | None:
    """The parse's table number a reference points into, from its ref id."""
    rid = ref.get("ref_id") or ""
    if rid.startswith("#/tables/"):
        return rid.rsplit("/", 1)[-1]
    if rid.startswith("t"):
        return rid.split("-", 1)[0].lstrip("t")
    return None


def reference_part(ref: dict) -> str:
    """One cited reference as a human-readable location on the page.

    A table cell reads page, table, row, column (row and column 1-based; the
    table keeps the parse's own number, taken from the cell's t<N>-<r>-<c>
    ref id). A figure value reads page, chart, segment. A whole table or a
    line item reads page and table number. Anything else reads page and kind.
    """
    page = ref.get("page")
    kind = ref.get("ref_type")
    table = table_number(ref)
    if kind == "tableCell" and ref.get("row") is not None and ref.get("col") is not None:
        return f"p{page} table {table} row {ref['row'] + 1} col {ref['col'] + 1}"
    if kind == "figureValue":
        chart = ref.get("chart") or "figure"
        segment = ref.get("segment") or ref.get("text") or ""
        return f"p{page} {chart}, {segment}".rstrip(", ")
    if table is not None:
        return f"p{page} table {table}"
    return f"p{page} {kind or 'reference'}"


def provenance(row: dict) -> str:
    """The cited references of one answered question, joined for the column.

    A citation of a whole table or of a line item says nothing a cited cell
    in the same table has not already said, so those are dropped when such a
    cell is present. Repeated renderings collapse to one, citation order kept.
    """
    refs = ((row or {}).get("response") or {}).get("references") or []
    covered = {
        (r.get("page"), table_number(r))
        for r in refs
        if r.get("ref_type") == "tableCell" and r.get("row") is not None
    }
    parts: list[str] = []
    for ref in refs:
        if ref.get("ref_type") in ("table", "chunkTable", "line_item"):
            if (ref.get("page"), table_number(ref)) in covered:
                continue
        part = reference_part(ref)
        if part not in parts:
            parts.append(part)
    return "; ".join(parts)


def cell_value(row: dict) -> str:
    if row.get("error"):
        return f"ERROR: {row['error']}"
    resp = row.get("response") or {}
    kind = resp.get("shape")
    payload = resp.get("value") or {}
    if kind == "scalar":
        return payload.get("value") or ""
    if kind == "unanswerable":
        return NOT_IN_DOCUMENT
    if kind == "prose":
        return payload.get("text") or resp.get("answer") or ""
    return json.dumps(payload, separators=(",", ":"))


def answer_status(value: str) -> str:
    """The cache status of one written cell value."""
    if value.startswith("ERROR:"):
        return "error"
    if value == NOT_IN_DOCUMENT:
        return "not_found"
    return "value"


def doc_key_for(conn: psycopg.Connection, doc_id: int) -> str:
    """The content-derived storage key the cache is keyed by.

    The documents API deliberately withholds doc_key (nothing in the UI
    needs it), so it is read from the table directly.
    """
    found = conn.execute("SELECT doc_key FROM ade_playground.documents WHERE id = %s", (doc_id,)).fetchone()
    if found is None:
        raise RuntimeError(f"document {doc_id} not in the library")
    return found[0]


def open_answers(conn: psycopg.Connection, doc_key: str) -> dict[str, dict]:
    """The document's current cached answers, question text -> row."""
    found = conn.execute(
        "SELECT question, status, response FROM ade_playground.sweep_answers"
        " WHERE doc_key = %s AND model = %s AND valid_to IS NULL",
        (doc_key, ANSWER_MODEL),
    ).fetchall()
    return {q: {"status": s, "response": resp} for q, s, resp in found}


def store_answer(conn: psycopg.Connection, doc_key: str, question: str, status: str, response: dict) -> None:
    """Record a fresh answer: close the key's open row, insert the new one."""
    conn.execute(
        "UPDATE ade_playground.sweep_answers SET valid_to = now()"
        " WHERE doc_key = %s AND question = %s AND model = %s AND valid_to IS NULL",
        (doc_key, question, ANSWER_MODEL),
    )
    conn.execute(
        "INSERT INTO ade_playground.sweep_answers (doc_key, question, model, status, response)"
        " VALUES (%s, %s, %s, %s, %s)",
        (doc_key, question, ANSWER_MODEL, status, Jsonb(response)),
    )


def run_batch(base: str, doc_id: int, questions: list[str]) -> dict[int, dict]:
    r = requests.post(
        f"{base}/api/batch",
        json={"doc_id": doc_id, "questions": "\n".join(questions), "want": "value"},
        timeout=60,
    )
    r.raise_for_status()
    job_id = r.json()["job_id"]
    while True:
        time.sleep(15)
        state = requests.get(f"{base}/api/batch/{job_id}", timeout=30).json()
        if "running" not in state:
            # The registry answers 404 detail-only when it no longer holds the
            # run - a server restart (e.g. a --reload triggered by files
            # changing under the repo) wipes the in-memory registry mid-run.
            raise RuntimeError(
                f"batch {job_id} lost by the server mid-run ({state}); "
                "the checkpoint holds every ticker finished before this one"
            )
        if not state["running"]:
            break
    return {row["index"]: row for row in state["rows"]}


def fetch(uri: str, dest_dir: Path) -> Path:
    if uri.startswith("s3://"):
        local = dest_dir / uri.rsplit("/", 1)[-1]
        subprocess.run(["aws", "s3", "cp", uri, str(local), "--quiet"], check=True)
        return local
    return Path(uri)


def main() -> int:
    p = argparse.ArgumentParser(description="Fill the metric sweep workbook from the playground RAG.")
    p.add_argument("workbook", help="source workbook: s3:// URI or local path")
    p.add_argument(
        "--output",
        default=None,
        help="destination for the filled copy (s3:// URI or path); default: source stem + '-filled.xlsx' beside the source",
    )
    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")
    p.add_argument("--dry-run", action="store_true", help="report scope per tab without answering anything")
    p.add_argument(
        "--tickers",
        default=None,
        help="comma-separated tab names to process (e.g. BRSP,KREF); default: every tab",
    )
    args = p.parse_args()

    workdir = Path(tempfile.mkdtemp(prefix="rj-reit-sweep-"))
    src = fetch(args.workbook, workdir)
    out_uri = args.output or args.workbook.rsplit(".", 1)[0] + "-filled.xlsx"
    local_out = workdir / "filled.xlsx" if out_uri.startswith("s3://") else Path(out_uri)

    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)
    # A misspelled ticker must abort here, before any answer is written or
    # cached, or the operator gets a partial run they did not notice asking for.
    tickers = [t.strip() for t in args.tickers.split(",") if t.strip()] if args.tickers else None
    if tickers is not None:
        unknown = [t for t in tickers if t not in wb.sheetnames]
        if unknown:
            print(f"ERROR: no such tab(s) in the workbook: {', '.join(unknown)}", flush=True)
            return 1
    total_written = total_cached = total_asked = errors = 0
    with playground_db.connect() as conn:
        for name in wb.sheetnames:
            if name == "Summary" or (tickers is not None and name not in tickers):
                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])}
            rows: list[tuple[int, str]] = []
            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 in IN_SCOPE:
                    if not metric:
                        print(f"SKIP {name} row {r}: in scope but Metric is empty", flush=True)
                        continue
                    rows.append((r, f"Find {metric}. Return {quarter} value"))
            if not rows:
                print(f"SKIP {name}: no in-scope rows", flush=True)
                continue
            questions = [q for _, q in rows]
            if args.dry_run:
                print(
                    f"{name}: {len(rows)} in-scope rows -> docs 10Q={ten_q['id']} Deck={deck['id']} (dry run)"
                )
                continue
            sources: dict[int, dict[str, str]] = {}
            for label, doc, column in (("10Q", ten_q, col["10Q"]), ("Deck", deck, col["Deck"])):
                t0 = time.time()
                doc_key = doc_key_for(conn, doc["id"])
                cached = open_answers(conn, doc_key)
                # A question is asked when it has no open row, its open row is
                # an error (errors are never reused), or the stored response
                # predates this column and cannot fill it.
                asked = [
                    i
                    for i, q in enumerate(questions)
                    if (c := cached.get(q)) is None or c["status"] == "error" or label not in c["response"]
                ]
                results = run_batch(args.base, doc["id"], [questions[i] for i in asked]) if asked else {}
                fresh = {asked[j]: row for j, row in results.items()}
                asked_set = set(asked)
                for i, (r, q) in enumerate(rows):
                    if i in asked_set:
                        row = fresh.get(i)
                        value = str(cell_value(row)) if row else "ERROR: row missing from run"
                        prov = provenance(row) if row else ""
                        response = {label: value}
                        if prov:
                            response[f"{label} Provenance"] = prov
                        store_answer(conn, doc_key, q, answer_status(value), response)
                    else:
                        stored = cached[q]["response"]
                        value = str(stored.get(label, ""))
                        prov = str(stored.get(f"{label} Provenance", ""))
                    if value.startswith("ERROR:"):
                        errors += 1
                    ws.cell(row=r, column=column, value=value)
                    total_written += 1
                    sources.setdefault(r, {})[label] = prov
                n_cached = len(rows) - len(asked)
                total_cached += n_cached
                total_asked += len(asked)
                print(
                    f"{name} {label}: {len(rows)} answers, {n_cached} cached, {len(asked)} asked"
                    f" in {time.time() - t0:.0f}s (doc {doc['id']})",
                    flush=True,
                )
            # Provenance lands in its own column per source, appended after the
            # sheet's existing headers so no formula column shifts. The Comment
            # column stays the reader's own.
            for label in ("10Q", "Deck"):
                header = f"{label} Provenance"
                if header not in col:
                    col[header] = ws.max_column + 1
                    ws.cell(row=1, column=col[header], value=header)
                for r, by_label in sources.items():
                    if by_label.get(label):
                        ws.cell(row=r, column=col[header], value=by_label[label])
            wb.save(local_out)  # checkpoint after each ticker so a crash loses little

    if args.dry_run:
        return 0
    wb.save(local_out)
    if out_uri.startswith("s3://"):
        subprocess.run(["aws", "s3", "cp", str(local_out), out_uri, "--quiet"], check=True)
    print(
        f"\nSWEEP DONE: {total_written} cells written"
        f" ({total_cached} cached, {total_asked} asked), {errors} errors -> {out_uri}",
        flush=True,
    )
    return 0


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