"""Reward Processor for U-MATH (university-level math) evaluation.

U-MATH golden answers are free-form natural language (multi-part tables,
solution sets joined by \\lor, integration constants C, words like "None"
or "Saddle point"), so exact / sympy matching fails on most samples.

Accuracy pipeline (two modes, controlled per-sample):
  - "sympy_only"  : sympy match → 0/1, no LLM call (explicit override only)
  - "llm_only"    : always use LLM-as-judge (default for all samples)

Mode selection per sample (priority):
  1. rollout_extra["judge_strategy"] explicit override (set during data prep)
  2. default: "llm_only"

Sub-rewards:
  - format   (weight 0.2): did the model use \\boxed{}? → -1 / 1
  - accuracy (weight 0.8): is the answer correct?      →  0 / 1
"""
from __future__ import annotations

import hashlib
import json
import os
import time
from typing import Dict

import httpx

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 dashscope.finetune.reinforcement.component.observability import (
    genai_llm_span,
)

from functions.reward.math_verify import (
    verify_math_answer,
    extract_answer_by_format,
    extract_answer,
)

_LLM_CONFIG_PATH = os.path.join(
    os.path.dirname(__file__), "..", "..", "llm_api_config.json"
)

_JUDGE_PROMPT = """\
You are a math grading assistant. Given the original problem, compare the \
student's answer with the reference answer and decide how much of it is \
mathematically correct.

**Original problem:**
{question}

**Reference answer:**
{golden_answer}

**Student's full response (may contain reasoning):**
{model_response}

Scoring rules:
  - 1.0  : fully correct (mathematically equivalent to the reference)
  - 0.0  : fully incorrect or no answer
  - 0 < score < 1.0 : partial credit, ONLY when the reference contains \
multiple distinct parts (e.g. several sub-questions, multiple items in a \
list, a function plus its domain, or a tuple of independent values). \
Set score = (number of parts the student got right) / (total parts).
Single-part questions only get 0.0 or 1.0.

First think briefly about whether the answers match, then output your decision.
Reply with ONLY a single JSON object, no markdown fences, no extra prose. The \
"reason" field MUST come before the "score" field:
{{"reason": "<one short sentence; if partial, mention how many of how many parts were correct>", "score": <number between 0.0 and 1.0>}}"""

_JUDGE_TIMEOUT_S = 60.0
_JUDGE_MAX_ATTEMPTS = 2  # original + 1 retry on timeout

# Strategy tokens
_STRAT_SYMPY_ONLY = "sympy_only"
_STRAT_LLM_ONLY = "llm_only"
_VALID_STRATEGIES = {_STRAT_SYMPY_ONLY, _STRAT_LLM_ONLY}

# Verbose log sampling: when an LLM judge runs, dump a full multi-line block
# (prompt + raw_response + golden + model preview) for ~SAMPLE_RATE of calls.
# Deterministic by rollout_id (same rollout always samples the same way on retry).
# Override via env var UMATH_JUDGE_LOG_SAMPLE_RATE (float 0..1, default 0.05).
_JUDGE_LOG_SAMPLE_RATE = float(os.environ.get("UMATH_JUDGE_LOG_SAMPLE_RATE", "0.05"))


def _should_sample_log(rollout_id: str | None) -> bool:
    """Deterministic ~rate sampler keyed on rollout_id (md5 first 8 hex chars / 2^32)."""
    if _JUDGE_LOG_SAMPLE_RATE >= 1.0:
        return True
    if _JUDGE_LOG_SAMPLE_RATE <= 0.0:
        return False
    key = rollout_id or "unknown"
    bucket = int(hashlib.md5(key.encode("utf-8")).hexdigest()[:8], 16) / 0x1_0000_0000
    return bucket < _JUDGE_LOG_SAMPLE_RATE



def _load_judge_config() -> dict:
    """Load API key / base_url / model.

    Resolution order:
      1. Environment variables (preferred, no file needed on the rollout pods)
         JUDGE_API_KEY, JUDGE_BASE_URL, JUDGE_MODEL
         (falls back to DASHSCOPE_API_KEY for the key)
      2. llm_api_config.json next to the project root.
    """
    api_key = os.environ.get("JUDGE_API_KEY") or os.environ.get("DASHSCOPE_API_KEY")
    base_url = os.environ.get("JUDGE_BASE_URL")
    model = os.environ.get("JUDGE_MODEL")

    if api_key and base_url and model:
        return {"api_key": api_key, "base_url": base_url, "model": model}

    try:
        with open(_LLM_CONFIG_PATH, encoding="utf-8") as f:
            cfg = json.load(f)
        return {
            "api_key": api_key or cfg["key"],
            "base_url": base_url or cfg["base_url"],
            "model": model or cfg.get("supported_models", ["qwen-plus"])[-1],
        }
    except Exception as e:
        logger.warning(f"[UMathReward] Cannot load judge config: {e}")
        return {}


def _choose_strategy(extra: dict) -> str:
    """Pick judge strategy per sample. See module docstring for priority rules."""
    explicit = (extra or {}).get("judge_strategy")
    if explicit in _VALID_STRATEGIES:
        return explicit
    if explicit is not None:
        logger.warning(
            f"[UMathReward] unknown judge_strategy={explicit!r}, "
            f"falling back to llm_only"
        )
    return _STRAT_LLM_ONLY


def _parse_judge_response(text: str) -> tuple[float, str, str]:
    """Parse a JSON judge response with a fractional score in [0, 1].

    Expected shape: {"reason": "...", "score": <0.0..1.0>}.

    Backward-compat fallbacks (older judge runs / models that ignore format):
      - "verdict" field: "correct" -> 1.0, "incorrect" -> 0.0, "partial" -> 0.5
      - JSON parse fails entirely: substring check on the raw text
        ("correct" without "incorrect" -> 1.0, otherwise 0.0)

    Returns (score, verdict, reason). `verdict` is a derived label
    ("correct" / "incorrect" / "partial") used only for log readability.
    """
    score: float | None = None
    reason = ""

    start = text.find("{")
    end = text.rfind("}")
    if start != -1 and end > start:
        try:
            data = json.loads(text[start:end + 1])
            if isinstance(data, dict):
                reason = str(data.get("reason", "")).strip()
                if "score" in data:
                    try:
                        score = float(data["score"])
                    except (TypeError, ValueError):
                        score = None
                if score is None and "verdict" in data:
                    v = str(data.get("verdict", "")).strip().lower()
                    if v == "correct":
                        score = 1.0
                    elif v == "incorrect":
                        score = 0.0
                    elif v == "partial":
                        score = 0.5
        except Exception:
            pass

    if score is None:
        low = text.lower()
        score = 1.0 if "correct" in low and "incorrect" not in low else 0.0

    score = max(0.0, min(1.0, score))

    if score >= 1.0:
        verdict = "correct"
    elif score <= 0.0:
        verdict = "incorrect"
    else:
        verdict = "partial"

    return score, verdict, reason


def _llm_judge(model_response: str, golden_answer: str, question: str, config: dict) -> dict:
    """Call a lightweight LLM to judge answer correctness.

    Returns:
        {
            "score": 0.0 | 1.0,
            "model": str | None,
            "prompt": str (truncated for storage),
            "raw_response": str | None,
            "reason": str | None,
            "error": str | None,
        }

    Retries once on httpx.TimeoutException; other exceptions return immediately.
    """
    info: dict = {
        "score": 0.0,
        "model": config.get("model") if config else None,
        "prompt": None,
        "raw_response": None,
        "reason": None,
        "error": None,
        "model_response_for_judge": None,
    }

    if not config:
        info["error"] = "judge config missing (set JUDGE_API_KEY/JUDGE_BASE_URL/JUDGE_MODEL or fill llm_api_config.json)"
        logger.warning(f"[UMathReward] LLM judge skipped: {info['error']}")
        return info

    judge_input = model_response[-3000:]
    info["model_response_for_judge"] = judge_input
    prompt = _JUDGE_PROMPT.format(
        question=question,
        golden_answer=golden_answer,
        model_response=judge_input,
    )
    info["prompt"] = prompt[:1000]  # keep storage bounded

    messages = [{"role": "user", "content": prompt}]
    t0 = time.monotonic()

    with genai_llm_span() as inv:
        inv.request_model = config["model"]
        inv.provider = "judge"
        try:
            from dashscope.finetune.reinforcement.component.observability.genai.messages import (
                openai_chat_messages_to_input_messages,
            )
            inv.input_messages = openai_chat_messages_to_input_messages(messages)
        except Exception:
            pass

        for attempt in range(1, _JUDGE_MAX_ATTEMPTS + 1):
            try:
                resp = httpx.post(
                    f"{config['base_url']}/chat/completions",
                    headers={
                        "Authorization": f"Bearer {config['api_key']}",
                        "Content-Type": "application/json",
                    },
                    json={
                        "model": config["model"],
                        "messages": messages,
                        "temperature": 0.0,
                        "max_tokens": 200,
                        "response_format": {"type": "json_object"},
                    },
                    timeout=_JUDGE_TIMEOUT_S,
                )
                resp.raise_for_status()
                text = resp.json()["choices"][0]["message"]["content"].strip()
                score, verdict, reason = _parse_judge_response(text)
                info["raw_response"] = text
                info["reason"] = reason
                info["score"] = score
                info["judge_latency_s"] = time.monotonic() - t0
                info["judge_success"] = 1.0

                try:
                    from dashscope.finetune.reinforcement.component.observability.genai.messages import (
                        openai_completion_to_output_messages,
                    )
                    inv.output_messages = openai_completion_to_output_messages(resp.json())
                    inv.response_model_name = config["model"]
                    usage = resp.json().get("usage")
                    if usage:
                        inv.input_tokens = usage.get("prompt_tokens")
                        inv.output_tokens = usage.get("completion_tokens")
                except Exception:
                    pass

                logger.info(
                    f"[UMathReward] LLM judge verdict='{verdict}' score={score:.3g} "
                    f"latency={info['judge_latency_s']:.2f}s reason={reason!r}"
                )
                return info
            except httpx.TimeoutException as e:
                if attempt < _JUDGE_MAX_ATTEMPTS:
                    logger.warning(
                        f"[UMathReward] LLM judge timeout attempt {attempt}/"
                        f"{_JUDGE_MAX_ATTEMPTS}, retrying ({type(e).__name__})"
                    )
                    continue
                info["judge_latency_s"] = time.monotonic() - t0
                info["judge_success"] = 0.0
                info["error"] = f"{type(e).__name__}: {e}"
                logger.warning(f"[UMathReward] LLM judge failed: {info['error']}")
                return info
            except Exception as e:
                info["judge_latency_s"] = time.monotonic() - t0
                info["judge_success"] = 0.0
                info["error"] = f"{type(e).__name__}: {e}"
                logger.warning(f"[UMathReward] LLM judge failed: {info['error']}")
                return info

    info["judge_latency_s"] = time.monotonic() - t0
    info["judge_success"] = 0.0
    return info


@reward_func("umath_reward")
class UMathRewardProcessor(AbstractRewardProcessor):

    def setup(self) -> None:
        self._judge_config = _load_judge_config()
        logger.info(
            f"[UMathRewardProcessor] setup() completed, "
            f"judge_model={self._judge_config.get('model', 'N/A')}"
        )

    @sub_reward_func("format", sub_weight=0.2)
    def format_score(self, input: RewardInput) -> RewardOutput:
        """Check whether \\boxed{} can be extracted."""
        content = self._get_last_content(input)
        extracted = extract_answer_by_format(content)
        score = 1.0 if extracted is not None else -1.0

        logger.info(f"[UMathReward] format 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:
        """Two-mode accuracy: sympy_only / llm_only (default).

        When LLM judge is invoked, the prompt / raw response / model / error
        are attached to RewardOutput.judge_extra, which the framework records
        in the REWARD span's output.value (trajectory UI -> Reward section).
        """
        content = self._get_last_content(input)
        extra = getattr(input.agent_output, "rollout_extra", None) or {}
        golden = extra.get("solution", "")

        if extract_answer_by_format(content) is None:
            logger.info("[UMathReward] accuracy=0 method=no_boxed (skipped judge)")
            return RewardOutput(
                reward=Reward(
                    reward_score=0.0,
                    reward_metrics={"method_is_judge": 0.0, "method_is_llm": 0.0,
                                    "no_boxed_skip": 1.0, "strategy/skipped": 1.0},
                ),
                status=TaskStatus.SUCCESS,
            )

        strategy = _choose_strategy(extra)

        question = ""
        for msg in (input.agent_output.messages or []):
            if isinstance(msg, dict) and msg.get("role") == "user":
                question = msg.get("content", "")
                break

        judge_info: dict | None = None

        if strategy == _STRAT_SYMPY_ONLY:
            method = "sympy_only"
            score = verify_math_answer(content, golden)
        else:  # llm_only (default)
            method = "llm_only"
            judge_info = _llm_judge(content, golden, question, self._judge_config)
            score = judge_info["score"]

        logger.info(
            f"[UMathReward] accuracy={score} strategy={strategy} method={method} | "
            f"golden={golden[:80]!r}"
        )

        judge_extra: dict | None = None
        if judge_info is not None:
            judge_extra = {
                "strategy": strategy,
                "method": method,
                "final_score": score,
                "model": judge_info.get("model"),
                "raw_response": judge_info.get("raw_response"),
                "reason": judge_info.get("reason"),
                "error": judge_info.get("error"),
                "prompt": judge_info.get("prompt"),
                "golden_preview": (golden or "")[:500],
                "model_response_for_judge": judge_info.get("model_response_for_judge"),
            }

            # Sampled verbose dump — searchable via tag `[UMathReward][JUDGE_SAMPLE]`.
            # Trajectory UI does not surface RewardOutput.judge_extra, so we mirror
            # a subset to worker stdout (controlled by UMATH_JUDGE_LOG_SAMPLE_RATE).
            rid = None
            meta = getattr(input, "request_metadata", None)
            if meta is not None:
                rid = getattr(meta, "rollout_id", None)
            if _should_sample_log(rid):
                logger.info(
                    "[UMathReward][JUDGE_SAMPLE] rollout_id=%s strategy=%s method=%s score=%s\n"
                    "  model=%s reason=%r raw_response=%r error=%s\n"
                    "  golden=%s\n"
                    "  model_response_for_judge=%s",
                    rid or "unknown",
                    strategy,
                    method,
                    score,
                    judge_extra["model"],
                    judge_extra["reason"],
                    judge_extra["raw_response"],
                    judge_extra["error"],
                    judge_extra["golden_preview"],
                    judge_extra["model_response_for_judge"],
                )

        method_is_llm = 1.0 if method in ("llm_only", "llm_judge") else 0.0
        reward_metrics: Dict[str, float] = {
            "method_is_judge": method_is_llm,
            "method_is_llm": method_is_llm,
            "partial_credit": 1.0 if 0.0 < score < 1.0 else 0.0,
            f"strategy/{strategy}": 1.0,
        }
        if judge_info is not None:
            reward_metrics["judge_latency_s"] = judge_info.get("judge_latency_s", 0.0)
            reward_metrics["judge_success"] = judge_info.get("judge_success", 0.0)
        subject = extra.get("subject", "")
        if subject:
            reward_metrics[f"subject/{subject}"] = score

        # RewardOutput.Config.extra = "allow" 允许挂额外字段；
        # 仅在确实跑了 LLM judge 时才挂 judge_extra，避免污染常规样本。
        output_kwargs = {
            "reward": Reward(reward_score=score, reward_metrics=reward_metrics),
            "status": TaskStatus.SUCCESS,
        }
        if judge_extra is not None:
            output_kwargs["judge_extra"] = judge_extra
        return RewardOutput(**output_kwargs)

    @aggregate_func
    def aggregate(self, sub_rewards: dict[str, RewardOutput]) -> RewardOutput:
        """Weighted sum: 0.2 * format + 0.8 * 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),
            "method_is_judge": acc_metrics.get("method_is_judge", 0.0),
            "method_is_llm": acc_metrics.get("method_is_llm", 0.0),
            "total": total,
        }

        logger.info(f"[UMathRewardProcessor] 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 ""


UMathRewardProcessor.process = observe_processor(UMathRewardProcessor.process)
