"""Re-answer on Opus with the same context minus parent tables whose row is present as a line item."""
import asyncio, json, os, re, sys, time
import psycopg
from quber.playground.answers import expectation
from quber.playground.retrieval import RetrievedChunk
from quber.playground.tracing import document_key
IN=sys.argv[1]; OUT=sys.argv[2]
d=json.load(open(IN)); doc_id=d["doc_id"]; doc_key=d["doc_key"]
conn=psycopg.connect(host=os.environ["POSTGRES_HOST"],user="quber",password=os.environ["POSTGRES_PASSWORD"],dbname="quber_rag")
cache={}
def chunk(cid):
    if cid not in cache:
        r=conn.execute("select chunk_id, chunk_type, page, content, parent_chunk_id from ade_playground.chunks where document_id=%s and chunk_id=%s",(doc_id,cid)).fetchone()
        cache[cid]=RetrievedChunk(chunk_id=r[0],chunk_type=r[1],page=r[2],content=r[3],score=0.0,parent_chunk_id=r[4])
    return cache[cid]
def dedupe(ids):
    chunks=[chunk(c) for c in ids]
    parents_with_rows={c.parent_chunk_id for c in chunks if c.chunk_type=="line_item" and c.parent_chunk_id}
    return [c for c in chunks if c.chunk_id not in parents_with_rows]
sem=asyncio.Semaphore(5)
async def one(r):
    async with sem:
        document_key.set(doc_key); t0=time.time()
        ctx=dedupe(r["context_ids"])
        try:
            ans=await expectation.answer(r["question"], ctx, "value")
            out={"shape":ans.payload.__class__.__name__.lower(),"payload":ans.payload.model_dump(),"cited_ids":ans.cited_ids,"error":None}
        except Exception as exc:
            out={"shape":"error","payload":{},"cited_ids":[],"error":repr(exc)}
        out.update({"index":r["index"],"question":r["question"],"context_ids":[c.chunk_id for c in ctx],"removed":len(r["context_ids"])-len(ctx),
                    "context_chars":[len(c.content) for c in ctx],"orig_chars":r["context_chars"],"picks":r["picks"],
                    "opus":{"shape":r["shape"],"payload":r["payload"],"cited_ids":r["cited_ids"]},"old":r["old"],"seconds":round(time.time()-t0,1)})
        print(f"[{r['index']:2d}] removed={out['removed']} {out['shape']:12s} {str(out['payload'].get('value') or out['error'] or '')[:30]!r} opus={r['shape']}:{str(r['payload'].get('value'))[:20]!r}", flush=True)
        return out
async def main():
    recs=await asyncio.gather(*(one(r) for r in d["rows"]))
    json.dump({"model":"claude-opus-5-dedupe","doc_key":doc_key,"doc_id":doc_id,"rows":recs}, open(OUT,"w"), indent=1, default=str)
    print("wrote", OUT)
asyncio.run(main())