"""Reward Processor for code generation evaluation.

Uses subprocess-based code execution to evaluate model-generated solutions
against test cases.  Supports both stdin-based (competitive programming)
and call-based (LeetCode functional) problems.

Sub-rewards:
  - format   (weight 0.1): does the response contain a ```python``` block?
  - accuracy (weight 0.9): passed / total (proportional score for training);
        pass@1 (strict) is always reported in reward_metrics for eval analysis.
"""
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.code_testing import (
    extract_code,
    run_tests,
    normalise_test_cases,
)


@reward_func("code_reward")
class CodeRewardProcessor(AbstractRewardProcessor):

    def setup(self) -> None:
        logger.info("[CodeRewardProcessor] setup() completed")

    @sub_reward_func("format", sub_weight=0.1)
    def format_score(self, input: RewardInput) -> RewardOutput:
        """Check whether a ```python``` code block can be extracted."""
        content = self._get_last_code_content(input)
        code = extract_code(content)
        score = 1.0 if code is not None else -1.0

        logger.info(f"[CodeRewardProcessor] format: code_extracted={code is not None} -> {score}")
        return RewardOutput(
            reward=Reward(reward_score=score),
            status=TaskStatus.SUCCESS,
        )

    @sub_reward_func("accuracy", sub_weight=0.9)
    def accuracy_score(self, input: RewardInput) -> RewardOutput:
        """Run extracted code against test cases.

        reward_score = passed / total  (proportional, smooth training signal)
        reward_metrics includes pass@1 (strict) for evaluation analysis.
        """
        content = self._get_last_code_content(input)
        code = extract_code(content)

        if code is None:
            logger.info("[CodeRewardProcessor] accuracy: no code extracted -> 0.0")
            return RewardOutput(
                reward=Reward(reward_score=0.0),
                status=TaskStatus.SUCCESS,
            )

        extra = getattr(input.agent_output, "rollout_extra", None) or {}
        test_cases, fn_name = normalise_test_cases(extra)
        score, meta = run_tests(code, test_cases, fn_name=fn_name)

        total = meta.get("total", 0)
        passed = meta.get("passed", 0)
        pass_rate = passed / total if total else 0.0
        pass_at_1 = 1.0 if passed == total and total > 0 else 0.0

        logger.info(
            f"[CodeRewardProcessor] accuracy={pass_rate:.3f} | "
            f"pass@1={pass_at_1} | {passed}/{total} | fn_name={fn_name}"
        )

        reward_metrics: Dict[str, float] = {
            "pass_rate": pass_rate,
            "pass@1": pass_at_1,
        }
        dataset = extra.get("dataset", "")
        difficulty = extra.get("difficulty", "")
        if dataset:
            reward_metrics[f"dataset/{dataset}"] = pass_rate
        if difficulty:
            reward_metrics[f"difficulty/{difficulty}"] = pass_rate

        return RewardOutput(
            reward=Reward(reward_score=pass_rate, reward_metrics=reward_metrics),
            status=TaskStatus.SUCCESS,
        )

    @aggregate_func
    def aggregate(self, sub_rewards: dict[str, RewardOutput]) -> RewardOutput:
        """Weighted sum: 0.1 * format + 0.9 * accuracy."""
        weights = self.get_weights()
        scores = self.get_scores(sub_rewards)
        total = sum(scores[k] * weights[k] for k in scores)

        acc_metrics = (sub_rewards.get("accuracy", RewardOutput(
            reward=Reward(reward_score=0.0), status=TaskStatus.SUCCESS,
        )).reward.reward_metrics or {})

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

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

    @staticmethod
    def _get_last_code_content(input: RewardInput) -> str:
        """Scan ALL assistant messages and return the last ```python``` block.

        In multi-turn rollouts the final code may not be in the very last
        message (which could be a summary or tool result).  We iterate in
        reverse and return the first extracted code, or fall back to the
        last assistant message's full text.
        """
        messages = input.agent_output.messages
        for msg in reversed(messages):
            if msg.get("role") != "assistant":
                continue
            content = msg.get("content", "")
            if not content:
                continue
            code = extract_code(content)
            if code is not None:
                return content
        for msg in reversed(messages):
            if msg.get("role") == "assistant" and msg.get("content"):
                return msg["content"]
        return ""


# @reward_func auto-generates process() but without @observe_processor;
# apply it post-hoc so the REWARD entry span is emitted.
CodeRewardProcessor.process = observe_processor(CodeRewardProcessor.process)
