"""RelevanceScorer - 检索结果相关性自动评估

对检索结果逐文档评估与查询的相关性，用于 CRAG 自纠正循环的评估阶段。
"""

from __future__ import annotations

import logging
import math
import re
from dataclasses import dataclass
from enum import Enum
from typing import Any

from agentkit.memory.base import MemoryItem

logger = logging.getLogger(__name__)


class RelevanceVerdict(str, Enum):
    """相关性判定结果"""

    CORRECT = "correct"
    AMBIGUOUS = "ambiguous"
    INCORRECT = "incorrect"


@dataclass
class RelevanceScore:
    """单个文档的相关性评分"""

    item: MemoryItem
    score: float  # 0.0 ~ 1.0
    verdict: RelevanceVerdict
    reason: str = ""


@dataclass
class RetrievalEvaluation:
    """一次检索的整体评估结果"""

    scores: list[RelevanceScore]
    overall_verdict: RelevanceVerdict
    avg_score: float
    relevant_count: int
    total_count: int


class RelevanceScorer:
    """检索结果相关性评估器

    基于查询-文档语义相似度和关键词重叠的轻量级评估器。
    不依赖 LLM 调用，适用于生产环境的低延迟评估。

    评分策略：
    1. 关键词重叠率（Jaccard 相似度）
    2. 查询词覆盖率（query term coverage）
    3. 原始检索分数加权
    4. 长度惩罚（过短或过长的文档降分）
    """

    def __init__(
        self,
        correct_threshold: float = 0.6,
        ambiguous_threshold: float = 0.35,
        keyword_weight: float = 0.3,
        coverage_weight: float = 0.3,
        retrieval_weight: float = 0.3,
        length_weight: float = 0.1,
        min_doc_length: int = 20,
        max_doc_length: int = 5000,
    ):
        self._correct_threshold = correct_threshold
        self._ambiguous_threshold = ambiguous_threshold
        self._keyword_weight = keyword_weight
        self._coverage_weight = coverage_weight
        self._retrieval_weight = retrieval_weight
        self._length_weight = length_weight
        self._min_doc_length = min_doc_length
        self._max_doc_length = max_doc_length

    def score_item(self, query: str, item: MemoryItem) -> RelevanceScore:
        """评估单个检索结果与查询的相关性"""
        doc_text = str(item.value)

        # 1. Keyword overlap (Jaccard similarity)
        query_terms = self._tokenize(query)
        doc_terms = self._tokenize(doc_text)
        keyword_score = self._jaccard_similarity(query_terms, doc_terms)

        # 2. Query term coverage
        coverage_score = self._query_coverage(query_terms, doc_terms)

        # 3. Original retrieval score
        retrieval_score = min(max(item.score, 0.0), 1.0)

        # 4. Length penalty
        length_score = self._length_score(len(doc_text))

        # Weighted combination
        final_score = (
            keyword_score * self._keyword_weight
            + coverage_score * self._coverage_weight
            + retrieval_score * self._retrieval_weight
            + length_score * self._length_weight
        )

        # Determine verdict
        verdict = self._determine_verdict(final_score)

        reason = (
            f"keyword={keyword_score:.2f}, coverage={coverage_score:.2f}, "
            f"retrieval={retrieval_score:.2f}, length={length_score:.2f}"
        )

        return RelevanceScore(
            item=item,
            score=final_score,
            verdict=verdict,
            reason=reason,
        )

    def evaluate(
        self, query: str, items: list[MemoryItem]
    ) -> RetrievalEvaluation:
        """评估一次检索的整体质量"""
        if not items:
            return RetrievalEvaluation(
                scores=[],
                overall_verdict=RelevanceVerdict.INCORRECT,
                avg_score=0.0,
                relevant_count=0,
                total_count=0,
            )

        scores = [self.score_item(query, item) for item in items]
        relevant_count = sum(
            1 for s in scores if s.verdict != RelevanceVerdict.INCORRECT
        )
        avg_score = sum(s.score for s in scores) / len(scores)

        # Overall verdict based on average score and relevant ratio
        relevant_ratio = relevant_count / len(scores)

        if avg_score >= self._correct_threshold and relevant_ratio >= 0.5:
            overall_verdict = RelevanceVerdict.CORRECT
        elif avg_score >= self._ambiguous_threshold or relevant_ratio >= 0.3:
            overall_verdict = RelevanceVerdict.AMBIGUOUS
        else:
            overall_verdict = RelevanceVerdict.INCORRECT

        return RetrievalEvaluation(
            scores=scores,
            overall_verdict=overall_verdict,
            avg_score=avg_score,
            relevant_count=relevant_count,
            total_count=len(scores),
        )

    def _determine_verdict(self, score: float) -> RelevanceVerdict:
        """根据分数判定相关性"""
        if score >= self._correct_threshold:
            return RelevanceVerdict.CORRECT
        elif score >= self._ambiguous_threshold:
            return RelevanceVerdict.AMBIGUOUS
        else:
            return RelevanceVerdict.INCORRECT

    @staticmethod
    def _tokenize(text: str) -> set[str]:
        """分词：中文按字符，英文按空格，统一小写"""
        tokens: set[str] = set()
        # Extract English words
        en_words = re.findall(r"[a-zA-Z]+", text.lower())
        tokens.update(en_words)
        # Extract Chinese characters (individual chars + bigrams)
        cn_chars = re.findall(r"[\u4e00-\u9fff]", text)
        tokens.update(cn_chars)
        # Add Chinese bigrams for better matching
        for i in range(len(cn_chars) - 1):
            tokens.add(cn_chars[i] + cn_chars[i + 1])
        return tokens

    @staticmethod
    def _jaccard_similarity(set_a: set[str], set_b: set[str]) -> float:
        """Jaccard 相似度"""
        if not set_a or not set_b:
            return 0.0
        intersection = len(set_a & set_b)
        union = len(set_a | set_b)
        if union == 0:
            return 0.0
        return intersection / union

    @staticmethod
    def _query_coverage(query_terms: set[str], doc_terms: set[str]) -> float:
        """查询词覆盖率：文档中出现的查询词比例"""
        if not query_terms:
            return 0.0
        covered = len(query_terms & doc_terms)
        return covered / len(query_terms)

    def _length_score(self, length: int) -> float:
        """长度评分：过短或过长的文档降分"""
        if length < self._min_doc_length:
            # Too short — likely insufficient context
            ratio = length / self._min_doc_length
            return ratio * 0.5
        elif length > self._max_doc_length:
            # Too long — may contain irrelevant information
            excess = (length - self._max_doc_length) / self._max_doc_length
            return max(0.3, 1.0 - excess * 0.5)
        else:
            # Good length range
            return 1.0