"""Origin ZeroGPU retrieval adapter: same scoring/chunking as frozen src/retrieval.py.
Corpus vectors MUST be exported by export_vectors.py from the frozen factory.
Never calculate corpus vectors at Space startup or inside a request.
"""
from __future__ import annotations
import hashlib
import re
import unicodedata
from pathlib import Path
import numpy as np
import pandas as pd
from rank_bm25 import BM25Okapi
from sklearn.feature_extraction.text import ENGLISH_STOP_WORDS
from sklearn.metrics.pairwise import cosine_similarity

MODEL_ID = 'sentence-transformers/paraphrase-multilingual-mpnet-base-v2'


def tokenize(text):
    text = unicodedata.normalize('NFKC', str(text)).casefold()
    tokens = re.findall(r'(?u)\b\w+\b', text)
    return [token for token in tokens if len(token) >= 2 and token not in ENGLISH_STOP_WORDS]


def safe_chunk_text(text, model, overlap=20):
    tokenizer = model.tokenizer
    token_ids = tokenizer.encode(text, add_special_tokens=False, truncation=False)
    special_tokens = tokenizer.num_special_tokens_to_add(pair=False)
    chunks = []
    start = 0
    while start < len(token_ids):
        end = min(start + model.max_seq_length - special_tokens, len(token_ids))
        while end > start:
            chunk = tokenizer.decode(token_ids[start:end], skip_special_tokens=True)
            length = len(tokenizer.encode(chunk, add_special_tokens=True, truncation=False))
            if length <= model.max_seq_length:
                break
            end -= 1
        if end == start:
            raise ValueError('Unable to create a valid chunk')
        chunks.append(chunk)
        if end == len(token_ids):
            break
        start = max(start + 1, end - overlap)
    return chunks


class OriginRetrieval:
    def __init__(self, csv_path: Path, vectors_path: Path):
        csv_path, vectors_path = Path(csv_path), Path(vectors_path)
        self.df = pd.read_csv(csv_path, encoding='utf-8-sig', keep_default_na=False)
        self.df.columns = self.df.columns.str.strip().str.lower()
        self.df['department'] = self.df['department'].astype(str).str.strip()
        if self.df.shape != (103, 7) or not self.df['project_id'].is_unique:
            raise ValueError('Corpus has diverged from frozen 103x7 dataset')
        with np.load(vectors_path, allow_pickle=False) as bundle:
            ids = bundle['project_ids'].astype(str).tolist()
            vectors = bundle['embeddings']
            csv_sha = str(bundle['csv_sha256'].item())
            model_name = str(bundle['model_name'].item())
        if csv_sha != hashlib.sha256(csv_path.read_bytes()).hexdigest():
            raise ValueError('Embeddings CSV fingerprint mismatch: re-export vectors')
        if model_name != MODEL_ID or ids != self.df['project_id'].astype(str).tolist():
            raise ValueError('Embedding model or ID ordering mismatch')
        if vectors.shape != (103, 768) or not np.isfinite(vectors).all():
            raise ValueError('Invalid embedding matrix')
        self.embeddings = vectors.astype(np.float32)
        self.project_texts = (
            self.df['title'].astype(str).str.strip() + '. ' +
            self.df['abstract'].astype(str).str.strip()
        ).str.replace(r'\s+', ' ', regex=True).str.strip().tolist()
        self.bm25 = BM25Okapi([tokenize(t) for t in self.project_texts])
        self.records = {str(r['project_id']): r for r in self.df.to_dict('records')}

    def query_vector(self, model, title, description):
        query = f'{title.strip()}. {description.strip()}'
        chunks = safe_chunk_text(query, model)
        vectors = model.encode(chunks, batch_size=16, convert_to_numpy=True, normalize_embeddings=True)
        vector = vectors.mean(axis=0, keepdims=True)
        vector /= np.linalg.norm(vector, axis=1, keepdims=True)
        return vector

    def search_from_vector(self, query_vector, title, description, lang, top_k=5, rrf_k=60):
        if lang not in {'ar', 'en'}: raise ValueError('Unsupported language')
        scores = cosine_similarity(query_vector, self.embeddings)[0]
        df = self.df
        semantic_order = np.argsort(-scores)
        if lang == 'ar':
            top_indices = semantic_order[:top_k]
            return [{'rank':i+1, 'project_id':str(df.iloc[k]['project_id']), 'retrieval_mode':'semantic_ar'}
                    for i,k in enumerate(top_indices)]
        # Equivalent to frozen bm25_search then hybrid_search using full-corpus semantic ranks.
        query_tokens = tokenize(f'{title.strip()}. {description.strip()}')
        bm25_scores = self.bm25.get_scores(query_tokens)
        matched_indices = np.flatnonzero(bm25_scores > 0)
        bm25_order = matched_indices[np.argsort(-bm25_scores[matched_indices], kind='stable')]
        semantic_ranks = {str(df.iloc[idx]['project_id']): i+1 for i,idx in enumerate(semantic_order)}
        bm25_ranks = {str(df.iloc[idx]['project_id']): i+1 for i,idx in enumerate(bm25_order)}
        ranking = df[['project_id']].copy()
        ranking['semantic_rank'] = ranking['project_id'].map(semantic_ranks)
        ranking['bm25_rank'] = ranking['project_id'].map(bm25_ranks)
        ranking['rrf_score'] = (1.0 / (rrf_k + ranking['semantic_rank']) +
                                (1.0 / (rrf_k + ranking['bm25_rank'])).fillna(0.0))
        ranking = ranking.sort_values(['rrf_score', 'project_id'], ascending=[False, True]).head(top_k)
        mode = 'hybrid_en'
        return [{'rank':i+1, 'project_id':str(pid), 'retrieval_mode':mode}
                for i,pid in enumerate(ranking['project_id'])]

    def public_retrieval(self, ranked, title, description, lang):
        results = []
        for item in ranked:
            row = self.records[item['project_id']]
            def both(s): return {'en':str(s), 'ar':str(s)}
            results.append({'rank':item['rank'],'project':{
                'id':str(row['project_id']), 'title':both(row['title']), 'abstract':both(row['abstract']),
                'department':both(row['department']), 'year':str(row['year']),
                'keywords':str(row['keywords']), 'sourcePage':int(row['source_page']),
                'isFictional':False}})
        return {'retrieval': {'idea':{'title':title,'description':description,'language':lang},
                   'detectedLanguage':lang, 'mode':'arabic_semantic' if lang=='ar' else 'english_hybrid_rrf',
                   'results':results, 'comparisons':[], 'isDemo':False}, 'suggestions':[]}