"""Label-separated retrieval evaluation on a LOCAL official LongMemEval JSON file.

No model/API or dataset download. This is NOT the official end-to-end QA score.
Only role/content/date/session identity are provided to memory. Gold answer,
answer_session_ids and has_answer are not used until after retrieval.
"""
from pathlib import Path
import argparse,json,sys,hashlib,time
ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
from phillnet_recall import RecallMemory,RecallPolicy
from phillnet_recall.evidence import semantic_retrieve


def public_input(row):
    required=('question_id','question','haystack_session_ids','haystack_dates','haystack_sessions')
    if any(k not in row for k in required):raise ValueError('missing official dataset fields')
    if not len(row['haystack_sessions'])==len(row['haystack_session_ids'])==len(row['haystack_dates']):
        raise ValueError('session/date shape mismatch')
    sessions=[]
    for sid,date,turns in zip(row['haystack_session_ids'],row['haystack_dates'],row['haystack_sessions']):
        sanitized=[]
        for turn in turns:
            if turn.get('role') not in ('user','assistant') or not isinstance(turn.get('content'),str):
                raise ValueError('invalid conversation turn')
            sanitized.append({'role':turn['role'],'content':turn['content']})
        sessions.append({'id':str(sid),'date':str(date),'turns':sanitized})
    return {'question_id':str(row['question_id']),'question':row['question'],
            'question_date':str(row.get('question_date','')),'sessions':sessions}


def run_case(inp,policy,backend,top_k):
    start=time.perf_counter()
    # This measures which conversation sessions are recalled, not the truth of
    # their utterances. Each stored item is an observation that someone said X.
    with RecallMemory(policy=policy,max_candidates=8192,max_entries=100000,clock=lambda:1000000.) as m:
        record_session={}
        for session in inp['sessions']:
            for i,turn in enumerate(session['turns']):
                # Deterministic full segmentation, no content is discarded.
                pieces=[turn['content'][j:j+8000] for j in range(0,len(turn['content']),8000)] or ['']
                for j,piece in enumerate(pieces):
                    ident=m.remember(f"Date: {session['date']}\nSpeaker: {turn['role']}\nUtterance: {piece}",
                        source=f"session:{session['id']}/turn:{i}/part:{j}",kind='measurement',valid_from=0)
                    record_session[ident]=session['id']
        if backend=='semantic':r=semantic_retrieve(m,inp['question'],top_k=128,max_scan=100000)
        else:r=m.retrieve(inp['question'],top_k=128,backend=backend)
        selected=[]
        for item in r['items']:
            sid=record_session[item['id']]
            if sid not in selected:selected.append(sid)
            if len(selected)>=top_k:break
    return {'question_id':inp['question_id'],'status':r['status'],'sessions':selected,
            'indexed_parts':len(record_session),'seconds_including_ingestion':time.perf_counter()-start}


def main():
    ap=argparse.ArgumentParser();ap.add_argument('--dataset',required=True,type=Path)
    ap.add_argument('--limit',type=int,default=20);ap.add_argument('--top-k',type=int,default=5)
    ap.add_argument('--output',type=Path,default=ROOT/'reports/longmemeval_retrieval.json');args=ap.parse_args()
    if not args.dataset.is_file():raise SystemExit('Provide a local official dataset JSON; no download is performed.')
    if 'oracle' in args.dataset.name.casefold():raise SystemExit('Oracle-only histories are not accepted for retrieval benchmarking.')
    if not 1<=args.limit<=500 or not 1<=args.top_k<=128:raise SystemExit('limit: 1..500; top-k: 1..128')
    if args.dataset.stat().st_size>2*1024**3:raise SystemExit('Dataset exceeds 2 GiB import limit')
    data=json.loads(args.dataset.read_text())
    if not isinstance(data,list):raise SystemExit('Expected official list of question instances')
    policy=RecallPolicy();rows=[];summary={}
    for backend in ('bm25','hybrid','semantic'):
        correct=any_match=count=abstention=0
        for original in data[:args.limit]:
            inp=public_input(original)
            out=run_case(inp,policy,backend,args.top_k) # Labels are not passed.
            expected=set(str(x) for x in original.get('answer_session_ids',[]))
            is_abstention=str(original['question_id']).endswith('_abs')
            if is_abstention:abstention+=1;all_ok=any_ok=None
            else:
                count+=1;all_ok=bool(expected) and expected.issubset(out['sessions']);any_ok=bool(expected&set(out['sessions']))
                correct+=all_ok;any_match+=any_ok
            rows.append({**out,'backend':backend,'gold_sessions_for_scoring_only':sorted(expected),
                         'all_support_at_k':all_ok,'any_support_at_k':any_ok,'abstention_question':is_abstention})
        summary[backend]={'answerable_questions':count,'all_support_at_k':correct,'any_support_at_k':any_match,'abstention_unscored':abstention}
    report={'status':'completed','dataset_sha256':hashlib.sha256(args.dataset.read_bytes()).hexdigest(),
            'dataset_filename':args.dataset.name,'source_questions':len(data),'evaluated_prefix':min(len(data),args.limit),
            'top_k_sessions':args.top_k,'metric':'session retrieval only; not official end-to-end QA',
            'model_generations':0,'labels_used_during_retrieval':False,'private_data_used':False,
            'ingestion':'all utterances preserved as conversation observations; no LLM extraction; deterministic 8000-character segmentation',
            'date_policy':'dates supplied in evidence text, not hard temporal filtering; no future-session exclusions',
            'summary':summary,'rows':rows}
    args.output.parent.mkdir(parents=True,exist_ok=True);args.output.write_text(json.dumps(report,indent=2)+'\n')
    print(json.dumps(summary,indent=2))
if __name__=='__main__':main()