from typing import Callable

from pydantic import BaseModel


class RewardConfig(BaseModel):
    tool_call_penalty: float = 0.0
    length_penalty: float = 0.0


def tool_call_penalty(messages: list[dict], penalty_per_error: float) -> float:
    num_errors = 0
    for msg in messages:
        if msg.get("role") != "assistant":
            continue
        content = msg.get("content")
        if content is None or not isinstance(content, dict):
            continue
        for tc in content.get("tool_calls", []):
            if tc.get("error") is not None:
                num_errors += 1
    return -penalty_per_error * num_errors


def length_penalty(messages: list[dict], penalty_per_message: float) -> float:
    count = sum(1 for msg in messages if msg.get("role") != "system")
    return -penalty_per_message * count


def cea_reward(output: dict | None, ground_truth: list[dict]) -> float:
    if output is None:
        return 0.0 if ground_truth else 1.0

    pred_annots = output.get("annotations", [])

    gt_map: dict[tuple[int, int], str] = {}
    for ann in ground_truth:
        gt_map[(ann["row"], ann["column"])] = ann["identifier"]

    pred_map: dict[tuple[int, int], str] = {}
    for ann in pred_annots:
        pred_map[(ann["row"], ann["column"])] = ann["identifier"]

    if not gt_map and not pred_map:
        return 1.0
    if not gt_map:
        return 0.0

    tp = sum(
        1 for key, ident in pred_map.items()
        if gt_map.get(key) == ident
    )

    precision = tp / len(pred_map) if pred_map else 0.0
    recall = tp / len(gt_map)

    if precision + recall == 0:
        return 0.0

    return 2 * precision * recall / (precision + recall)


_REWARD_REGISTRY: dict[str, Callable[[dict | None, list[dict]], float]] = {
    "cea": cea_reward,
}


def get_reward_fn(task: str) -> Callable[[dict | None, list[dict]], float]:
    if task not in _REWARD_REGISTRY:
        raise ValueError(f"No reward function for task {task!r}")
    return _REWARD_REGISTRY[task]


def compute_reward(
    output: dict | None,
    ground_truth: list[dict],
    messages: list[dict],
    task: str,
    reward_config: RewardConfig,
) -> float:
    reward = get_reward_fn(task)(output, ground_truth)
    if reward_config.tool_call_penalty > 0:
        reward += tool_call_penalty(messages, reward_config.tool_call_penalty)
    if reward_config.length_penalty > 0:
        reward += length_penalty(messages, reward_config.length_penalty)
    return reward


def compute_step_rewards(
    output: dict | None,
    ground_truth: list[dict],
    messages: list[dict],
    task: str,
    reward_config: RewardConfig,
) -> list[float]:
    task_reward = get_reward_fn(task)(output, ground_truth)

    # One reward per assistant turn
    step_rewards: list[float] = []
    for msg in messages:
        if msg.get("role") != "assistant":
            continue

        step_reward = 0.0

        # Per-step tool call error penalty
        if reward_config.tool_call_penalty > 0:
            content = msg.get("content")
            if isinstance(content, dict):
                for tc in content.get("tool_calls", []):
                    if tc.get("error") is not None:
                        step_reward -= reward_config.tool_call_penalty

        # Per-step length penalty
        if reward_config.length_penalty > 0:
            step_reward -= reward_config.length_penalty

        step_rewards.append(step_reward)

    # Add task reward to the final step
    if step_rewards:
        step_rewards[-1] += task_reward
    else:
        step_rewards = [task_reward]

    return step_rewards
