"""Supplementary LLM check with graph context supplied by the local harness.

This exercises the school's OpenWebUI chat API and generated claims, NOT the
WebUI browser, installed tools, or autonomous routing. Gold never enters a
request. The primary experiment remains the complete frozen retrieval TEST.
"""
from __future__ import annotations

import argparse
import json
import os
import time
from collections import defaultdict
from pathlib import Path

from tools.rag.evaluation.rag_runner import DEFAULT_MODEL, OPENWEBUI_URL, request_json
from ..graph_db import create_neo4j_driver, get_neo4j_settings
from ..graph_search import search_graph_with_session
from .benchmark_v2 import BASE, academics, supporting_documents
from .reproducibility import graph_fingerprint, run_metadata, verify_freeze
from .source_oracle import load_source_oracle, norm
from .source_validation import claim_issues
from tools.answer_evaluation.context_policy import citation_check

SYSTEM = '''Odpovedaj po slovensky výhradne podľa dodaného výsledku grafového vyhľadávania.
Nedoplňuj vedomosti ani údaje z otázky bez potvrdenia v kontexte. Pri no_answer oznám,
že údaj nebol nájdený; nevymýšľaj názov. Zachovaj presné znenie mien, názvov, tagov a kategórií.
Rozlišuj rok práce (year) od začiatku štúdia (start_year), autora dokumentu od študenta.
Vráť JSON objekt s poľami answer (stručný prirodzený text), status (ok/no_answer/ambiguous),
claims (zoznam objektov s person a len relevantnými faktami: title, year, start_year,
keyword, category, document_author), source_urls (zoznam citovaných URL), count (číslo alebo null).
Každá claim musí patriť jednej osobe; nespoj podmienky z rôznych osôb. Pri porovnaní
uveď samostatné claims pre každú osobu; pri počte iba zodpovedajúcich kandidátov.
Pri otázke na názov uveď aj rok, ak je známy. Pri no_answer vráť prázdne claims.
V odpovedi aj source_urls cituj skutočné relevantné zdroje z kontextu.'''


def grade_claims(case: dict, answer: dict, oracle: dict) -> dict:
    graph = case['graph']; claims = answer.get('claims', [])
    if not isinstance(claims, list) or any(not isinstance(c, dict) for c in claims):
        return {'strict': False, 'reason': 'invalid_claim_schema'}
    fields = ['person', 'title', 'year', 'start_year', 'keyword', 'category', 'document_author']
    def matches(binding, claim):
        return all(norm(claim.get(k, '')) == norm(binding[k]) for k in fields if k in binding)
    covered = all(any(matches(b, c) for c in claims) for b in graph['expected_bindings'])
    names = {norm(name): name for name in oracle['people']}
    unsupported = []
    for claim in claims:
        person = names.get(norm(claim.get('person')))
        valid = person is not None
        if valid and 'start_year' in claim:
            valid = norm(claim['start_year']) == norm(oracle['people'][person]['start_year'])
        metadata = {k: claim[k] for k in ['keyword', 'category', 'document_author'] if claim.get(k) is not None}
        if valid and metadata:
            valid = all(supporting_documents(oracle, person, **{k: v}) for k, v in metadata.items())
            if 'category' in metadata and 'document_author' in metadata:
                valid = valid and bool(supporting_documents(oracle, person, **metadata))
        if valid and any(claim.get(k) is not None for k in ['title', 'year']):
            records = academics(oracle, person)
            valid = any(all(norm(r.get(k)) == norm(claim[k]) for k in ['title', 'year']
                            if claim.get(k) is not None) for r in records)
        if not valid:
            unsupported.append(claim)
        elif claim_issues(claim, oracle):
            unsupported.append(claim)
    urls = answer.get('source_urls', [])
    sources_ok = isinstance(urls, list) and all(any(a['source_url'] in urls for a in proof['any_of'])
                       for proof in graph.get('expected_evidence_documents', []))
    sources_ok = sources_ok and citation_check(answer,{
        doc['source_url'] for doc in oracle.get('documents',{}).values() if doc.get('source_url')})['citations_ok']
    negative = graph.get('expected_no_answer', False)
    status_ok = (answer.get('status') == 'no_answer' and not claims) if negative else answer.get('status') == 'ok'
    count_ok = 'expected_count' not in graph or answer.get('count') == graph['expected_count']
    set_ok = True
    if graph.get('expected_entity_set') is not None:
        field = graph['entity_set_field']
        actual = {norm(c[field]) for c in claims if c.get(field) is not None}
        set_ok = actual == {norm(v) for v in graph['expected_entity_set']}
    return {'strict': bool(covered and not unsupported and sources_ok and status_ok and count_ok and set_ok),
            'bindings_ok': covered, 'unsupported_claims': unsupported, 'sources_ok': sources_ok,
            'status_ok': status_ok, 'count_ok': count_ok, 'set_ok': set_ok}


def main():
    # Import the common adapter here: it imports grade_claims from this module.
    from tools.answer_evaluation.evaluate_graph_answers import graph_context, grade_answer, HERE
    from tools.answer_evaluation.checked_generation import generate_checked_answer
    from tools.graphrag.build_graphrag import graph_version

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--questions', type=Path, required=True)
    parser.add_argument('--freeze-manifest', type=Path, required=True)
    parser.add_argument('--audit', type=Path, required=True)
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--per-category', type=int, default=2)
    parser.add_argument('--model', default=DEFAULT_MODEL)
    args = parser.parse_args()
    if args.output.exists():
        parser.error('Refusing to overwrite a previous external measurement.')
    if not os.getenv('OPENWEBUI_API_KEY'):
        parser.error('OPENWEBUI_API_KEY is not configured.')
    questions = json.loads(args.questions.read_text()); selected = []; counts = defaultdict(int)
    for case in questions:
        if case['split'] == 'test' and counts[case['category']] < args.per_category:
            selected.append(case); counts[case['category']] += 1
    oracle = load_source_oracle(); manifest = json.loads(args.freeze_manifest.read_text())
    settings = get_neo4j_settings(); rows = []
    prompt = (HERE / 'graph_answer_prompt.txt').read_text()
    payload = {'method': 'Supplementary known-query check with common context guards; no autonomous tool routing.',
               'model': args.model, 'selected_ids': [q['id'] for q in selected], 'system_prompt': prompt,
               'corrective_request_limit': 1,
               'limitations': ['Small stratified supplementary sample; not the primary 320-case accuracy.',
                   'Structured claims, source citations and binding checks; free prose requires separate review.'],
               'questions': rows}
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with create_neo4j_driver(settings) as driver, driver.session(database=settings.database) as session:
        metadata = run_metadata(args.questions)
        verify_freeze(manifest, metadata, graph_fingerprint(session), audit=args.audit, **manifest['parameters'])
        payload['run_metadata'] = metadata
        version = graph_version(session)
        for case in selected:
            started = time.perf_counter()
            result = search_graph_with_session(session, case['question'], **manifest['parameters'])
            result['index_version'] = version
            context = graph_context(result)
            row = {'id': case['id'], 'category': case['category'], 'question': case['question'],
                   'graph_context': context, 'strict': False}
            try:
                def complete(messages):
                    return request_json(OPENWEBUI_URL, method='POST', timeout=90, max_attempts=1,
                        headers={'Authorization': f'Bearer {os.environ["OPENWEBUI_API_KEY"]}', 'Content-Type': 'application/json'},
                        payload={'model': args.model, 'stream': False, 'response_format': {'type': 'json_object'},
                                 'messages': messages})
                row.update(generate_checked_answer(case['question'], context, prompt, complete))
                if row.get('answer'):
                    row.update(grade_answer(case, row['answer'], oracle, context=context))
                row['strict'] = bool(row['strict'] and row['accepted_for_delivery'])
            except Exception as error:
                row['error'] = f'{type(error).__name__}: {str(error)[:500]}'
            row['latency_seconds'] = time.perf_counter() - started
            rows.append(row)
            payload['summary'] = {'questions': len(rows), 'correct': sum(r['strict'] for r in rows),
                                  'accuracy': sum(r['strict'] for r in rows) / len(rows)}
            args.output.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + '\n')
            print(f'{case["id"]}: {"OK" if row["strict"] else "FAIL"} ({row["latency_seconds"]:.1f}s)', flush=True)
        verify_freeze(manifest, run_metadata(args.questions), graph_fingerprint(session), audit=args.audit, **manifest['parameters'])
        assert graph_version(session) == version, 'Index version changed during generation.'
    print(json.dumps(payload['summary']))


if __name__ == '__main__':
    main()