import os
import json
import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer
from typing import List, Tuple
from src.ontology.models import OntologyRecord

from src.runtime.device_manager import DeviceManager

class ONNXEmbeddingEngine:
    def __init__(self, model_path: str = "JairoDanielMT/paraphrase-multilingual-MiniLM-L12-v2-onnx-int8", index_dir: str = "data/faiss_indices_onnx"):
        self.model_path = model_path
        self.index_dir = index_dir
        self.tokenizer = None
        self.session = None
        self.indices = {}
        self.records = {}
        self.provider = DeviceManager.get_optimal_device()
        self.batch_size = DeviceManager.get_optimal_batch_size()
        
    def _init_model(self):
        if not self.tokenizer:
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
        if not self.session:
            from optimum.onnxruntime import ORTModelForFeatureExtraction
            self.model = ORTModelForFeatureExtraction.from_pretrained(
                self.model_path,
                file_name="model_quantized.onnx",
                provider=self.provider
            )
            # We assign the underlying inference session to maintain compatibility with the rest of the code
            self.session = self.model.model
            
    def load_index(self):
        import faiss
        if not os.path.exists(self.index_dir):
            raise FileNotFoundError(f"Index directory not found at {self.index_dir}")
            
        for filename in os.listdir(self.index_dir):
            if filename.endswith(".index"):
                category = filename.replace(".index", "")
                index_path = os.path.join(self.index_dir, filename)
                mapping_path = os.path.join(self.index_dir, f"{category}_records.json")
                
                if os.path.exists(mapping_path):
                    self.indices[category] = faiss.read_index(index_path)
                    with open(mapping_path, "r", encoding="utf-8") as f:
                        data = json.load(f)
                        self.records[category] = [OntologyRecord(**item) for item in data]
        self._init_model()

    def build_index(self, ontology_dir: str):
        import faiss
        self._init_model()
        os.makedirs(self.index_dir, exist_ok=True)
        
        for filename in os.listdir(ontology_dir):
            if filename.endswith(".json"):
                category = filename.replace(".json", "")
                path = os.path.join(ontology_dir, filename)
                
                cat_records = []
                texts = []
                
                with open(path, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    for item in data:
                        record = OntologyRecord(**item)
                        cat_records.append(record)
                        texts.append(record.canonical)
                        for alias in record.aliases:
                            cat_records.append(record)
                            texts.append(alias)
                
                if not texts:
                    continue
                    
                embeddings = self.encode(texts)
                
                dimension = embeddings.shape[1]
                index = faiss.IndexFlatL2(dimension)
                index.add(embeddings)
                
                index_path = os.path.join(self.index_dir, f"{category}.index")
                faiss.write_index(index, index_path)
                
                mapping_path = os.path.join(self.index_dir, f"{category}_records.json")
                with open(mapping_path, "w", encoding="utf-8") as f:
                    json.dump([r.model_dump() for r in cat_records], f, ensure_ascii=False)
                    
                self.indices[category] = index
                self.records[category] = cat_records

    def encode(self, texts: List[str], batch_size: int = 128) -> np.ndarray:
        self._init_model()
        all_embeddings = []
        import tqdm
        
        for i in tqdm.tqdm(range(0, len(texts), batch_size), desc="Encoding batches (ONNX)"):
            batch_texts = texts[i:i + batch_size]
            inputs = self.tokenizer(batch_texts, padding=True, truncation=True, return_tensors="np")
            outputs = self.session.run(None, {
                "input_ids": inputs["input_ids"],
                "attention_mask": inputs["attention_mask"],
                "token_type_ids": inputs["token_type_ids"]
            })
            token_embeddings = outputs[0]
            attention_mask = inputs["attention_mask"]
            
            # Mean pooling
            input_mask_expanded = np.broadcast_to(np.expand_dims(attention_mask, -1), token_embeddings.shape)
            sum_embeddings = np.sum(token_embeddings * input_mask_expanded, 1)
            sum_mask = np.clip(np.sum(input_mask_expanded, 1), a_min=1e-9, a_max=None)
            embeddings = sum_embeddings / sum_mask
            
            # Normalize
            norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
            embeddings = embeddings / norms
            all_embeddings.append(embeddings.astype('float32'))
            
        return np.vstack(all_embeddings)

    def calibrate_score(self, l2_distance: float) -> float:
        cos_sim = 1.0 - (l2_distance / 2.0)
        k = 15.0
        x0 = 0.8
        conf = 1.0 / (1.0 + np.exp(-k * (cos_sim - x0)))
        return float(np.clip(conf, 0.0, 1.0))

    def search(self, query: str, category: str = None, top_k: int = 5) -> List[Tuple[OntologyRecord, float]]:
        if not self.indices:
            self.load_index()
            
        query_vector = self.encode([query])
        results = []
        
        categories_to_search = [category] if category and category in self.indices else self.indices.keys()
        
        for cat in categories_to_search:
            distances, indices = self.indices[cat].search(query_vector, top_k)
            for dist, idx in zip(distances[0], indices[0]):
                if idx < len(self.records[cat]):
                    conf = self.calibrate_score(float(dist))
                    results.append((self.records[cat][idx], conf))
                    
        results.sort(key=lambda x: x[1], reverse=True)
        return results[:top_k]