import asyncio, sys
from quber.playground.retrieval import retrieve
from quber.playground.db import connect
QS = [
 "Find Portfolio Details: Concentration of Credit Risk - collateral underlying commercial real estate loans, Collateral Property Type: Office (% of Principal). Return Q1 2026 value",
 "Find  Concentration of Credit Risk - collateral underlying commercial real estate loans, Geography: Florida (% of Principal). . Return Q1 2026 value",
 "Find  Concentration of Credit Risk - collateral underlying commercial real estate loans, Geography: United Kingdom (% of Principal). . Return Q1 2026 value",
]
label = sys.argv[1]
with connect() as c:
    row = c.execute("select content from ade_playground.chunks ch join ade_playground.documents d on d.id=ch.document_id where d.doc_key='998c8ed06c0be22d' and ch.chunk_type='text' and ch.page=22 and content ilike '%concentration%' limit 1").fetchone()
    txt = row[0]; i = txt.find('Concentration'); seg = txt[i:i+30]
    print(label, "intro phrase bytes:", seg.encode('utf-8'))
    hits = c.execute("select count(*) from ade_playground.chunks ch join ade_playground.documents d on d.id=ch.document_id where d.doc_key='998c8ed06c0be22d' and ch.page=22 and content ilike '%Concentration of Credit Risk%'").fetchone()[0]
    print(label, "page-23 chunks containing the phrase 'Concentration of Credit Risk':", hits)
for q in QS:
    chunks = asyncio.run(retrieve('998c8ed06c0be22d', q, 10))
    ids = [f"{c.chunk_id}@p{c.page+1}" for c in chunks]
    print(label, "|", q.split('Geography: ')[-1].split('Collateral Property Type: ')[-1][:22], "|", ids)