import os
from typing import Dict, List

import numpy as np
import torch
from transformers import AutoTokenizer, AutoModel
from huggingface_hub import InferenceClient


class RAGEngine:
    def __init__(
        self,
        documents: List[Dict[str, str]],
        embedding_model: str = "BAAI/bge-m3",
        llm_model: str = "meta-llama/Llama-3.1-8B-Instruct",
        batch_size: int = 64,
    ):
        """
        Initialise le moteur RAG avec les documents (contenant chacun 'url' et 'text'),
        les paramètres de configuration et les clients nécessaires.

        Args:
            documents: Liste de documents, chacun un dictionnaire avec les clés 'url' et 'text'.
            embedding_model: Nom du modèle pour calculer les embeddings en local.
            llm_model: Nom du modèle LLM pour les complétions.
            batch_size: Nombre de documents à traiter par lot.
        """
        self.documents = documents
        self.embedding_model = embedding_model  # Nom du modèle pour embeddings (local)
        self.llm_model = llm_model
        self.batch_size = batch_size
        self.embeddings: List[List[float]] = []
        # Filtrer les documents dont le texte est vide pour éviter les erreurs
        self.indexed_documents = [doc for doc in self.documents if doc["text"].strip()]

        # Initialiser le modèle et le tokenizer en local pour le calcul des embeddings
        self.embedding_tokenizer = AutoTokenizer.from_pretrained(self.embedding_model)
        self.embedding_model_local = AutoModel.from_pretrained(self.embedding_model)

        # Initialiser le client pour le LLM (l'inférence reste à distance pour le LLM)
        self._init_client_hf()

    def _init_client_hf(self) -> None:
        self.client = InferenceClient(
            model=self.llm_model,
            token=os.environ.get("HF_TOKEN"),
        )

    def index_documents(self) -> None:
        """Calcule les embeddings par lots en local avec le modèle Hugging Face."""
        texts = [doc["text"] for doc in self.indexed_documents]
        
        for i in range(0, len(texts), self.batch_size):
            batch = texts[i:i + self.batch_size]
            if not batch:
                continue

            # Tokenisation et préparation des tenseurs
            inputs = self.embedding_tokenizer(batch, padding=True, truncation=True, return_tensors="pt")
            with torch.no_grad():
                outputs = self.embedding_model_local(**inputs)
            # Calcul du pooling moyen sur la dernière couche
            batch_embeddings_tensor = outputs.last_hidden_state.mean(dim=1)
            batch_embeddings = batch_embeddings_tensor.cpu().tolist()
            self.embeddings.extend(batch_embeddings)
            
            print(f"Batch {i//self.batch_size + 1} traité, {len(batch_embeddings)} embeddings obtenus")

    @staticmethod
    def cosine_similarity(query_vec: np.ndarray, matrix: np.ndarray) -> np.ndarray:
        """
        Calcule la similarité cosinus entre un vecteur de requête et chaque vecteur d'une matrice.
        """
        query_norm = np.linalg.norm(query_vec)
        query_normalized = query_vec / (query_norm + 1e-10)
        matrix_norm = np.linalg.norm(matrix, axis=1, keepdims=True)
        matrix_normalized = matrix / (matrix_norm + 1e-10)
        return np.dot(matrix_normalized, query_normalized)

    def search(self, query_embedding: List[float], top_k: int = 5) -> List[Dict]:
        """
        Recherche des documents sur la base de la similarité cosinus.

        Args:
            query_embedding: L'embedding de la requête.
            top_k: Nombre de résultats à renvoyer.

        Returns:
            Une liste de dictionnaires avec les clés "url", "text" et "score".
        """
        query_vec = np.array(query_embedding)
        emb_matrix = np.array(self.embeddings)
        scores = self.cosine_similarity(query_vec, emb_matrix)
        top_indices = np.argsort(scores)[::-1][:top_k]
        results = []
        for idx in top_indices:
            doc = self.indexed_documents[idx]
            results.append(
                {"url": doc["url"], "text": doc["text"], "score": float(scores[idx])}
            )
        return results

    def ask_llm(self, prompt: str) -> str:
        """
        Appelle le LLM avec l'invite construite et renvoie la réponse générée.
        """
        messages = [{"role": "user", "content": prompt}]
        response = self.client.chat.completions.create(
            model=self.llm_model, messages=messages
        )
        return response.choices[0].message.content

    def rag(self, question: str, top_k: int = 4) -> Dict[str, str]:
        """
        Effectue une génération augmentée par récupération (RAG) pour une question donnée.

        Args:
            question: La question posée.
            top_k: Nombre de documents de contexte à inclure.

        Returns:
            Un dictionnaire avec les clés "response", "prompt" et "urls".
        """
        # 1. Calculer l'embedding de la question en local.
        inputs = self.embedding_tokenizer(question, return_tensors="pt")
        with torch.no_grad():
            outputs = self.embedding_model_local(**inputs)
        question_embedding_tensor = outputs.last_hidden_state.mean(dim=1)[0]
        question_embedding = question_embedding_tensor.cpu().tolist()

        # 2. Récupérer les documents les plus similaires.
        results = self.search(query_embedding=question_embedding, top_k=top_k)
        context = "\n\n".join([f"URL: {res['url']}\n{res['text']}" for res in results])

        # 3. Construire l'invite.
        prompt = (
            "You are a highly capable, thoughtful, and precise assistant. Your goal is to deeply understand the user's intent, ask clarifying questions when needed, think step-by-step through complex problems, provide clear and accurate answers, and proactively anticipate helpful follow-up information. "
            "Based on the following context, answer the question precisely and concisely. "
            "If you do not know the answer, do not make it up.\n\n"
            f"Context:\n{context}\n\n"
            f"Question: {question}\n\n"
            "Answer:"
        )

        urls = [res['url'] for res in results]

        # 4. Appeler le LLM avec l'invite construite.
        llm_response = self.ask_llm(prompt)
        return {"response": llm_response, "prompt": prompt, "urls": urls}