"""Independent FTS5/BM25 and E5 candidates over witnessed source assertions.

The checkpoint is an evidence store, not an alternative to text retrieval.
Named finite scopes are expanded after ranking so a bounded candidate list
cannot silently remove an assignment goal or one side of a comparison.
"""
from functools import lru_cache
import base64
import sqlite3
import threading
import zlib

from tools.answers.context_policy import normalized

_lock = threading.RLock()


@lru_cache(maxsize=2)
def lexical_index(version, texts):
    connection = sqlite3.connect(':memory:', check_same_thread=False)
    connection.execute('CREATE VIRTUAL TABLE evidence USING fts5(id UNINDEXED, text, tokenize="unicode61 remove_diacritics 2")')
    connection.executemany('INSERT INTO evidence VALUES (?,?)',texts)
    return connection


RERANK_DEPTH = 60
_lemma_memo = {}


def _unit_lemmas(index, data_dir):
    """Lemma strings of all units, read from the disk cache once per checkpoint."""
    from tools.retrieval import lemmatize
    key=(index['version'],str(data_dir))
    with _lock:
        if key not in _lemma_memo:
            _lemma_memo.clear()
            _lemma_memo[key]=lemmatize.LemmaCache(lemmatize.cache_path(data_dir)).many([r['text'] for r in index['units']])
        return _lemma_memo[key]


def retrieve_units(question, index, spec, *, limit=5, data_dir=None):
    from tools.knowledge.query import _terms
    from tools.knowledge.evidence import claim_eligible
    from tools.answers.question_intent import discovery_topic
    search_question=discovery_topic(question) if spec.intent=='discovery' else question
    people={p.name for p in spec.people if p.role=='student'}
    similar=spec.intent=='similar'
    def in_scope(row):
        if not spec.matches_evidence(row):return False
        if spec.period and spec.period.field!='work_year':
            value=row.get(spec.period.field)
            if not value or not spec.period.start<=value<=spec.period.end:return False
        return similar or not people or row['person'] in people
    candidates=[r for r in index['units'] if in_scope(r)]
    eligible={r['id'] for r in candidates}
    from tools.retrieval import lemmatize, rerank
    # Lemmas extend the indexed text ("chatbotom" also matches "chatbot").
    lemma_on=lemmatize.available()
    lemma_strings=_unit_lemmas(index,data_dir) if lemma_on else ['']*len(index['units'])
    texts=tuple((r['id'],r['text']+' '+' '.join(r.get('concepts',[]))+' '+l) for r,l in zip(index['units'],lemma_strings))
    connection=lexical_index((index['version'],lemma_on),texts)
    terms=_terms(search_question)
    if lemma_on and not similar:
        mapped=lemmatize.lemma_map(search_question)
        terms=terms|{mapped[t] for t in terms if t in mapped and len(mapped[t])>=3}
    if similar:
        # The reference's actual substantive topics seed the peer search.
        from tools.knowledge.concepts import substantive_concept
        terms={normalized(c) for r in candidates if r['person'] in people and claim_eligible(r)
               for c in r.get('concepts',[]) if substantive_concept(c)} or terms
    expression=' OR '.join('"'+term.replace('"','""')+'"*' for term in sorted(terms))
    with _lock:
        lexical=[r[0] for r in connection.execute(
            'SELECT id FROM evidence WHERE evidence MATCH ? ORDER BY bm25(evidence),id',
            (expression,)) if r[0] in eligible] if expression else []
    semantic={}
    if index.get('vectors'):
        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('Content 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 content vector count')
        anchors=[i for i,r in enumerate(index['units']) if r['person'] in people and in_scope(r) and claim_eligible(r)] if similar else []
        vector=np.mean(matrix[anchors],axis=0) if anchors else embed_query(search_question)
        norm=float(np.linalg.norm(vector))
        if not norm:raise ValueError('Invalid query/reference embedding')
        scores=matrix @ (vector/norm)
        semantic={r['id']:float(s) for r,s in zip(index['units'],scores) if r['id'] in eligible}
    vector_rank=sorted(semantic,key=lambda key:(-semantic[key],key))
    # A cross-student question needs a broader pool than a named person.
    budget=min(256,max(160 if spec.intent=="discovery" else 80,limit*16))
    fused={}
    for ranking in (lexical[:budget],vector_rank[:budget]):
        for position,key in enumerate(ranking,1):fused[key]=fused.get(key,0)+1/(60+position)
    ranked=sorted(fused,key=lambda key:(-fused[key],key))[:budget]
    selected=set(ranked)
    # Whole named source/work scopes are finite, explicit evidence expansion;
    # this is not an assumption that a global top-k result is exhaustive.
    if people:
        selected.update(r['id'] for r in candidates if r['person'] in people)
    rows=[r for r in candidates if r['id'] in selected]
    rows.sort(key=lambda r:(-fused.get(r['id'],0),r['source_path'],r['line']))
    # The cross-encoder reorders the strongest fused candidates; nothing is
    # added or removed, and a named scope keeps every one of its units.
    reranked=False
    if rows and not similar and rerank.available():
        # Reciprocal-rank fusion of the retrieval order and the cross-encoder
        # order: the reranker corrects the fused list instead of replacing it
        # (alone it promoted "hate speech" units for "speech recognition").
        head=rows[:RERANK_DEPTH]
        query=getattr(spec,'topic',None) or search_question
        values=rerank.scores(query,[r['text'] for r in head])
        if values is not None:
            by_rerank=sorted(range(len(head)),key=lambda i:(-values[i],i))
            fused_rank={i:1/(60+pos) for pos,i in enumerate(range(len(head)),1)}
            for pos,i in enumerate(by_rerank,1):fused_rank[i]+=1/(60+pos)
            head=[head[i] for i in sorted(fused_rank,key=lambda i:(-fused_rank[i],i))]
            rows=head+rows[RERANK_DEPTH:]
            reranked=True
    return rows,semantic,{'strategy':'bm25_e5_rrf_assertions'+('_lemma' if lemma_on else '')+('_rerank' if reranked else ''),
        'candidate_count':len(rows),
        'lexical_candidates':len(lexical[:budget]),'vector_candidates':len(vector_rank[:budget]),
        'named_scope_expansion':bool(people),'bounded':not people and len(candidates)>len(rows)}