"""Bounded, read-only graph neighborhoods with checkpoint-validated provenance."""
from __future__ import annotations

from dataclasses import dataclass, field
import json

from neo4j import READ_ACCESS

from tools.knowledge import store
from tools.knowledge.concepts import concept_key, substantive_concept
from tools.knowledge.evidence import claim_eligible, relation_has_local_support
from tools.graphrag.graph_db import create_neo4j_driver, get_neo4j_settings
from tools.answers.context_policy import normalized

TIMEOUT = 8
SEED_LIMIT = 32
CONCEPT_LIMIT = 24
EVIDENCE_LIMIT = 800

# Fixed templates only. User/model strings enter exclusively as parameters.
ROOT_QUERY = '''MATCH (p:KnowledgePerson {graph_scope:$scope})-[:HAS_EVIDENCE]->(e:KnowledgeEvidence {graph_scope:$scope})
OPTIONAL MATCH (p)-[:HAS_WORK]->(w:KnowledgeWork {graph_scope:$scope})-[:EVIDENCED_BY]->(e)
WITH p,w,e WHERE (p.source_path IN $profiles OR
    EXISTS { MATCH (w)-[:SUPERVISED_BY]->(v:KnowledgeAdvisor {graph_scope:$scope}) WHERE v.id IN $supervisors } OR
    EXISTS { MATCH (w)-[:CONSULTED_BY]->(v:KnowledgeAdvisor {graph_scope:$scope}) WHERE v.id IN $consultants })
AND ((w.id IN $works AND e.work_id=w.id) OR ($profile_records AND e.work_id IS NULL))
AND ($date_field IS NULL OR ($start <= e[$date_field] AND e[$date_field] <= $end))
OPTIONAL MATCH (w)-[:HAS_ACTIVITY]->(a:KnowledgeActivity {graph_scope:$scope})-[:SUPPORTED_BY]->(e)
OPTIONAL MATCH (e)-[:ABOUT]->(c:KnowledgeConcept {graph_scope:$scope})
RETURN properties(p) AS person,properties(w) AS work,properties(e) AS evidence,a.id AS activity,
       collect(DISTINCT c.key) AS concepts
ORDER BY CASE WHEN evidence.kind='work_record' THEN 0 ELSE 1 END,work.year DESC,evidence.line
LIMIT $budget'''

SEED_QUERY = '''MATCH (p:KnowledgePerson {graph_scope:$scope})-[:HAS_WORK]->(w:KnowledgeWork {graph_scope:$scope})
    -[:EVIDENCED_BY]->(e:KnowledgeEvidence {graph_scope:$scope})
WHERE e.id IN $ids AND e.work_id=w.id AND w.id IN $works
  AND ($date_field IS NULL OR ($start <= e[$date_field] AND e[$date_field] <= $end))
OPTIONAL MATCH (w)-[:HAS_ACTIVITY]->(a:KnowledgeActivity {graph_scope:$scope})-[:SUPPORTED_BY]->(e)
OPTIONAL MATCH (e)-[:ABOUT]->(c:KnowledgeConcept {graph_scope:$scope})
RETURN properties(p) AS person,properties(w) AS work,properties(e) AS evidence,a.id AS activity,
       collect(DISTINCT c.key) AS concepts LIMIT $budget'''
PERIOD_QUERY = SEED_QUERY.replace('e.id IN $ids AND ', '')

NEIGHBOR_QUERY = '''MATCH (p:KnowledgePerson {graph_scope:$scope})-[:HAS_WORK]->(w:KnowledgeWork {graph_scope:$scope})
    -[:EVIDENCED_BY]->(e:KnowledgeEvidence {graph_scope:$scope})-[:ABOUT]->(c:KnowledgeConcept {graph_scope:$scope})
WHERE c.key IN $keys AND NOT p.source_path IN $exclude AND w.id IN $works
  AND e.work_id=w.id AND e.actor=p.name AND e.annotation_status='knowledge'
  AND e.actor_status IN ['explicit_owner','profile_context']
  AND ($date_field IS NULL OR ($start <= e[$date_field] AND e[$date_field] <= $end))
OPTIONAL MATCH (w)-[:HAS_ACTIVITY]->(a:KnowledgeActivity {graph_scope:$scope})-[:SUPPORTED_BY]->(e)
WITH p,w,e,a,collect(DISTINCT c.key) AS concepts
ORDER BY CASE WHEN e.phase IN ['result','progress'] THEN 0 ELSE 1 END,e.record_date DESC,e.line
WITH p,collect({work:properties(w),evidence:properties(e),activity:a.id,concepts:concepts})[..4] AS hits,
     reduce(shared=[],group IN collect(concepts) | shared + [key IN group WHERE NOT key IN shared]) AS shared
ORDER BY size(shared) DESC,CASE WHEN any(hit IN hits WHERE hit.evidence.phase IN ['result','progress']) THEN 0 ELSE 1 END,
         p.source_path LIMIT $peers
UNWIND hits AS hit
RETURN properties(p) AS person,hit.work AS work,hit.evidence AS evidence,hit.activity AS activity,hit.concepts AS concepts'''

RELATION_QUERY = '''MATCH (r:KnowledgeRelation {graph_scope:$scope,origin:'llm_checked',verification:'supported'})
    -[:EVIDENCED_BY]->(e:KnowledgeEvidence {graph_scope:$scope}),
    (r)-[:SUBJECT]->(s:KnowledgeConcept {graph_scope:$scope}),
    (r)-[:OBJECT]->(o:KnowledgeConcept {graph_scope:$scope})
WHERE e.id IN $ids AND r.evidence=e.id AND r.work_id=e.work_id AND r.source_path=e.source_path
RETURN properties(r) AS relation,e.id AS evidence,s.key AS subject,o.key AS object LIMIT $budget'''

EXPAND_QUERY = '''MATCH (r:KnowledgeRelation {graph_scope:$scope,origin:'llm_checked',verification:'supported'})
    -[:EVIDENCED_BY]->(e:KnowledgeEvidence {graph_scope:$scope}),
    (r)-[:SUBJECT]->(s:KnowledgeConcept {graph_scope:$scope}),
    (r)-[:OBJECT]->(o:KnowledgeConcept {graph_scope:$scope}),
    (p:KnowledgePerson {graph_scope:$scope})-[:HAS_WORK]->(w:KnowledgeWork {graph_scope:$scope})-[:EVIDENCED_BY]->(e)
WHERE (s.key IN $keys OR o.key IN $keys) AND w.id IN $works AND r.evidence=e.id
  AND ($any_owner OR p.source_path IN $profiles)
  AND e.work_id=w.id AND r.work_id=w.id AND r.source_path=p.source_path
  AND ($date_field IS NULL OR ($start <= e[$date_field] AND e[$date_field] <= $end))
OPTIONAL MATCH (w)-[:HAS_ACTIVITY]->(a:KnowledgeActivity {graph_scope:$scope})-[:SUPPORTED_BY]->(e)
RETURN properties(p) AS person,properties(w) AS work,properties(e) AS evidence,a.id AS activity,
       [s.key,o.key] AS concepts,properties(r) AS relation,e.id AS evidence_id,s.key AS subject,o.key AS object
LIMIT $budget'''


@dataclass
class LocalGraph:
    units: list[dict] = field(default_factory=list)
    paths: list[dict] = field(default_factory=list)
    relations: list[dict] = field(default_factory=list)
    seeds: list[dict] = field(default_factory=list)
    truncation_reasons: list[str] = field(default_factory=list)


def validated_row(record, by_id):
    actual = dict(record['evidence'])
    original = by_id.get(actual.get('id'))
    if not original:
        raise ValueError('Unknown graph evidence')
    fields = ('person','actor','actor_status','source_path','source_url','text','line','end_line',
              'work_id','work_label','work_type','work_year','phase','section_scope','kind',
              'annotation_status','record_date','event_date','event_year','target_date','target_year')
    if any(actual.get(k) != original.get(k) for k in fields):
        raise ValueError('Graph assertion differs from its checkpoint')
    person, work = dict(record['person']), dict(record['work'] or {})
    if person.get('source_path') != original['source_path'] or person.get('name') != original['person']:
        raise ValueError('Cross-person graph edge')
    if original.get('work_id') and any(work.get(k) != original.get(v) for k,v in
            [('id','work_id'),('source_path','source_path'),('label','work_label'),('type','work_type'),('year','work_year')]):
        raise ValueError('Cross-work graph edge')
    if record.get('activity') and record['activity'] != original['id']:
        raise ValueError('Cross-activity graph edge')
    if (original.get('annotation_status')=='knowledge' and original.get('kind') in {'assignment','plan','progress','result'}
            and claim_eligible(original) and record.get('activity')!=original['id']):
        raise ValueError('Missing work-activity-evidence path')
    keys = set(record['concepts'])
    if not keys <= {concept_key(c) for c in original.get('concepts',[])}:
        raise ValueError('Unverified concept edge')
    return {**original, '_graph_keys': sorted(keys), '_graph_activity': record.get('activity')}


def seed_ids(question, index, spec):
    """Classical retrieval proposes entry points; Neo4j must prove every path."""
    from tools.knowledge.query import _terms, activity_supported, concept_families
    from tools.answers.content_selection import administrative_only
    terms = _terms(question); families = concept_families(question)
    candidates = [u for u in index['units'] if u.get('work_id') and claim_eligible(u)
                  and u.get('annotation_status')=='knowledge' and not administrative_only(u['text'])
                  and spec.matches_work(u['person'],u.get('work_label'),u.get('work_type'),u.get('work_year'))
                  and activity_supported(question,u)]
    if spec.period and spec.period.field!='work_year':
        candidates=[u for u in candidates if u.get(spec.period.field) and
                    spec.period.start<=u[spec.period.field]<=spec.period.end]
    scores = {}
    if index.get('vectors') and candidates:
        import base64, zlib
        import numpy as np
        from tools.rag.embedding_utils import embed_query, model_profile, embedding_model_name
        if model_profile(embedding_model_name())['fingerprint'] != index['embedding_fingerprint']:
            raise RuntimeError('Graph seed embeddings belong to another model')
        matrix=np.frombuffer(zlib.decompress(base64.b64decode(index['vectors'])),dtype='<f4').reshape(index['vector_shape'])
        if len(matrix)!=len(index['units']):raise ValueError('Invalid seed vector count')
        scores={u['id']:float(s) for u,s in zip(index['units'],matrix @ embed_query(question))}
    ranked=[]
    for row in candidates:
        text_terms=_terms(row['text']+' '+' '.join(row.get('concepts',[])))
        hits=sum(any(a.startswith(b) or b.startswith(a) for a in text_terms) for b in terms)
        equivalent=bool(families & concept_families(row['text']+' '+' '.join(row.get('concepts',[]))))
        semantic=scores.get(row['id'],0)
        if hits or equivalent or semantic>=.74:
            ranked.append((hits+int(equivalent)+semantic,row))
    ranked.sort(key=lambda item:(-item[0],item[1]['source_path'],item[1]['line']))
    return [row['id'] for _,row in ranked[:SEED_LIMIT+1]]


def validated_relation(record, by_id):
    row=by_id[record.get('evidence_id') or record['evidence']]; actual=dict(record['relation'])
    match=next((r for r in row.get('relations',[]) if r.get('verification')=='supported'
        and concept_key(r['subject'])==record['subject'] and concept_key(r['object'])==record['object']
        and r['predicate']==actual.get('predicate') and r['quote']==actual.get('quote')
        and relation_has_local_support(row,r)),None)
    if not match or any(actual.get(k)!=row.get(k) for k in ('actor','work_id','phase','source_path','source_url','line')):
        raise ValueError('Unverified semantic graph relation')
    return {'evidence_id':row['id'],**match}


def retrieve_local_graph(question, index, bound, *, limit=5):
    spec=bound.spec; result=LocalGraph(); by_id={u['id']:u for u in index['units']}
    people=[p for p in spec.people if p.role=='student']
    profiles=[p.source_document for p in people if p.source_document]
    works=sorted({u['work_id'] for u in index['units'] if u.get('work_id') and
        spec.matches_work(u['person'],u.get('work_label'),u.get('work_type'),u.get('work_year'))})
    ids=seed_ids(question,index,spec) if not profiles and not any(p.role in {'supervisor','consultant'} for p in spec.people) else []
    if len(ids)>SEED_LIMIT:result.truncation_reasons.append('graph_seed_limit')
    ids=ids[:SEED_LIMIT]
    period=spec.period if spec.period and spec.period.field!='work_year' else None
    params={'scope':store.SCOPE,'profiles':profiles,'works':works,'ids':ids,
        'supervisors':[store.identifier(normalized(p.name)) for p in spec.people if p.role=='supervisor'],
        'consultants':[store.identifier(normalized(p.name)) for p in spec.people if p.role=='consultant'],
        'date_field':period.field if period else None,'start':period.start if period else None,'end':period.end if period else None,
        'profile_records':not any(w.work_type or w.years or w.labels or w.numbers or
                                 w.excluded_labels or w.excluded_numbers or w.unnumbered for w in spec.works),
        'budget':EVIDENCE_LIMIT+1,'peers':min(61,max(11,limit*3+1))}
    settings=get_neo4j_settings()
    with create_neo4j_driver(settings,connection_timeout=TIMEOUT,max_transaction_retry_time=0) as driver,\
            driver.session(database=settings.database,default_access_mode=READ_ACCESS) as session, session.begin_transaction(timeout=TIMEOUT*4) as tx:
        state=tx.run('MATCH (s:KnowledgeState {scope:$scope}) RETURN properties(s) AS state',scope=store.SCOPE).single()
        state=dict(state['state']) if state else {}
        if (state.get('version')!=index['version'] or state.get('projection_schema')!=store.CONTENT_GRAPH_SCHEMA
                or state.get('index_version_json')!=json.dumps(index['index_version'],sort_keys=True)
                or index.get('_checkpoint_sha256') and state.get('checkpoint_sha256')!=index['_checkpoint_sha256']):
            raise RuntimeError('Content graph projection is stale or differs from checkpoint')
        def fetch(query, **extra):
            return [validated_row(dict(r),by_id) for r in tx.run(query,**{**params,**extra})]
        root=fetch(ROOT_QUERY) if profiles or params['supervisors'] or params['consultants'] else fetch(
            PERIOD_QUERY if spec.intent=='period' else SEED_QUERY)
        if len(root)>EVIDENCE_LIMIT:result.truncation_reasons.append('graph_evidence_limit')
        root=root[:EVIDENCE_LIMIT]
        result.seeds=root
        result.units=root
        if spec.intent in {'similar','discovery','period'}:
            ranked={}
            for row in root:
                if row.get('annotation_status')!='knowledge' or not claim_eligible(row):continue
                for key in row['_graph_keys']:
                    if spec.intent!='similar' or substantive_concept(key):
                        ranked.setdefault(key,[]).append(row)
            keys=sorted(ranked,key=lambda k:(-len(ranked[k]),k))[:CONCEPT_LIMIT]
            if len(ranked)>CONCEPT_LIMIT:result.truncation_reasons.append('graph_concept_limit')
            expansions=list(tx.run(EXPAND_QUERY,**{**params,'keys':keys,
                                  'any_owner':spec.intent!='similar','budget':SEED_LIMIT+1})) if keys else []
            if len(expansions)>SEED_LIMIT:result.truncation_reasons.append('graph_relation_limit')
            for record in expansions[:SEED_LIMIT]:
                row=validated_row(dict(record),by_id)
                relation=validated_relation(dict(record),by_id)
                result.relations.append(relation)
                for key in row['_graph_keys']:
                    if (spec.intent!='similar' or substantive_concept(key)) and key not in keys:
                        if len(keys)>=CONCEPT_LIMIT+8:
                            result.truncation_reasons.append('graph_concept_limit');continue
                        keys.append(key);ranked.setdefault(key,[]).append(row)
                root.append(row)
            peers=fetch(NEIGHBOR_QUERY,keys=keys,exclude=profiles if spec.intent=='similar' else []) if keys else []
            if len({r['source_path'] for r in peers})>=params['peers']:
                result.truncation_reasons.append('graph_peer_limit')
            for row in peers:
                for key in row['_graph_keys']:
                    if key not in keys:continue
                    anchors=[a for a in ranked[key] if spec.intent!='similar' or a['source_path']!=row['source_path']]
                    for anchor in sorted(anchors,key=lambda a:(a['phase'] not in {'result','progress'},a['line']))[:2]:
                        result.paths.append({'kind':'shared_topic','concept':key,'anchor_id':anchor['id'],
                            'peer_id':row['id'],'anchor_work_id':anchor['work_id'],'peer_work_id':row['work_id'],
                            'anchor_activity_id':anchor.get('_graph_activity'),'peer_activity_id':row.get('_graph_activity'),
                            'concept_match':'approved_equivalent_key'})
            result.units=list({r['id']:r for r in [*root,*peers]}.values())
        records=tx.run(RELATION_QUERY,scope=store.SCOPE,
                       ids=[u['id'] for u in result.units],budget=EVIDENCE_LIMIT+1)
        for record in records:
            relation=validated_relation(dict(record),by_id)
            if relation not in result.relations:result.relations.append(relation)
        tx.commit()
    return result