# -*- coding: utf-8 -*-
"""
Event Identification Grader for Finance Domain

Evaluates the quality of event identification in financial responses by comparing
two responses based on accuracy, time precision, and key element identification.
"""

import textwrap
from typing import Optional

from loguru import logger

from openjudge.evaluation_strategy import BaseEvaluationStrategy
from openjudge.graders.base_grader import GraderError, GraderMode, GraderRank
from openjudge.graders.llm_grader import LLMGrader
from openjudge.models.base_chat_model import BaseChatModel
from openjudge.models.schema.oai.message import ChatMessage
from openjudge.models.schema.prompt_template import LanguageEnum, PromptTemplate

# English Prompt
EVENT_IDENTIFICATION_PROMPT_EN = textwrap.dedent(
    """
You are an expert AI Assistant Evaluator. Your task is to evaluate the **Event Identification Accuracy** of the AI's response based strictly on the provided evaluation criteria.

<Rubrics>
Please evaluate the response based on **Event Identification**:
Criterion 1. Accurate Identification: Accurately retrieve and identify the event mentioned in the question

Criterion 2. Time Range Precision: Strictly match the time range specified in the question, avoid vague expressions;

Criterion 3. Core Elements Accuracy: Whether the key information of the event has been accurately identified and described. For example, when answering questions about "US CPI data", must accurately cite the latest CPI value, month-over-month and year-over-year growth rates, and core CPI data
</Rubrics>

<Steps>
1. First, determine whether the event in the question has been identified in the responses. A response that identifies the event is better;
2. For each evaluation criterion, select the better response and provide your reasoning;
3. Comprehensively evaluate the results under each criterion, then select the overall best answer and provide reasoning;
</Steps>

<Constraints>
- **Objectivity**: Do not be influenced by the politeness or length of the response.
- **Focus**: Only evaluate based on the provided evaluation criteria.
- A response that identifies the event is better than one that does not.
</Constraints>

<Scale>
- If Answer 1 is better: rank = [1, 2]
- If Answer 2 is better: rank = [2, 1]
</Scale>

<Query>
{query}
</Query>

<Response 1>
{answer_1}
</Response 1>

<Response 2>
{answer_2}
</Response 2>

<Output Schema>
Please output your assessment in the following JSON format strictly:
{{
    "reason": "<detailed explanation of your evaluation reasoning, including performance comparison under each evaluation criterion>",
    "rank": <[1, 2] or [2, 1]>
}}
</Output Schema>

JSON:
"""
).strip()

# Chinese Prompt
EVENT_IDENTIFICATION_PROMPT_ZH = textwrap.dedent(
    """
你是一个专业的AI助手评估专家。你的任务是根据提供的评估标准，严格评估AI回答的**事件识别准确性**。

<评分标准>
请根据**事件识别**评估回答：
评价标准1. 准确识别：准确检索识别问题中的事件

评价标准2. 时间范围精准：严格匹配问题指定的时间范围，避免模糊表述；

评价标准3. 核心要素准确：是否准确无误地识别和描述了事件的关键信息。例如，在回答"美国CPI数据"问题时，必须准确引用最新的CPI数值、环比与同比增长率、以及核心CPI数据
</评分标准>

<评估步骤>
1. 先判断回答中是否识别到了问题中的事件，能识别到的答案更好;
2. 针对每个评估标准选择更好的回答，并给出你的理由;
3. 综合评估每个评估标准下的结果，然后选出综合最优的答案，并给出理由；
</评估步骤>

<注意事项>
- **客观性**：不要被回答的礼貌程度或长度所影响。
- **聚焦**：仅根据提供的评估标准进行评估。
- 能识别到事件的答案优于不能识别的答案。
</注意事项>

<评分量表>
- 如果回答1更好：rank = [1, 2]
- 如果回答2更好：rank = [2, 1]
</评分量表>

<查询>
{query}
</查询>

<回复1>
{answer_1}
</回复1>

<回复2>
{answer_2}
</回复2>

<输出格式>
请按以下结构化 JSON 格式严格输出你的评估：
{{
    "reason": "<详细解释你的评估理由，包括在各个评估标准下的表现对比>",
    "rank": <[1, 2] 或 [2, 1]>
}}
</输出格式>

JSON:
"""
).strip()

# Build default template from prompts
DEFAULT_EVENT_IDENTIFICATION_TEMPLATE = PromptTemplate(
    messages={
        LanguageEnum.EN: [
            ChatMessage(
                role="user",
                content=EVENT_IDENTIFICATION_PROMPT_EN,
            ),
        ],
        LanguageEnum.ZH: [
            ChatMessage(
                role="user",
                content=EVENT_IDENTIFICATION_PROMPT_ZH,
            ),
        ],
    },
)


class EventIdentificationGrader(LLMGrader):
    """
    Event Identification Grader for Finance Domain

    Purpose:
        Evaluates the quality of event identification in financial responses by comparing
        two responses based on accuracy, time precision, and key element identification.

    What it evaluates:
        - Accurate Identification: Whether the event in the question is correctly identified
        - Time Range Precision: Strictly matching the specified time range
        - Core Elements Accuracy: Accurate identification and description of key event information

    When to use:
        - Evaluating financial event identification responses
        - Comparing accuracy of event retrieval
        - Assessing precision of time-sensitive financial data

    Scoring:
        - rank = [1, 2]: Answer 1 is better
        - rank = [2, 1]: Answer 2 is better

    Args:
        model: BaseChatModel instance or dict config for OpenAIChatModel
        template: Custom evaluation template (default: DEFAULT_EVENT_IDENTIFICATION_TEMPLATE)
        language: Prompt language - EN or ZH (default: LanguageEnum.ZH)

    Returns:
        GraderRank object with:
            - rank: [1, 2] or [2, 1] indicating which answer is better
            - reason: Explanation of the evaluation reasoning
            - metadata: Evaluation details

    Example:
        >>> import asyncio
        >>> from openjudge.models.openai_chat_model import OpenAIChatModel
        >>> from cookbooks.finance_grader.event_interpretation.event_identification import EventIdentificationGrader
        >>>
        >>> model = OpenAIChatModel(api_key="sk-...", model="qwen3-32b")
        >>> grader = EventIdentificationGrader(model=model)
        >>>
        >>> result = asyncio.run(grader.aevaluate(
        ...     query="2024年3月美国CPI数据如何？",
        ...     answer_1="CPI数据有所上涨。",
        ...     answer_2="2024年3月美国CPI同比上涨3.5%，环比上涨0.4%，核心CPI同比上涨3.8%。"
        ... ))
        >>> print(result.rank)  # [2, 1] means answer_2 is better
    """

    def __init__(
        self,
        model: BaseChatModel | dict,
        template: Optional[PromptTemplate] = None,
        language: LanguageEnum = LanguageEnum.ZH,
        strategy: BaseEvaluationStrategy | None = None,
    ):
        """
        Initialize EventIdentificationGrader.

        Args:
            model: The chat model to use for evaluation, either as a BaseChatModel instance or config dict
            template: The prompt template for event identification evaluation.
                     Defaults to DEFAULT_EVENT_IDENTIFICATION_TEMPLATE.
            language: The language for the evaluation prompt. Defaults to LanguageEnum.ZH (Chinese).
            strategy: The evaluation strategy to use. Defaults to DirectEvaluationStrategy.
        """
        super().__init__(
            name="event_identification",
            mode=GraderMode.LISTWISE,
            description="Evaluate financial event identification quality by comparing two responses",
            model=model,
            template=template or DEFAULT_EVENT_IDENTIFICATION_TEMPLATE,
            language=language,
            strategy=strategy,
        )

    async def _aevaluate(
        self,
        query: str,
        answer_1: str,
        answer_2: str,
        **kwargs,
    ) -> GraderRank:
        """
        Evaluate two financial event identification responses

        Args:
            query: The financial question or prompt
            answer_1: First response to evaluate
            answer_2: Second response to evaluate
            **kwargs: Additional arguments

        Returns:
            GraderRank: Rank result with [1, 2] if answer_1 is better, [2, 1] if answer_2 is better

        Example:
            >>> result = await grader.aevaluate(
            ...     query="2024年3月美国CPI数据如何？",
            ...     answer_1="CPI数据有所上涨。",
            ...     answer_2="2024年3月美国CPI同比上涨3.5%，环比上涨0.4%。"
            ... )
        """
        try:
            result = await super()._aevaluate(
                query=query,
                answer_1=answer_1,
                answer_2=answer_2,
            )

            rank = result.rank
            reason = result.reason

            # Validate rank format
            if not isinstance(rank, list) or len(rank) != 2:
                logger.warning(f"Invalid rank format: {rank}, defaulting to [1, 2]")
                rank = [1, 2]

            # Ensure rank is either [1, 2] or [2, 1]
            if set(rank) != {1, 2}:
                logger.warning(f"Invalid rank values: {rank}, defaulting to [1, 2]")
                rank = [1, 2]

            return GraderRank(
                name=self.name,
                rank=rank,
                reason=reason,
                metadata={
                    "evaluation_type": "event_identification",
                    "criteria": ["accurate_identification", "time_precision", "core_elements_accuracy"],
                },
            )

        except Exception as e:
            logger.exception(f"Error evaluating event identification: {e}")
            return GraderError(
                name=self.name,
                error=f"Evaluation error: {str(e)}",
            )


__all__ = ["EventIdentificationGrader", "DEFAULT_EVENT_IDENTIFICATION_TEMPLATE"]
