from __future__ import annotations

import random
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Sequence

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F

from ..game import GameConfig
from ..sdk import Action, ActionType, CatanEnv
from .encoding import encode_action, encode_observation
from .agents import finalize_action


@dataclass
class Transition:
    observation: np.ndarray
    action_vectors: np.ndarray
    action_index: int
    log_prob: float
    value: float
    reward: float
    done: bool


@dataclass
class PPOConfig:
    total_steps: int = 200_000
    rollout_size: int = 4096
    mini_batch_size: int = 512
    epochs: int = 4
    gamma: float = 0.99
    gae_lambda: float = 0.95
    clip_coef: float = 0.2
    entropy_coef: float = 0.01
    value_coef: float = 0.5
    lr: float = 3e-4
    hidden_sizes: Sequence[int] = (512, 512)
    device: str = "cpu"
    win_reward: float = 10.0
    loss_penalty: float = 1.0
    seed: Optional[int] = None
    player_pool: Optional[Sequence[str]] = None
    min_players: Optional[int] = None
    max_players: Optional[int] = None


class ActionScoringNetwork(nn.Module):
    def __init__(self, obs_dim: int, action_dim: int, hidden_sizes: Sequence[int]) -> None:
        super().__init__()
        layers: List[nn.Module] = []
        input_dim = obs_dim + action_dim
        for hidden in hidden_sizes:
            layers.extend([nn.Linear(input_dim, hidden), nn.ReLU()])
            input_dim = hidden
        layers.append(nn.Linear(input_dim, 1))
        self.network = nn.Sequential(*layers)

    def forward(self, obs: torch.Tensor, action_batch: torch.Tensor) -> torch.Tensor:
        if obs.dim() == 1:
            obs = obs.unsqueeze(0)
        obs_expanded = obs.repeat(action_batch.shape[0], 1)
        logits = self.network(torch.cat([obs_expanded, action_batch], dim=-1))
        return logits.squeeze(-1)


class ValueNetwork(nn.Module):
    def __init__(self, obs_dim: int, hidden_sizes: Sequence[int]) -> None:
        super().__init__()
        layers: List[nn.Module] = []
        input_dim = obs_dim
        for hidden in hidden_sizes:
            layers.extend([nn.Linear(input_dim, hidden), nn.ReLU()])
            input_dim = hidden
        layers.append(nn.Linear(input_dim, 1))
        self.network = nn.Sequential(*layers)

    def forward(self, obs: torch.Tensor) -> torch.Tensor:
        return self.network(obs).squeeze(-1)


class SelfPlayPPOTrainer:
    def __init__(
        self,
        game_config: GameConfig,
        config: PPOConfig,
    ) -> None:
        self.cfg = config
        self.game_config = game_config
        self.device = torch.device(config.device)
        self.rng = random.Random(config.seed)
        self.player_pool = list(config.player_pool or game_config.player_names)
        self.min_players = max(3, config.min_players or len(self.player_pool))
        self.max_players = min(4, config.max_players or len(self.player_pool))
        if self.min_players > self.max_players:
            raise ValueError("min_players cannot exceed max_players.")
        self.color_map = self._build_color_map(game_config)
        self.env = CatanEnv(game_config)
        obs_dim = encode_observation(self.env.observe()).shape[0]
        dummy_action = encode_action(Action(ActionType.END_TURN, {}))
        action_dim = dummy_action.shape[0]
        self.actor = ActionScoringNetwork(obs_dim, action_dim, config.hidden_sizes).to(self.device)
        self.critic = ValueNetwork(obs_dim, config.hidden_sizes[::-1]).to(self.device)
        self.optimizer = torch.optim.Adam(
            list(self.actor.parameters()) + list(self.critic.parameters()),
            lr=config.lr,
        )
        self.buffer: List[Transition] = []
        self.total_steps = 0
        self.update_steps = 0
        self.episode_counter = 0
        self.current_episode_reward = 0.0
        self.recent_rewards: List[float] = []
        self.win_counts: Dict[str, int] = {name: 0 for name in self.player_pool}
        self._reset_env_for_episode()

    def train(self, log_interval: int = 10) -> None:
        while self.total_steps < self.cfg.total_steps:
            steps = self._collect_rollout()
            self.total_steps += steps
            metrics = self._update_policy()
            self.update_steps += 1
            if self.update_steps % log_interval == 0:
                avg_reward = np.mean(self.recent_rewards[-log_interval:]) if self.recent_rewards else 0.0
                typer_msg = (
                    f"[PPO] step={self.total_steps} updates={self.update_steps} "
                    f"loss={metrics['loss']:.4f} avg_adv={metrics['adv']:.4f} "
                    f"reward={avg_reward:.3f}"
                )
                print(typer_msg)

    def save(self, path: str) -> None:
        state = {
            "actor": self.actor.state_dict(),
            "critic": self.critic.state_dict(),
            "config": self.cfg,
            "game_config": self.game_config,
        }
        torch.save(state, path)

    def _collect_rollout(self) -> int:
        self.buffer.clear()
        steps = 0
        while steps < self.cfg.rollout_size:
            observation = self.env.observe()
            obs_vec = encode_observation(observation)
            player_name = self.env.game.current_player.name
            legal_actions = self.env.legal_actions()
            if not legal_actions:
                break
            encoded_actions = np.stack([encode_action(action) for action in legal_actions])
            value = (
                self.critic(torch.tensor(obs_vec, dtype=torch.float32, device=self.device))
                .detach()
                .item()
            )
            action_index, log_prob = self._sample_action(obs_vec, encoded_actions, legal_actions)
            selected_template = legal_actions[action_index]
            resolved_action = finalize_action(self.env, selected_template, self.rng)
            _, reward, done, info = self.env.step(resolved_action)
            if done and info.get("winner"):
                if info["winner"] == player_name:
                    reward += self.cfg.win_reward
                else:
                    reward -= self.cfg.loss_penalty
                self.win_counts[info["winner"]] += 1
            self.buffer.append(
                Transition(
                    observation=obs_vec,
                    action_vectors=encoded_actions,
                    action_index=action_index,
                    log_prob=log_prob,
                    value=value,
                    reward=reward,
                    done=done,
                )
            )
            steps += 1
            self.current_episode_reward += reward
            if done:
                self.episode_counter += 1
                self.recent_rewards.append(self.current_episode_reward)
                self.current_episode_reward = 0.0
                self._reset_env_for_episode()
        return steps

    def _build_color_map(self, config: GameConfig) -> Dict[str, str]:
        if not config.colors:
            return {}
        return {name: color for name, color in zip(config.player_names, config.colors)}

    def _select_player_names(self) -> List[str]:
        count = self.rng.randint(self.min_players, self.max_players)
        if count >= len(self.player_pool):
            return list(self.player_pool)
        return self.rng.sample(self.player_pool, k=count)

    def _reset_env_for_episode(self) -> None:
        names = self._select_player_names()
        colors = [self.color_map.get(name, "player") for name in names] if self.color_map else None
        seed = self.rng.randint(0, 2**31 - 1)
        config = GameConfig(player_names=names, colors=colors, seed=seed)
        self.env = CatanEnv(config)

    def _sample_action(
        self,
        obs_vec: np.ndarray,
        action_vectors: np.ndarray,
        legal_actions: List[Action],
    ) -> tuple[int, float]:
        obs_tensor = torch.tensor(obs_vec, dtype=torch.float32, device=self.device)
        actions_tensor = torch.tensor(action_vectors, dtype=torch.float32, device=self.device)
        logits = self.actor(obs_tensor, actions_tensor)
        probs = torch.softmax(logits, dim=0)
        dist = torch.distributions.Categorical(probs=probs)
        index = dist.sample().item()
        log_prob = dist.log_prob(torch.tensor(index, device=self.device)).item()
        return index, log_prob

    def _update_policy(self) -> Dict[str, float]:
        advantages, returns = self._compute_advantages()
        advantages = torch.tensor(advantages, dtype=torch.float32, device=self.device)
        returns = torch.tensor(returns, dtype=torch.float32, device=self.device)
        indices = list(range(len(self.buffer)))
        random.shuffle(indices)
        losses: List[float] = []
        avg_adv: List[float] = []
        for start in range(0, len(indices), self.cfg.mini_batch_size):
            batch_idx = indices[start : start + self.cfg.mini_batch_size]
            if not batch_idx:
                continue
            obs_batch = torch.tensor(
                np.stack([self.buffer[i].observation for i in batch_idx]),
                dtype=torch.float32,
                device=self.device,
            )
            old_log_probs = torch.tensor(
                [self.buffer[i].log_prob for i in batch_idx],
                dtype=torch.float32,
                device=self.device,
            )
            batch_advantages = advantages[batch_idx]
            batch_advantages = (batch_advantages - batch_advantages.mean()) / (
                batch_advantages.std() + 1e-8
            )
            log_probs, entropy = self._evaluate_log_probs(batch_idx)
            ratios = torch.exp(log_probs - old_log_probs)
            unclipped = ratios * batch_advantages
            clipped = torch.clamp(ratios, 1.0 - self.cfg.clip_coef, 1.0 + self.cfg.clip_coef) * batch_advantages
            policy_loss = -torch.min(unclipped, clipped).mean()
            values = self.critic(obs_batch)
            value_loss = F.mse_loss(values, returns[batch_idx])
            entropy_bonus = entropy.mean()
            loss = (
                policy_loss
                + self.cfg.value_coef * value_loss
                - self.cfg.entropy_coef * entropy_bonus
            )
            self.optimizer.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(self.actor.parameters(), max_norm=0.5)
            nn.utils.clip_grad_norm_(self.critic.parameters(), max_norm=0.5)
            self.optimizer.step()
            losses.append(loss.item())
            avg_adv.append(batch_advantages.mean().item())
        return {"loss": float(np.mean(losses)), "adv": float(np.mean(avg_adv))}

    def _evaluate_log_probs(self, batch_idx: List[int]) -> tuple[torch.Tensor, torch.Tensor]:
        log_probs: List[torch.Tensor] = []
        entropies: List[torch.Tensor] = []
        for idx in batch_idx:
            transition = self.buffer[idx]
            obs_tensor = torch.tensor(
                transition.observation, dtype=torch.float32, device=self.device
            )
            actions_tensor = torch.tensor(
                transition.action_vectors, dtype=torch.float32, device=self.device
            )
            logits = self.actor(obs_tensor, actions_tensor)
            probs = torch.softmax(logits, dim=0)
            log_probs.append(torch.log(probs[transition.action_index] + 1e-8))
            entropies.append(-(probs * torch.log(probs + 1e-8)).sum())
        return torch.stack(log_probs), torch.stack(entropies)

    def _compute_advantages(self) -> tuple[List[float], List[float]]:
        advantages: List[float] = []
        returns: List[float] = []
        gae = 0.0
        next_value = 0.0
        for transition in reversed(self.buffer):
            mask = 0.0 if transition.done else 1.0
            delta = transition.reward + self.cfg.gamma * next_value * mask - transition.value
            gae = delta + self.cfg.gamma * self.cfg.gae_lambda * mask * gae
            advantages.insert(0, gae)
            returns.insert(0, gae + transition.value)
            next_value = transition.value
        return advantages, returns
