"""Baseline contexts (10 chunks, from the Opus rerun) re-answered on Opus with the document line in the prompt."""
import asyncio, json, os, sys, time
import psycopg
from quber.playground.answers import declared
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")
_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=declared.prompt_for
declared.prompt_for=lambda q, chunks: DOC_LINE+"\n\n"+_orig(q, chunks)
from quber.playground.answers import expectation
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]
sem=asyncio.Semaphore(5)
async def one(r):
    async with sem:
        document_key.set(doc_key)
        try:
            ans=await expectation.answer(r["question"], [chunk(c) for c in r["context_ids"]], "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"],"picks":r["picks"],"context_ids":r["context_ids"],"context_chars":r["context_chars"],
                    "baseline":{"shape":r["shape"],"payload":r["payload"],"cited_ids":r["cited_ids"],"context_chars":r["context_chars"]},"old":r["old"]})
        print(f"[{r['index']:2d}] {out['shape']:12s} {str(out['payload'].get('value') or out['error'] or '')[:24]!r} base={r['shape']}:{str(r['payload'].get('value'))[:18]!r}", flush=True)
        return out
async def main():
    recs=await asyncio.gather(*(one(r) for r in d["rows"]))
    json.dump({"mode":"E","doc_key":doc_key,"doc_id":doc_id,"rows":recs}, open(OUT,"w"), indent=1, default=str)
    print("wrote", OUT)
asyncio.run(main())