"""
RAG evaluation module.

Evaluates retrieval quality, context relevance, and answer faithfulness.
"""

import json
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

from ..infra.logging import get_logger
from ..infra.tracing import trace_function

logger = get_logger(__name__)


@dataclass
class RetrievalMetrics:
    """Metrics for retrieval evaluation."""
    precision: float = 0.0
    recall: float = 0.0
    f1: float = 0.0
    mrr: float = 0.0  # Mean Reciprocal Rank
    ndcg: float = 0.0  # Normalized Discounted Cumulative Gain
    hit_rate: float = 0.0
    
    def to_dict(self) -> Dict[str, float]:
        return {
            "precision": self.precision,
            "recall": self.recall,
            "f1": self.f1,
            "mrr": self.mrr,
            "ndcg": self.ndcg,
            "hit_rate": self.hit_rate,
        }


@dataclass
class RAGMetrics:
    """Complete RAG evaluation metrics."""
    retrieval: RetrievalMetrics = field(default_factory=RetrievalMetrics)
    context_relevance: float = 0.0
    context_coverage: float = 0.0
    answer_faithfulness: float = 0.0
    answer_relevance: float = 0.0
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "retrieval": self.retrieval.to_dict(),
            "context_relevance": self.context_relevance,
            "context_coverage": self.context_coverage,
            "answer_faithfulness": self.answer_faithfulness,
            "answer_relevance": self.answer_relevance,
        }


class RAGEvaluator:
    """
    Evaluator for RAG pipeline quality.
    
    Supports both automatic metrics and LLM-based evaluation.
    """
    
    def __init__(
        self,
        llm_client=None,
        use_llm_eval: bool = True
    ):
        """
        Initialize evaluator.
        
        Args:
            llm_client: LLM client for LLM-based evaluation
            use_llm_eval: Whether to use LLM-based evaluation
        """
        self.llm_client = llm_client
        self.use_llm_eval = use_llm_eval and llm_client is not None
    
    @trace_function(name="evaluate_retrieval")
    def evaluate_retrieval(
        self,
        retrieved_docs: List[Dict[str, Any]],
        relevant_docs: List[str],
        k: int = 10
    ) -> RetrievalMetrics:
        """
        Evaluate retrieval quality.
        
        Args:
            retrieved_docs: Retrieved documents with content
            relevant_docs: List of relevant document IDs or content
            k: Number of documents to evaluate
            
        Returns:
            RetrievalMetrics with precision, recall, etc.
        """
        if not retrieved_docs or not relevant_docs:
            return RetrievalMetrics()
        
        # Get retrieved IDs/content
        retrieved = []
        for doc in retrieved_docs[:k]:
            doc_id = doc.get("id") or doc.get("content", "")[:100]
            retrieved.append(doc_id)
        
        # Calculate precision
        relevant_set = set(relevant_docs)
        retrieved_relevant = sum(1 for d in retrieved if d in relevant_set)
        
        precision = retrieved_relevant / len(retrieved) if retrieved else 0.0
        recall = retrieved_relevant / len(relevant_docs) if relevant_docs else 0.0
        f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0
        
        # Calculate MRR
        mrr = 0.0
        for i, doc in enumerate(retrieved):
            if doc in relevant_set:
                mrr = 1.0 / (i + 1)
                break
        
        # Calculate hit rate
        hit_rate = 1.0 if retrieved_relevant > 0 else 0.0
        
        # Calculate NDCG
        ndcg = self._calculate_ndcg(retrieved, relevant_docs)
        
        return RetrievalMetrics(
            precision=precision,
            recall=recall,
            f1=f1,
            mrr=mrr,
            ndcg=ndcg,
            hit_rate=hit_rate,
        )
    
    def _calculate_ndcg(
        self,
        retrieved: List[str],
        relevant: List[str],
        k: Optional[int] = None
    ) -> float:
        """Calculate NDCG@k."""
        import math
        
        if not retrieved or not relevant:
            return 0.0
        
        k = k or len(retrieved)
        relevant_set = set(relevant)
        
        # DCG
        dcg = 0.0
        for i, doc in enumerate(retrieved[:k]):
            if doc in relevant_set:
                dcg += 1.0 / math.log2(i + 2)  # +2 because i is 0-indexed
        
        # Ideal DCG
        idcg = sum(1.0 / math.log2(i + 2) for i in range(min(len(relevant), k)))
        
        return dcg / idcg if idcg > 0 else 0.0
    
    @trace_function(name="evaluate_context_relevance")
    def evaluate_context_relevance(
        self,
        query: str,
        contexts: List[str]
    ) -> float:
        """
        Evaluate context relevance to query.
        
        Uses LLM to assess relevance on 0-1 scale.
        """
        if not contexts:
            return 0.0
        
        if not self.use_llm_eval:
            # Fallback to simple overlap metric
            return self._simple_relevance(query, contexts)
        
        # LLM-based evaluation
        prompt = f"""Evaluate the relevance of the following contexts to the query.
For each context, rate its relevance from 0 (not relevant) to 1 (highly relevant).

Query: {query}

Contexts:
{self._format_contexts(contexts)}

Provide your evaluation as a JSON object with format:
{{"scores": [0.8, 0.5, ...], "average": 0.65}}

Only return the JSON, no other text."""

        try:
            response = self.llm_client.generate(prompt)
            result = json.loads(response)
            return result.get("average", 0.0)
        except Exception as e:
            logger.warning(f"LLM evaluation failed: {e}")
            return self._simple_relevance(query, contexts)
    
    def _simple_relevance(self, query: str, contexts: List[str]) -> float:
        """Simple word overlap relevance metric."""
        query_words = set(query.lower().split())
        
        scores = []
        for context in contexts:
            context_words = set(context.lower().split())
            overlap = len(query_words & context_words)
            score = overlap / len(query_words) if query_words else 0.0
            scores.append(min(1.0, score))
        
        return sum(scores) / len(scores) if scores else 0.0
    
    @trace_function(name="evaluate_faithfulness")
    def evaluate_faithfulness(
        self,
        answer: str,
        contexts: List[str]
    ) -> float:
        """
        Evaluate answer faithfulness to source contexts.
        
        Checks if claims in answer are supported by contexts.
        """
        if not answer or not contexts:
            return 0.0
        
        if not self.use_llm_eval:
            return self._simple_faithfulness(answer, contexts)
        
        # LLM-based evaluation
        context_text = "\n\n".join(contexts[:5])  # Limit context
        
        prompt = f"""Evaluate whether the answer is faithful to the source contexts.
Check if all claims in the answer are supported by the contexts.

Contexts:
{context_text}

Answer:
{answer}

Rate faithfulness from 0 (contains unsupported claims) to 1 (fully supported).

Provide your evaluation as JSON:
{{"score": 0.85, "unsupported_claims": ["claim1", "claim2"], "reasoning": "..."}}

Only return the JSON."""

        try:
            response = self.llm_client.generate(prompt)
            result = json.loads(response)
            return result.get("score", 0.0)
        except Exception as e:
            logger.warning(f"Faithfulness evaluation failed: {e}")
            return self._simple_faithfulness(answer, contexts)
    
    def _simple_faithfulness(self, answer: str, contexts: List[str]) -> float:
        """Simple n-gram overlap faithfulness metric."""
        # Check if answer sentences can be traced to contexts
        combined_context = " ".join(contexts).lower()
        answer_sentences = answer.split(".")
        
        supported = 0
        for sent in answer_sentences:
            sent = sent.strip().lower()
            if len(sent) < 10:  # Skip very short sentences
                continue
            
            # Check for significant overlap
            words = sent.split()
            if len(words) < 3:
                continue
            
            # Look for trigram matches
            trigrams = [" ".join(words[i:i+3]) for i in range(len(words)-2)]
            matches = sum(1 for t in trigrams if t in combined_context)
            
            if matches >= len(trigrams) * 0.3:  # 30% trigram match
                supported += 1
        
        total = len([s for s in answer_sentences if len(s.strip()) >= 10])
        return supported / total if total > 0 else 0.0
    
    @trace_function(name="evaluate_answer_relevance")
    def evaluate_answer_relevance(
        self,
        query: str,
        answer: str
    ) -> float:
        """
        Evaluate answer relevance to query.
        """
        if not query or not answer:
            return 0.0
        
        if not self.use_llm_eval:
            return self._simple_relevance(query, [answer])
        
        prompt = f"""Evaluate how well the answer addresses the query.

Query: {query}

Answer:
{answer}

Rate relevance from 0 (does not address query) to 1 (fully addresses query).

Provide your evaluation as JSON:
{{"score": 0.9, "missing_aspects": ["aspect1"], "reasoning": "..."}}

Only return the JSON."""

        try:
            response = self.llm_client.generate(prompt)
            result = json.loads(response)
            return result.get("score", 0.0)
        except Exception as e:
            logger.warning(f"Answer relevance evaluation failed: {e}")
            return self._simple_relevance(query, [answer])
    
    def _format_contexts(self, contexts: List[str]) -> str:
        """Format contexts for prompt."""
        formatted = []
        for i, ctx in enumerate(contexts[:10]):  # Limit
            formatted.append(f"Context {i+1}:\n{ctx[:500]}...")
        return "\n\n".join(formatted)
    
    @trace_function(name="evaluate_rag")
    def evaluate(
        self,
        query: str,
        retrieved_docs: List[Dict[str, Any]],
        answer: str,
        relevant_docs: Optional[List[str]] = None
    ) -> RAGMetrics:
        """
        Run complete RAG evaluation.
        
        Args:
            query: Original query
            retrieved_docs: Retrieved documents
            answer: Generated answer
            relevant_docs: Ground truth relevant docs (optional)
            
        Returns:
            Complete RAGMetrics
        """
        contexts = [doc.get("content", "") for doc in retrieved_docs]
        
        metrics = RAGMetrics()
        
        # Retrieval metrics (if ground truth available)
        if relevant_docs:
            metrics.retrieval = self.evaluate_retrieval(retrieved_docs, relevant_docs)
        
        # Context relevance
        metrics.context_relevance = self.evaluate_context_relevance(query, contexts)
        
        # Context coverage (simple metric)
        if contexts:
            avg_len = sum(len(c) for c in contexts) / len(contexts)
            metrics.context_coverage = min(1.0, avg_len / 500)  # Normalize
        
        # Answer faithfulness
        metrics.answer_faithfulness = self.evaluate_faithfulness(answer, contexts)
        
        # Answer relevance
        metrics.answer_relevance = self.evaluate_answer_relevance(query, answer)
        
        return metrics


def evaluate_with_ragas(
    questions: List[str],
    contexts: List[List[str]],
    answers: List[str],
    ground_truths: Optional[List[str]] = None
) -> Dict[str, float]:
    """
    Evaluate using RAGAS framework (if available).
    
    Args:
        questions: List of questions
        contexts: List of context lists
        answers: Generated answers
        ground_truths: Ground truth answers (optional)
        
    Returns:
        RAGAS metrics dictionary
    """
    try:
        from ragas import evaluate
        from ragas.metrics import (
            faithfulness,
            answer_relevancy,
            context_relevancy,
            context_recall,
        )
        from datasets import Dataset
        
        # Prepare dataset
        data = {
            "question": questions,
            "contexts": contexts,
            "answer": answers,
        }
        
        if ground_truths:
            data["ground_truth"] = ground_truths
        
        dataset = Dataset.from_dict(data)
        
        # Select metrics
        metrics = [faithfulness, answer_relevancy, context_relevancy]
        if ground_truths:
            metrics.append(context_recall)
        
        # Evaluate
        results = evaluate(dataset, metrics=metrics)
        
        return dict(results)
        
    except ImportError:
        logger.warning("RAGAS not installed. Install with: pip install ragas")
        return {}
    except Exception as e:
        logger.error(f"RAGAS evaluation failed: {e}")
        return {}
