"""
数学 Agentic RL 的 Reward 处理器。

Reward 由两个子奖励加权组成：
  - format   (权重 0.2): 模型是否使用了 \\boxed{} 格式输出答案？
                         是 → +1.0，否 → -1.0
                         负分惩罚鼓励模型学会使用指定格式
  - accuracy (权重 0.8): 模型的答案是否数学正确？
                         正确 → 1.0，错误 → 0.0

装饰器模式说明：
  @reward_func("math_reward")   : 将类注册为 reward 函数，自动生成 process() 入口
  @sub_reward_func(name, weight): 定义子奖励函数及其权重
  @aggregate_func               : 定义子奖励的聚合逻辑
  @observe_processor            : 追踪 REWARD span 的可观测性装饰器

reward_metrics 中记录的指标会被训练框架自动采集：
  - format   : 格式分（-1 或 1）
  - accuracy : 正确性分（0 或 1）
  - total    : 加权总分
"""
from __future__ import annotations

from typing import Dict

from dashscope.finetune.reinforcement import logger
from dashscope.finetune.reinforcement import (
    RewardInput, RewardOutput, Reward, TaskStatus,
    reward_func, sub_reward_func, aggregate_func,
    AbstractRewardProcessor,
)
from dashscope.finetune.reinforcement.component.observability import observe_processor

from functions.reward.math_verify import verify_math_answer, extract_answer_by_format


@reward_func("math_reward")
class MathRewardProcessor(AbstractRewardProcessor):
    """数学奖励处理器，评估模型的格式规范性和答案正确性。

    被 @reward_func 装饰后，框架会自动生成 process() 方法，
    按 sub_reward_func 的定义依次调用子奖励函数，最后调用 aggregate_func 聚合。

    继承自 AbstractRewardProcessor，需实现 setup() 和各子奖励方法。
    """

    def setup(self) -> None:
        """初始化回调，可用于加载模型或资源。当前无需额外初始化。"""
        logger.info("[MathRewardProcessor] setup() completed")

    @sub_reward_func("format", sub_weight=0.2)
    def format_score(self, input: RewardInput) -> RewardOutput:
        """格式子奖励：检查模型输出是否使用了 \\boxed{} 格式。

        通过 extract_answer_by_format() 严格检测 \\boxed{} 格式，
        成功提取 → +1.0，未使用格式 → -1.0（惩罚分）。

        Args:
            input: RewardInput，包含：
                - agent_output.messages: rollout 生成的完整对话轨迹
                - agent_output.rollout_extra: 附加信息（含 dataset 名称等）

        Returns:
            RewardOutput: 包含格式分（+1.0 或 -1.0）。
        """
        rollout_extra = getattr(input.agent_output, "rollout_extra", None) or {}
        dataset = rollout_extra.get("dataset", "")

        content = self._get_last_content(input)
        extracted = extract_answer_by_format(content, dataset)
        has_format = extracted is not None

        score = 1.0 if has_format else -1.0

        logger.info(
            f"[MathRewardProcessor] format={dataset or 'default'} "
            f"extracted={extracted!r} → {score}"
        )
        return RewardOutput(
            reward=Reward(reward_score=score),
            status=TaskStatus.SUCCESS,
        )

    @sub_reward_func("accuracy", sub_weight=0.8)
    def accuracy_score(self, input: RewardInput) -> RewardOutput:
        """正确性子奖励：验证模型答案是否与标准答案数学等价。

        调用 verify_math_answer() 进行两层验证：
          1. Hendrycks 标准字符串归一化 + 精确匹配
          2. SymPy 代数等价性回退

        标准答案来源优先级：input.ground_truth > rollout_extra["solution"]

        Args:
            input: RewardInput，包含 ground_truth 和对话轨迹。

        Returns:
            RewardOutput: 包含正确性分（1.0 或 0.0）。
        """
        content = self._get_last_content(input)
        ground_truth = input.ground_truth
        if ground_truth is None:
            extra = getattr(input.agent_output, "rollout_extra", None) or {}
            ground_truth = extra.get("solution")
        ground_truth = str(ground_truth) if ground_truth else ""
        score = verify_math_answer(content, ground_truth)

        logger.info(
            f"[MathRewardProcessor] accuracy={score} | "
            f"ground_truth={ground_truth!r} | content_tail={content[-100:]!r}"
        )
        return RewardOutput(
            reward=Reward(reward_score=score),
            status=TaskStatus.SUCCESS,
        )

    @aggregate_func
    def aggregate(self, sub_rewards: dict[str, RewardOutput]) -> RewardOutput:
        """聚合子奖励：total = 0.2 × format + 0.8 × accuracy。

        通过 self.get_weights() 获取各子奖励权重（由 @sub_reward_func 的 sub_weight 定义），
        通过 self.get_scores() 提取各子奖励分数。

        reward_metrics 中的指标会被框架自动收集并展示在训练监控面板中。

        Args:
            sub_rewards: 子奖励名称 → RewardOutput 的映射。

        Returns:
            RewardOutput: 包含加权总分和各项指标。
        """
        weights = self.get_weights()
        scores = self.get_scores(sub_rewards)
        total = sum(scores[k] * weights[k] for k in scores)

        accuracy = scores.get("accuracy", 0.0)

        reward_metrics: Dict[str, float] = {
            "format": scores.get("format", 0.0),
            "accuracy": accuracy,
            "total": total,
        }

        logger.info(f"[MathRewardProcessor] aggregate: {reward_metrics}")
        return RewardOutput(
            reward=Reward(
                reward_score=total,
                reward_metrics=reward_metrics,
            ),
            status=TaskStatus.SUCCESS,
        )

    @staticmethod
    def _get_last_content(input: RewardInput) -> str:
        """从对话轨迹中提取最后一条消息的内容（即模型最终回答）。"""
        messages = input.agent_output.messages
        return messages[-1].get("content", "") if messages else ""


# @reward_func 自动生成的 process() 不包含 @observe_processor 装饰，
# 这里手动补加，以便 REWARD span 被可观测性系统记录。
MathRewardProcessor.process = observe_processor(MathRewardProcessor.process)
