from typing_extensions import Annotated, TypedDict
from langchain.chat_models import ChatOpenAI
from config import client, llm
from config import dataset_id

# Define the output schema for grading
class CorrectnessGrade(TypedDict):
    explanation: Annotated[str, ..., "Expliquez votre raisonnement pour la note."]
    score: Annotated[float, ..., "Un score de 0 à 10 basé sur la justesse de la réponse"]

# Grading prompt with expert evaluation criteria
correctness_instructions = """Vous êtes un expert en intelligence artificielle chargé d'évaluer la pertinence d'une réponse générée par un système RAG.

**Instructions** :
- Vous recevrez une QUESTION, une RÉPONSE DE RÉFÉRENCE (vérité terrain) et une RÉPONSE GÉNÉRÉE.
- Évaluez uniquement la précision factuelle de la réponse générée par rapport à la référence.
- Une réponse peut contenir plus d'informations que la référence, tant qu'elles sont exactes.
- Si la réponse générée contient des erreurs ou des informations contradictoires, pénalisez-la.

**Notation (score sur 10)** :
- **0 (Très insuffisant)** : Réponse hors sujet ou contenant des erreurs majeures.
- **2.5 (Insuffisant)** : Réponse partiellement correcte mais incomplète ou floue.
- **5 (Correct)** : Réponse pertinente mais manquant de clarté ou de détails.
- **7.5 (Bon)** : Réponse pertinente, claire et complète avec une légère amélioration possible.
- **10 (Très bon)** : Réponse exacte, détaillée et bien structurée.

**Expliquez votre notation étape par étape.**"""

def get_reference_answer(question: str) -> str:
    """Retrieve the reference answer from a LangSmith dataset."""
    
    # Récupérer les exemples stockés dans le dataset
    dataset = client.list_examples(dataset_id)

    # Chercher la correspondance exacte
    for example in dataset:
        if example.inputs.get("question") == question:
            return example.outputs.get("answer")  # Retourne la réponse de référence
    
    return None  # Si aucune correspondance n'est trouvée

def correctness(inputs: dict, outputs: dict) -> float:
    """Evaluates the correctness of a RAG-generated answer and returns a score (0-10)."""
    question = inputs["question"]
    reference_answer = get_reference_answer(question)
    
    if not reference_answer:
        raise ValueError(f"No reference answer found for question: {question}")

    answers = f"""\nQUESTION: {question}
    RÉPONSE DE RÉFÉRENCE: {reference_answer}
    RÉPONSE GÉNÉRÉE: {outputs['answer']}"""

    # Invoke the LLM for grading
    grade = llm.invoke([
        {"role": "system", "content": correctness_instructions}, 
        {"role": "user", "content": answers}
    ])
    
    return grade["score"]