"""Row resolver test: one table + question -> which row. Compared to the row Opus cited."""
import asyncio, json, glob, os, re, time
import psycopg
from pydantic import BaseModel, Field
from pydantic_ai import Agent
from pydantic_ai.settings import ModelSettings
from quber.playground.agent import anthropic_model
from quber.agents.langsmith_tracer import usage_metadata_from
from quber.playground.tracing import run_metadata, tracer, document_key
S=os.path.dirname(os.path.abspath(__file__)); MODEL="claude-haiku-4-5"
conn=psycopg.connect(host=os.environ["POSTGRES_HOST"],user="quber",password=os.environ["POSTGRES_PASSWORD"],dbname="quber_rag")
CELL=re.compile(r"^t(\d+)-(\d+)-(\d+)$")
SYSTEM="""\
You are given ONE table from a financial document and a QUESTION that asks for a
single figure. Your only job is to say which ROW of the table holds that figure.

Each cell is tagged like <td id="tN-R-C">, where R is the row number. Answer with
the id of any cell in the row that holds the figure, for example "t32-4-3".

Rules:
- Do not state the figure. Do not compute anything. Name the row.
- If the question needs several rows (a whole table, a series across rows), set
  whole_table to true and leave row empty.
- If no row in this table holds the asked-for figure, set not_found to why and
  leave row empty. Do not pick a similar-looking row.
- Only return an id that literally appears in the table.
"""
class RowPick(BaseModel):
    row: str | None = Field(default=None, description="A cell id in the row that holds the figure, e.g. 't32-4-3'.")
    whole_table: bool = False
    not_found: str | None = None
agent=Agent(anthropic_model(MODEL), output_type=RowPick, system_prompt=SYSTEM, model_settings=ModelSettings(temperature=0.0))
cases=[]
for f in sorted(glob.glob(f"{S}/rerun_*.json")):
    if any(s in f for s in (".haiku",".sonnet",".dedupe")): continue
    d=json.load(open(f))
    for r in d["rows"]:
        if r["shape"]!="scalar": continue
        oc=r["cited_ids"][0] if r["cited_ids"] else ""; m=CELL.match(oc)
        if not m: continue
        cases.append((d["doc_key"], d["doc_id"], r["index"], r["question"], f"#/tables/{m.group(1)}", int(m.group(2)), oc, r["payload"].get("value")))
print("cases", len(cases), flush=True)
sem=asyncio.Semaphore(5)
async def one(c):
    dk, doc_id, idx, q, table, exp_row, oc, val = c
    async with sem:
        document_key.set(dk)
        content=conn.execute("select content from ade_playground.chunks where document_id=%s and chunk_id=%s",(doc_id,table)).fetchone()[0]
        prompt=f"TABLE:\n{content}\n\nQUESTION: {q}\n\nWhich row holds the asked-for figure?"
        inputs={"messages":[{"role":"system","content":SYSTEM},{"role":"user","content":prompt}]}
        async with tracer().llm_run("resolve_row", inputs, model=MODEL, extra_metadata=run_metadata()) as run:
            res=await agent.run(prompt)
            run.outputs={"messages":[{"role":"assistant","content":res.output.model_dump_json()}],"usage_metadata":usage_metadata_from(res.usage)}
        out=res.output; m=CELL.match(out.row or "")
        got=int(m.group(2)) if m and m.group(1)==table.split("/")[-1] else None
        verdict="MATCH" if got==exp_row else ("WHOLE" if out.whole_table else ("NONE" if out.row is None else "WRONG_ROW"))
        print(f"[{dk[:6]} {idx:2d}] expected row {exp_row:3d} got {got!s:>4} {verdict:9s} {q[:60]!r}", flush=True)
        return {"doc_key":dk,"index":idx,"question":q,"table":table,"expected_row":exp_row,"expected_cell":oc,"opus_value":val,"got_row":got,"got_cell":out.row,"whole_table":out.whole_table,"not_found":out.not_found,"verdict":verdict,"table_chars":len(content)}
async def main():
    recs=await asyncio.gather(*(one(c) for c in cases))
    json.dump(recs, open(f"{S}/row_resolver_haiku.json","w"), indent=1)
    from collections import Counter; print(Counter(r["verdict"] for r in recs))
asyncio.run(main())