# vecta_sdk/vecta/core/evaluator.py
"""Evaluation logic for retrieval performance metrics and generation quality assessment."""

from __future__ import annotations

import concurrent.futures
import json
import logging
import time
from typing import Any, Callable, Dict, List, Optional, Set, Tuple

from openai import OpenAI

from vecta.core.schemas import (
    BenchmarkEntry,
    BenchmarkResults,
    EvaluationMetrics,
    EvaluationResultRow,
    GenerationMetrics,
    GenerationOnlyResults,
    RetrievalAndGenerationResults,
    UsageInfo,
)

logger = logging.getLogger("vecta.evaluator")


class Evaluator:
    """Evaluates AI performance against benchmark (optionally using KB metadata)."""

    def __init__(
        self,
        benchmark_entries: List[BenchmarkEntry],
        meta_lookup: Optional[
            Dict[str, Tuple[Optional[List[int]], Optional[str]]]
        ] = None,
        openai_api_key: Optional[str] = None,
        openai_base_url: str = "https://openrouter.ai/api/v1",
        model: str = "google/gemini-2.5-flash-lite",
        usage_service: Any = None,
    ):
        self.benchmark_entries = benchmark_entries
        self.meta_lookup = meta_lookup or {}
        self.openai_api_key = openai_api_key
        self.openai_base_url = openai_base_url
        self.model = model
        self.usage_service = usage_service
        self.total_usage = UsageInfo()
        self._client: Optional[Any] = None
        self._llm_failures: int = 0

        if self.openai_api_key and OpenAI is not None:
            try:
                self._client = OpenAI(
                    base_url=self.openai_base_url,
                    api_key=self.openai_api_key,
                    max_retries=3,
                    timeout=60.0,
                )
            except Exception:
                logger.warning("Failed to initialise OpenAI client", exc_info=True)
                self._client = None

    # ------------------------------------------------------------------
    # Public evaluation methods
    # ------------------------------------------------------------------

    def evaluate_retrieval(
        self,
        retrieval_function: Callable[[str], List[str]],
        evaluation_name: str,
    ) -> BenchmarkResults:
        """Evaluate retrieval-only function performance."""
        if not self.meta_lookup:
            raise ValueError(
                "meta_lookup is required for retrieval evaluation and must contain "
                "{chunk_id: (page_nums, source_path)}"
            )

        start_time = time.time()

        all_results: List[Tuple[BenchmarkEntry, List[str], str]] = []
        for entry in self.benchmark_entries:
            retrieved_chunk_ids = retrieval_function(entry.question) or []
            normalized = [cid for cid in retrieved_chunk_ids if cid]
            all_results.append((entry, normalized, ""))

        chunk_metrics = self._compute_retrieval_metrics(
            all_results,
            lambda e: set(e.chunk_ids) if e.chunk_ids else set(),
            lambda ids: set(ids),
        )

        page_metrics = self._compute_page_level_metrics(all_results)

        doc_metrics = self._compute_retrieval_metrics(
            all_results,
            lambda e: set(e.source_paths) if e.source_paths else set(),
            lambda ids: self._map_chunks_to_docs(ids),
        )

        duration_seconds = int(time.time() - start_time)
        detailed_results = self._build_detailed_results(all_results)

        return BenchmarkResults(
            chunk_level=chunk_metrics,
            page_level=page_metrics,
            document_level=doc_metrics,
            detailed_results=detailed_results,
            total_questions=len(self.benchmark_entries),
            duration_seconds=duration_seconds,
            evaluation_name=evaluation_name,
            usage_info=self.total_usage if self.total_usage.total_tokens > 0 else None,
        )

    def evaluate_retrieval_and_generation(
        self,
        retrieval_generation_function: Callable[[str], Tuple[List[str], str]],
        evaluation_name: str,
    ) -> RetrievalAndGenerationResults:
        """Evaluate a function that returns both retrieved chunk_ids and generated text."""
        if not self.meta_lookup:
            raise ValueError(
                "meta_lookup is required for retrieval+generation evaluation and must "
                "contain {chunk_id: (page_nums, source_path)}"
            )

        start_time = time.time()

        all_results: List[Tuple[BenchmarkEntry, List[str], str]] = []
        for entry in self.benchmark_entries:
            chunk_ids, gen_text = retrieval_generation_function(entry.question)
            normalized = [cid for cid in chunk_ids if cid]
            all_results.append((entry, normalized, gen_text))

        duration_seconds = int(time.time() - start_time)

        chunk_metrics = self._compute_retrieval_metrics(
            all_results,
            lambda e: set(e.chunk_ids) if e.chunk_ids else set(),
            lambda ids: set(ids),
        )
        page_metrics = self._compute_page_level_metrics(all_results)
        doc_metrics = self._compute_retrieval_metrics(
            all_results,
            lambda e: set(e.source_paths) if e.source_paths else set(),
            lambda ids: self._map_chunks_to_docs(ids),
        )

        questions = [entry.question for entry, _, _ in all_results]
        expected_answers = [entry.answer for entry, _, _ in all_results]
        generated_answers = [gen_text for _, _, gen_text in all_results]

        accuracy, groundedness, per_question_scores = self._evaluate_generation_batch(
            questions, expected_answers, generated_answers
        )

        gen_metrics = GenerationMetrics(accuracy=accuracy, groundedness=groundedness)
        detailed_results = self._build_detailed_results(
            all_results, generation_scores=per_question_scores
        )

        return RetrievalAndGenerationResults(
            chunk_level=chunk_metrics,
            page_level=page_metrics,
            document_level=doc_metrics,
            generation_metrics=gen_metrics,
            detailed_results=detailed_results,
            total_questions=len(self.benchmark_entries),
            duration_seconds=duration_seconds,
            evaluation_name=evaluation_name,
            usage_info=self.total_usage if self.total_usage.total_tokens > 0 else None,
        )

    def evaluate_generation_only(
        self,
        generation_function: Callable[[str], str],
        evaluation_name: str,
    ) -> GenerationOnlyResults:
        """Evaluate generation-only quality versus expected answers."""
        questions: List[str] = []
        expected_answers: List[str] = []
        generated_answers: List[str] = []

        start_time = time.time()

        for entry in self.benchmark_entries:
            questions.append(entry.question)
            expected_answers.append(entry.answer)
            generated_answers.append(generation_function(entry.question))

        duration_seconds = int(time.time() - start_time)

        accuracy, groundedness, per_question_scores = self._evaluate_generation_batch(
            questions, expected_answers, generated_answers
        )

        gen_metrics = GenerationMetrics(accuracy=accuracy, groundedness=groundedness)

        all_results: List[Tuple[BenchmarkEntry, List[str], str]] = [
            (entry, [], generated_answers[idx])
            for idx, entry in enumerate(self.benchmark_entries)
        ]

        detailed_results = self._build_detailed_results(
            all_results,
            generation_scores=per_question_scores,
            include_retrieval_metrics=False,
        )

        if self._llm_failures > 0:
            logger.warning(
                "Generation scoring completed with %d/%d LLM failures "
                "(those questions scored 0.0)",
                self._llm_failures,
                len(questions),
            )

        return GenerationOnlyResults(
            generation_metrics=gen_metrics,
            detailed_results=detailed_results,
            total_questions=len(self.benchmark_entries),
            duration_seconds=duration_seconds,
            evaluation_name=evaluation_name,
            usage_info=self.total_usage if self.total_usage.total_tokens > 0 else None,
        )

    # ------------------------------------------------------------------
    # Retrieval metrics
    # ------------------------------------------------------------------

    def _compute_retrieval_metrics(
        self,
        all_results: List[Tuple[BenchmarkEntry, List[str], str]],
        get_ground_truth: Callable[[BenchmarkEntry], Set[Any]],
        get_retrieved: Callable[[List[str]], Set[Any]],
    ) -> EvaluationMetrics:
        total_precision = 0.0
        total_recall = 0.0
        total_f1 = 0.0

        for entry, chunk_ids, _ in all_results:
            ground_truth = get_ground_truth(entry)
            retrieved = get_retrieved(chunk_ids)
            precision, recall, f1 = self._calculate_metrics(ground_truth, retrieved)
            total_precision += precision
            total_recall += recall
            total_f1 += f1

        n = len(all_results) or 1
        return EvaluationMetrics(
            precision=total_precision / n,
            recall=total_recall / n,
            f1_score=total_f1 / n,
        )

    def _compute_page_level_metrics(
        self, all_results: List[Tuple[BenchmarkEntry, List[str], str]]
    ) -> Optional[EvaluationMetrics]:
        has_page_info = any(
            entry.page_nums and len(entry.page_nums) > 0
            for entry, _, _ in all_results
        )
        if not has_page_info:
            return None

        total_precision = 0.0
        total_recall = 0.0
        total_f1 = 0.0
        valid_entries = 0

        for entry, chunk_ids, _ in all_results:
            if not entry.page_nums or not entry.source_paths:
                continue
            valid_entries += 1

            ground_truth_pages = {
                (source_path, page_num)
                for source_path in entry.source_paths
                for page_num in entry.page_nums
            }
            retrieved_pages = self._map_chunks_to_pages_with_overlap(chunk_ids, entry)

            precision, recall, f1 = self._calculate_metrics(
                ground_truth_pages, retrieved_pages
            )
            total_precision += precision
            total_recall += recall
            total_f1 += f1

        if valid_entries == 0:
            return None

        return EvaluationMetrics(
            precision=total_precision / valid_entries,
            recall=total_recall / valid_entries,
            f1_score=total_f1 / valid_entries,
        )

    # ------------------------------------------------------------------
    # Detailed results builder
    # ------------------------------------------------------------------

    def _build_detailed_results(
        self,
        all_results: List[Tuple[BenchmarkEntry, List[str], str]],
        generation_scores: Optional[List[Tuple[float, float]]] = None,
        include_retrieval_metrics: bool = True,
    ) -> List[EvaluationResultRow]:
        detailed_results: List[EvaluationResultRow] = []

        for index, (entry, chunk_ids, generated_answer) in enumerate(all_results):
            ordered_chunks = list(dict.fromkeys([cid for cid in chunk_ids if cid]))
            expected_chunks = list(dict.fromkeys(entry.chunk_ids or []))

            expected_chunk_set = set(expected_chunks)
            retrieved_chunk_set = set(ordered_chunks)

            chunk_precision: Optional[float] = None
            chunk_recall: Optional[float] = None
            chunk_f1: Optional[float] = None
            document_precision: Optional[float] = None
            document_recall: Optional[float] = None
            document_f1: Optional[float] = None
            page_precision: Optional[float] = None
            page_recall: Optional[float] = None
            page_f1: Optional[float] = None
            matched_chunks: List[str] = []
            missed_chunks: List[str] = []
            expected_source_paths = list(entry.source_paths or [])
            expected_source_set = set(expected_source_paths)
            retrieved_source_paths: List[str] = []
            matched_source_paths: List[str] = []

            if include_retrieval_metrics:
                if expected_chunk_set or retrieved_chunk_set:
                    c_p, c_r, c_f = self._calculate_metrics(
                        expected_chunk_set, retrieved_chunk_set
                    )
                    chunk_precision, chunk_recall, chunk_f1 = c_p, c_r, c_f

                matched_chunks = [
                    cid for cid in ordered_chunks if cid in expected_chunk_set
                ]
                missed_chunks = [
                    cid for cid in expected_chunks if cid not in retrieved_chunk_set
                ]

                for cid in ordered_chunks:
                    meta = self.meta_lookup.get(cid)
                    if not meta:
                        continue
                    _, source_path = meta
                    if source_path and source_path not in retrieved_source_paths:
                        retrieved_source_paths.append(source_path)

                matched_source_paths = [
                    doc for doc in retrieved_source_paths if doc in expected_source_set
                ]

                retrieved_docs = self._map_chunks_to_docs(ordered_chunks)
                if expected_source_set or retrieved_docs:
                    d_p, d_r, d_f = self._calculate_metrics(
                        expected_source_set, retrieved_docs
                    )
                    document_precision, document_recall, document_f1 = d_p, d_r, d_f

                if entry.page_nums and entry.source_paths:
                    gt_pages = {
                        (sp, p) for sp in entry.source_paths for p in entry.page_nums
                    }
                    ret_pages = self._map_chunks_to_pages_with_overlap(
                        ordered_chunks, entry
                    )
                    if gt_pages or ret_pages:
                        p_p, p_r, p_f = self._calculate_metrics(gt_pages, ret_pages)
                        page_precision, page_recall, page_f1 = p_p, p_r, p_f

            accuracy: Optional[float] = None
            groundedness: Optional[float] = None
            if generation_scores and index < len(generation_scores):
                accuracy, groundedness = generation_scores[index]

            detailed_results.append(
                EvaluationResultRow(
                    benchmark_entry_id=entry.id,
                    question=entry.question,
                    expected_answer=entry.answer,
                    generated_answer=generated_answer or None,
                    retrieved_chunk_ids=ordered_chunks,
                    expected_chunk_ids=expected_chunks,
                    matched_chunk_ids=matched_chunks,
                    missed_chunk_ids=missed_chunks,
                    expected_source_paths=expected_source_paths,
                    retrieved_source_paths=retrieved_source_paths,
                    matched_source_paths=matched_source_paths,
                    chunk_precision=chunk_precision,
                    chunk_recall=chunk_recall,
                    chunk_f1=chunk_f1,
                    page_precision=page_precision,
                    page_recall=page_recall,
                    page_f1=page_f1,
                    document_precision=document_precision,
                    document_recall=document_recall,
                    document_f1=document_f1,
                    accuracy=accuracy,
                    groundedness=groundedness,
                )
            )

        if not include_retrieval_metrics:
            for detail in detailed_results:
                detail.matched_chunk_ids = []
                detail.missed_chunk_ids = []
                detail.chunk_precision = None
                detail.chunk_recall = None
                detail.chunk_f1 = None
                detail.document_precision = None
                detail.document_recall = None
                detail.document_f1 = None
                detail.page_precision = None
                detail.page_recall = None
                detail.page_f1 = None
                detail.retrieved_source_paths = []
                detail.matched_source_paths = []

        return detailed_results

    # ------------------------------------------------------------------
    # Generation scoring
    # ------------------------------------------------------------------

    def _evaluate_generation_batch(
        self,
        questions: List[str],
        expected_answers: List[str],
        generated_answers: List[str],
    ) -> Tuple[float, float, List[Tuple[float, float]]]:
        if not self._client:
            logger.info(
                "No OpenAI client configured — generation metrics will be 0.0"
            )
            return 0.0, 0.0, [(0.0, 0.0) for _ in questions]

        self._llm_failures = 0
        all_accuracy: List[float] = []
        all_groundedness: List[float] = []
        per_question_scores: List[Tuple[float, float]] = [
            (0.0, 0.0) for _ in range(len(questions))
        ]
        batch_size = 100

        for i in range(0, len(questions), batch_size):
            batch_q = questions[i : i + batch_size]
            batch_e = expected_answers[i : i + batch_size]
            batch_g = generated_answers[i : i + batch_size]

            with concurrent.futures.ThreadPoolExecutor(max_workers=100) as executor:
                futures: List[
                    Tuple[int, concurrent.futures.Future[Tuple[float, float]]]
                ] = []
                for offset, (q, e, g) in enumerate(
                    zip(batch_q, batch_e, batch_g)
                ):
                    future = executor.submit(
                        self._evaluate_generation_single, q, e, g
                    )
                    futures.append((i + offset, future))

                for absolute_index, future in futures:
                    try:
                        acc, fac = future.result(timeout=90)
                    except Exception:
                        logger.warning(
                            "Generation scoring timed out for question #%d",
                            absolute_index,
                        )
                        acc, fac = 0.0, 0.0
                        self._llm_failures += 1

                    per_question_scores[absolute_index] = (acc, fac)
                    all_accuracy.append(acc)
                    all_groundedness.append(fac)

        n = len(all_accuracy) or 1
        return (
            sum(all_accuracy) / n,
            sum(all_groundedness) / n,
            per_question_scores,
        )

    @staticmethod
    def _truncate(text: str, max_len: int = 10_000) -> str:
        """Truncate text for LLM prompts, adding ellipsis only when needed."""
        if len(text) <= max_len:
            return text
        return text[:max_len] + "…"

    def _evaluate_generation_single(
        self, question: str, expected_answer: str, generated_answer: str
    ) -> Tuple[float, float]:
        if not self._client:
            return 0.0, 0.0

        system_prompt = (
            "You are an expert evaluator assessing AI-generated answers. "
            "Score ACCURACY (match to expected answer) and GROUNDEDNESS "
            "(consistency with retrieved data, no invented facts). "
            'Return JSON: {"accuracy_score": 0.0-1.0, "groundedness_score": 0.0-1.0}.'
        )

        user_prompt = (
            f"Question: {self._truncate(question)}\n\n"
            f"Expected: {self._truncate(expected_answer)}\n\n"
            f"Generated: {self._truncate(generated_answer)}\n\n"
            "Evaluate accuracy (vs expected) and groundedness "
            "(self-consistent, no unsupported claims)."
        )

        if self.usage_service:
            estimated_tokens = self.usage_service.estimate_tokens(
                system_prompt + user_prompt
            )
            self.usage_service.check_usage_limit(estimated_tokens)

        try:
            response = self._client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                response_format={
                    "type": "json_schema",
                    "json_schema": {
                        "name": "generation_evaluation",
                        "strict": True,
                        "schema": {
                            "type": "object",
                            "properties": {
                                "accuracy_score": {
                                    "type": "number",
                                    "minimum": 0,
                                    "maximum": 1,
                                },
                                "groundedness_score": {
                                    "type": "number",
                                    "minimum": 0,
                                    "maximum": 1,
                                },
                            },
                            "required": ["accuracy_score", "groundedness_score"],
                            "additionalProperties": False,
                        },
                    },
                },
                timeout=60,
                max_tokens=1000,
            )

            if response.usage and self.usage_service:
                usage_info = UsageInfo(
                    prompt_tokens=response.usage.prompt_tokens,
                    completion_tokens=response.usage.completion_tokens,
                    total_tokens=response.usage.total_tokens,
                )
                self.usage_service.log_usage(usage_info)

            content = response.choices[0].message.content
            if content is None:
                logger.warning("LLM returned empty content for generation scoring")
                self._llm_failures += 1
                return 0.0, 0.0

            data = json.loads(content)
            return (
                float(data.get("accuracy_score", 0.0)),
                float(data.get("groundedness_score", 0.0)),
            )

        except Exception:
            logger.warning(
                "Generation scoring failed (defaulting to 0.0)", exc_info=True
            )
            self._llm_failures += 1
            return 0.0, 0.0

    # ------------------------------------------------------------------
    # Retrieval helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _calculate_metrics(
        ground_truth: Set[Any], retrieved: Set[Any]
    ) -> Tuple[float, float, float]:
        if not retrieved or not ground_truth:
            return 0.0, 0.0, 0.0
        inter = ground_truth & retrieved
        precision = len(inter) / len(retrieved)
        recall = len(inter) / len(ground_truth)
        f1 = (
            0.0
            if (precision + recall == 0)
            else 2 * precision * recall / (precision + recall)
        )
        return precision, recall, f1

    def _map_chunks_to_pages_with_overlap(
        self, chunk_ids: List[str], entry: BenchmarkEntry
    ) -> Set[Tuple[str, int]]:
        retrieved_pages: Set[Tuple[str, int]] = set()

        expected_pages: Set[Tuple[str, int]] = set()
        if entry.source_paths and entry.page_nums:
            for sp in entry.source_paths:
                for p in entry.page_nums:
                    expected_pages.add((sp, p))

        for cid in chunk_ids:
            if cid not in self.meta_lookup:
                continue
            page_nums, source_path = self.meta_lookup[cid]
            if page_nums is None or source_path is None:
                continue

            chunk_pages = {(source_path, p) for p in page_nums}
            expected_for_doc = {(d, p) for d, p in expected_pages if d == source_path}
            overlap = chunk_pages & expected_for_doc
            if overlap:
                retrieved_pages.update(overlap)

        return retrieved_pages

    def _map_chunks_to_docs(self, chunk_ids: List[str]) -> Set[str]:
        out: Set[str] = set()
        for cid in chunk_ids:
            if cid not in self.meta_lookup:
                continue
            _, source_path = self.meta_lookup[cid]
            if source_path is not None:
                out.add(source_path)
        return out