
from typing import List, Optional, Type
from pydantic import Field
from rm_gallery.core.data.schema import DataSample
from rm_gallery.core.model.message import format_messages
from rm_gallery.core.reward.base import BaseLLMReward
from rm_gallery.core.reward.schema import RewardDimensionWithRank, RewardDimensionWithScore, RewardResult
from rm_gallery.core.reward.template import BasePromptTemplate


class BaseRubricPairWiseTemplate(BasePromptTemplate):
    reason: Optional[str] = Field(
        default=None, description="你的分析和理由"
    )
    best: Optional[int] = Field(
        default=None, description="你的最佳答案,1 or 2"
    )

    @classmethod
    def parse(cls, text: str):
        """
        Parses text input containing principle violation information.

        Processes standard template format and converts violation field
        from string representation to Python list.

        Args:
            text (str): Input string containing XML-style tagged content

        Returns:
            PrinciplePointWiseTemplate: Constructed instance with parsed values
        """
        contents = cls._parse(text)
        try:
            contents["violation"] = eval(contents["violation"])
        except Exception:
            contents["violation"] = []
        return cls(**contents)

    @classmethod
    def format(
        cls,
        desc: str,
        rubrics: str,
        examples: str,
        query: str,
        answer_1: str,
        answer_2: str,
        **kwargs,
    ) -> str:
        if examples:
            examples = f"\n# 样例\n{examples}\n"

        return f"""# 任务描述
{desc}

# 评价标准
{rubrics}

# 示例
{examples}

# 金融相关问题
{query}

# 回答1
[回答1开始]
{answer_1}
[回答1结束]

# 回答2
[回答2开始]
{answer_2}
[回答2结束]

# 输出格式要求
{cls.schema(**kwargs)}
"""
    

class BaseRubricReward(BaseLLMReward):
    rubrics: List[str] = Field(default=..., description="评估标准")
    examples: List[str] = Field(default=[], description="示例")
    template: Type[BasePromptTemplate] = Field(
        default=BaseRubricPairWiseTemplate, description="prompt模板"
    )
    desc: str = Field(default=..., description="任务描述")

    def _before_evaluate(self, sample:DataSample, **kwargs):
        rubrics_str = ""
        for i, rubric in enumerate(self.rubrics):
            rubrics_str += f"评价标准{i + 1}. {rubric}\n"
        query = format_messages(sample.input)
        return {
            "desc": self.desc,
            "rubrics": rubrics_str,
            "examples": "\n".join(self.examples),
            "query": query,
        }

class BasePointWiseRubricReward(BaseRubricReward):
    def _before_evaluate(self, sample:DataSample, **kwargs):
        params = super()._before_evaluate(sample, **kwargs)
        params["answer"] = sample.output[0].answer.content
        return params

    def _after_evaluate(self, response, **kwargs):

        score = (
            1 - len(response.violation) / len(self.rubrics)
            if response.violation
            else 1.0
        )
        return RewardResult(
            name=self.name,
            details=[
                RewardDimensionWithScore(
                    name=self.name, reason=response.reason, score=score
                )
            ],
        )
class BasePairWiseRubricReward(BaseRubricReward):
    def _before_evaluate(self, sample, **kwargs):
        params = super()._before_evaluate(sample, **kwargs)
        params["answer_1"] = sample.output[0].answer.content
        params["answer_2"] = sample.output[1].answer.content
        return params

    def _after_evaluate(
        self, response: BaseRubricPairWiseTemplate, sample: DataSample, **kwargs
    ) -> RewardResult:
        scores = [0.0] * len(sample.output)
        scores[response.best - 1] = 1.0
        return RewardResult(
            name=self.name,
            details=[
                RewardDimensionWithRank(
                    name=self.name, reason=response.reason, rank=scores
                )
            ],
            extra_data={
                "best": response.best - 1,
                "response": response.model_dump(),
            },
        )
    
    

        