"""Proposed pipeline end to end. MODE=B: folded window, picks-only, whole tables. MODE=A: + row resolver."""
import asyncio, json, os, re, sys, time
import psycopg
from pydantic import BaseModel, Field
from pydantic_ai import Agent
from pydantic_ai.settings import ModelSettings
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rowctx import row_chunk
from quber.playground import retrieval
from quber.playground.retrieval import WINDOW, RetrievedChunk
from quber.playground.selection import select
from quber.playground.agent import anthropic_model
from quber.playground.answers import expectation
from quber.agents.langsmith_tracer import usage_metadata_from
from quber.playground.tracing import run_metadata, tracer, document_key
MODE=sys.argv[1]; IN=sys.argv[2]; OUT=sys.argv[3]; FLOOR=3
RESOLVE = MODE in ("A","D")
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+)$")
RESOLVER_SYSTEM=open(os.path.join(os.path.dirname(os.path.abspath(__file__)),"row_resolver_v2.py")).read().split('SYSTEM="""\\\n')[1].split('"""')[0]
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
resolver=Agent(anthropic_model("claude-haiku-4-5"), output_type=RowPick, system_prompt=RESOLVER_SYSTEM, model_settings=ModelSettings(temperature=0.0))
d=json.load(open(IN)); doc_id=d["doc_id"]; doc_key=d["doc_key"]
from quber.playground.answers import declared
_title,_ftype,_year,_period=conn.execute("select title, filing_type, year, period from ade_playground.documents where id=%s",(doc_id,)).fetchone()
DOC_LINE=f"Document: {_title} (form {_ftype}; reporting period {_period} {_year}). Every figure in the context below is from this document unless a chunk states another period."
_orig_prompt_for=declared.prompt_for
def prompt_with_document(question, chunks):
    return DOC_LINE + "\n\n" + _orig_prompt_for(question, chunks)
if MODE in ("C","D"): declared.prompt_for=prompt_with_document
tcache={}
def table_chunk(tid):
    if tid not in tcache:
        r=conn.execute("select chunk_id, chunk_type, page, content from ade_playground.chunks where document_id=%s and chunk_id=%s",(doc_id,tid)).fetchone()
        tcache[tid]=RetrievedChunk(chunk_id=r[0],chunk_type=r[1],page=r[2],content=r[3],score=0.0)
    return tcache[tid]
hcache={}
def page_heading(page):
    if page not in hcache:
        r=conn.execute("select chunk_id, chunk_type, page, content from ade_playground.chunks where document_id=%s and page=%s and chunk_type='text' order by length(content) asc limit 1",(doc_id,page)).fetchone()
        hcache[page]=RetrievedChunk(chunk_id=r[0],chunk_type=r[1],page=r[2],content=r[3],score=0.0) if r and len(r[3])<=300 else None
    return hcache[page]
def fold(window):
    seen=set(); out=[]
    for x in window:
        if x.chunk_type=="line_item" and x.parent_chunk_id:
            if x.parent_chunk_id in seen: continue
            seen.add(x.parent_chunk_id); out.append(table_chunk(x.parent_chunk_id))
        else:
            if x.chunk_id in seen: continue
            seen.add(x.chunk_id); out.append(x)
    return out
async def resolve(q, table):
    prompt=f"TABLE:\n{table.content}\n\nQUESTION: {q}\n\nWhich row holds the asked-for figure?"
    inputs={"messages":[{"role":"system","content":RESOLVER_SYSTEM},{"role":"user","content":prompt}]}
    async with tracer().llm_run("resolve_row_v2", inputs, model="claude-haiku-4-5", extra_metadata=run_metadata()) as run:
        res=await resolver.run(prompt)
        run.outputs={"messages":[{"role":"assistant","content":res.output.model_dump_json()}],"usage_metadata":usage_metadata_from(res.usage)}
    m=CELL.match(res.output.row or ""); t=table.chunk_id.split("/")[-1]
    return (int(m.group(2)) if m and m.group(1)==t else None), res.output
sem=asyncio.Semaphore(5)
async def one(r):
    async with sem:
        document_key.set(doc_key); t0=time.time(); q=r["question"]
        window=retrieval._fused_window(doc_key, q, WINDOW, None)
        cands=retrieval._cap_line_records(window, retrieval.PARENT_CAP) if MODE in ('C','D') else fold(window)
        try: picked=await select(q, cands)
        except Exception as exc: picked=[]
        chosen=[cands[j] for j in picked]
        if MODE in ('C','D'):
            seen={c.chunk_id for c in chosen}
            for c in cands:
                if len(chosen) >= len(picked)+FLOOR: break
                if c.chunk_id not in seen: seen.add(c.chunk_id); chosen.append(c)
        elif not chosen: chosen=cands[:FLOOR]
        ctx=[]; resolved=[]
        for c in chosen:
            if RESOLVE and c.chunk_type=="table":
                row,out=await resolve(q, c)
                rc=row_chunk(conn, doc_id, c.chunk_id, row, with_labels=True) if row is not None else None
                if rc is not None: ctx.append(rc); resolved.append({"table":c.chunk_id,"row":row,"cell":out.row})
                else: ctx.append(c); resolved.append({"table":c.chunk_id,"row":None,"whole":out.whole_table,"not_found":out.not_found})
            else: ctx.append(c)
        try:
            ans=await expectation.answer(q, ctx, "value")
            res={"shape":ans.payload.__class__.__name__.lower(),"payload":ans.payload.model_dump(),"cited_ids":ans.cited_ids,"error":None}
        except Exception as exc:
            res={"shape":"error","payload":{},"cited_ids":[],"error":repr(exc)}
        res.update({"index":r["index"],"question":q,"n_cands":len(cands),"cand_chars":sum(len(c.content) for c in cands),"picks":len(picked),
                    "context_ids":[c.chunk_id for c in ctx],"context_chars":[len(c.content) for c in ctx],"resolved":resolved,
                    "baseline":{"shape":r["shape"],"payload":r["payload"],"cited_ids":r["cited_ids"],"context_chars":r["context_chars"]},"old":r["old"],"seconds":round(time.time()-t0,1)})
        print(f"[{r['index']:2d}] picks={len(picked)} ctx={sum(res['context_chars'])//4:5d}tok {res['shape']:12s} {str(res['payload'].get('value') or res['error'] or '')[:24]!r} base={r['shape']}:{str(r['payload'].get('value'))[:18]!r}", flush=True)
        return res
async def main():
    recs=await asyncio.gather(*(one(r) for r in d["rows"]))
    json.dump({"mode":MODE,"doc_key":doc_key,"doc_id":doc_id,"rows":recs}, open(OUT,"w"), indent=1, default=str)
    print("wrote", OUT)
asyncio.run(main())