from __future__ import annotations

import random
from typing import Dict, List, Optional, Protocol

from ..data import Resource
from ..sdk import Action, ActionType, CatanEnv, parse_resource


class Agent(Protocol):
    def choose_action(self, env: CatanEnv, legal_actions: List[Action]) -> Action:
        """Select an action from the legal action list."""


def _resource_from_value(value: str | Resource | None) -> Optional[Resource]:
    if value is None:
        return None
    if isinstance(value, Resource):
        return value
    return parse_resource(value)


def finalize_action(env: CatanEnv, template: Action, rng: Optional[random.Random] = None) -> Action:
    payload = dict(template.payload or {})
    action_type = template.type

    if action_type in {ActionType.MOVE_ROBBER, ActionType.PLAY_KNIGHT}:
        options = payload.get("options", [])
        if not options:
            raise ValueError("No robber options available.")
        choice = rng.choice(options) if rng else options[0]
        new_payload: Dict[str, object] = {"hex": choice["hex"]}
        victims = choice.get("victims") or []
        if victims:
            new_payload["victim"] = rng.choice(victims) if rng else victims[0]
        return Action(action_type, new_payload)

    if action_type == ActionType.PLAY_ROAD_BUILDING:
        edges = payload.get("edges", [])
        if not edges:
            return Action(ActionType.END_TURN, {})
        if rng:
            picks = rng.sample(edges, k=min(2, len(edges)))
            while len(picks) < 2:
                picks.append(rng.choice(edges))
        else:
            picks = edges[:2] if len(edges) >= 2 else edges + edges[: 2 - len(edges)]
        return Action(action_type, {"edges": picks[:2]})

    if action_type == ActionType.PLAY_YEAR_OF_PLENTY:
        bank = payload.get("bank", {})
        available = [res for res, amount in bank.items() if amount > 0]
        if not available:
            available = list(bank.keys())
        if not available:
            raise ValueError("No bank resources available for Year of Plenty.")
        def pick() -> str:
            return rng.choice(available) if rng else available[0]

        return Action(action_type, {"resources": [pick(), pick()]})

    if action_type == ActionType.PLAY_MONOPOLY:
        choices = payload.get("resources") or [res.value for res in Resource if res != Resource.DESERT]
        if not choices:
            raise ValueError("No resources listed for Monopoly.")
        selection = rng.choice(choices) if rng else choices[0]
        return Action(action_type, {"resource": selection})

    if action_type == ActionType.DISCARD:
        required = payload.get("required")
        if not isinstance(required, int):
            raise ValueError("Discard action missing required count.")
        player_name = payload.get("player", env.game.current_player.name)
        player = env.game.player_by_name(player_name)
        cards: Dict[Resource, int] = {res: 0 for res in Resource if res != Resource.DESERT}
        total = required
        resource_list: List[Resource] = []
        for resource, count in player.resources.items():
            if resource == Resource.DESERT or count <= 0:
                continue
            resource_list.extend([resource] * count)
        if rng:
            rng.shuffle(resource_list)
        else:
            resource_list.sort(key=lambda r: (-player.resources[r], r.value))
        for res in resource_list:
            if total <= 0:
                break
            cards[res] += 1
            total -= 1
        card_payload = {res.value: amount for res, amount in cards.items() if amount > 0}
        return Action(
            ActionType.DISCARD,
            {"player": player_name, "cards": card_payload},
        )

    return Action(action_type, payload)


class RandomAgent:
    def __init__(self, seed: Optional[int] = None):
        self.rng = random.Random(seed)

    def choose_action(self, env: CatanEnv, legal_actions: List[Action]) -> Action:
        if not legal_actions:
            raise ValueError("No legal actions available.")
        template = self.rng.choice(legal_actions)
        return finalize_action(env, template, self.rng)


class PolicyAgent:
    def __init__(self, policy, device: str = "cpu", stochastic: bool = True):
        self.policy = policy
        self.device = device
        self.stochastic = stochastic

    def choose_action(self, env: CatanEnv, legal_actions: List[Action]) -> Action:
        from .encoding import encode_action, encode_observation  # Lazy import to avoid cycles
        import torch

        if not legal_actions:
            raise ValueError("No legal actions available.")
        obs = encode_observation(env.observe())
        obs_tensor = torch.tensor(obs, dtype=torch.float32, device=self.device)
        action_vectors = torch.tensor(
            [encode_action(action) for action in legal_actions],
            dtype=torch.float32,
            device=self.device,
        )
        logits = self.policy(obs_tensor, action_vectors)
        probs = torch.softmax(logits, dim=0)
        if self.stochastic:
            dist = torch.distributions.Categorical(probs=probs)
            idx = dist.sample().item()
        else:
            idx = torch.argmax(probs).item()
        selected = legal_actions[idx]
        return finalize_action(env, selected, None)
