"""Row resolver, loosened label matching. Two sets: tables that hold the answer (expected row known),
and tables that do not (every row is wrong; the only right output is a decline)."""
import asyncio, json, glob, os, re, collections
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.
- The row may be labelled differently from the question when it is the same line
  under ordinary financial naming: "term loans" for "Secured term loan, net",
  "first mortgages" for "Senior loans", "corporate debt" for a corporate revolving
  credit facility. Pick that row.
- Decline when the figure would need arithmetic across rows, or when no row is
  that line. A row that merely mentions a related item, a component, or a
  different measure is not the line. Set not_found to why and leave row empty.
- If the question needs several rows (a whole table, a series across rows), set
  whole_table to true and leave row empty.
- 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))
# Set A: tables that hold the answer (from the Opus rerun, cited cell known)
A=[]; B=[]
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"]:
        picked_tables=[c for c in r["context_ids"][:r["picks"]] if c.startswith("#/tables/")]
        if r["shape"]=="scalar":
            oc=r["cited_ids"][0] if r["cited_ids"] else ""; m=CELL.match(oc)
            if m: A.append((d["doc_key"], d["doc_id"], r["index"], r["question"], f"#/tables/{m.group(1)}", int(m.group(2))))
        else:
            # Set B: the tables the agent picked for questions Opus found unanswerable
            for t in picked_tables[:1]:
                B.append((d["doc_key"], d["doc_id"], r["index"], r["question"], t, None))
print(f"set A (answer present): {len(A)}   set B (no answer in document per Opus): {len(B)}", flush=True)
sem=asyncio.Semaphore(5)
async def one(c, which):
    dk, doc_id, idx, q, table, exp_row = 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_v2", 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
        if which=="A": verdict="MATCH" if got==exp_row else ("WHOLE" if out.whole_table else ("DECLINE" if got is None else "WRONG_ROW"))
        else: verdict="DECLINE" if (got is None and not out.whole_table) else ("WHOLE" if out.whole_table else "NAMED_A_ROW")
        return {"set":which,"doc_key":dk,"doc_id":doc_id,"index":idx,"question":q,"table":table,"expected_row":exp_row,"got_row":got,"got_cell":out.row,"whole_table":out.whole_table,"not_found":out.not_found,"verdict":verdict}
async def main():
    recs=await asyncio.gather(*([one(c,"A") for c in A]+[one(c,"B") for c in B]))
    json.dump(recs, open(f"{S}/row_resolver_v2.json","w"), indent=1)
    for which in ("A","B"):
        print(which, dict(collections.Counter(r["verdict"] for r in recs if r["set"]==which)))
asyncio.run(main())