from __future__ import annotations

from typing import List, Optional

import torch
import typer

from .game import GameConfig
from .sdk import CatanEnv
from .ml.trainers import ReinforcementLearningTrainer
from .ml.selfplay import SelfPlayPPOTrainer, PPOConfig

app = typer.Typer(help="Reinforcement learning utilities for Catan.")


def _parse_players(players: str) -> List[str]:
    names = [name.strip() for name in players.split(",") if name.strip()]
    if len(names) < 3:
        raise typer.BadParameter("Нужно от трёх до четырёх игроков.")
    if len(names) > 4:
        raise typer.BadParameter("Базовая игра поддерживает максимум четырёх игроков.")
    return names


def _parse_layers(hidden: str) -> List[int]:
    try:
        layers = [int(value) for value in hidden.split(",") if value.strip()]
    except ValueError as exc:
        raise typer.BadParameter("Слои должны быть целыми числами через запятую.") from exc
    return layers or [256, 256]


@app.command()
def reinforce(
    players: str = typer.Option("Alpha,Beta,Gamma,Delta", help="Игроки через запятую для тренировки"),
    episodes: int = typer.Option(50, help="Количество эпизодов обучения"),
    seed: Optional[int] = typer.Option(None, help="Seed генератора"),
    lr: float = typer.Option(3e-4, help="Learning rate для оптимизации"),
    gamma: float = typer.Option(0.99, help="Коэффициент дисконтирования"),
    device: str = typer.Option("cpu", help="Устройство PyTorch (cpu/cuda)"),
    hidden_layers: str = typer.Option("256,256", help="Слои сети через запятую"),
    output: Optional[str] = typer.Option(
        None, help="Путь для сохранения state_dict обученной политики"
    ),
) -> None:
    names = _parse_players(players)
    layers = _parse_layers(hidden_layers)
    config = GameConfig(player_names=names, seed=seed)
    env = CatanEnv(config)
    trainer = ReinforcementLearningTrainer(
        env,
        hidden_layers=layers,
        lr=lr,
        gamma=gamma,
        device=device,
    )
    typer.echo(f"Запуск обучения на {episodes} эпизодов...")
    rewards = trainer.train(episodes=episodes)
    typer.echo("Средние награды по эпизодам:")
    for idx, reward in enumerate(rewards, start=1):
        typer.echo(f"{idx:03d}: {reward:.3f}")
    if output:
        torch.save(trainer.policy.state_dict(), output)
        typer.echo(f"Модель сохранена в {output}")


@app.command()
def selfplay(
    players: str = typer.Option("Alpha,Beta,Gamma,Delta", help="Пул игроков через запятую"),
    min_players: int = typer.Option(3, help="Минимум игроков в эпизоде"),
    max_players: int = typer.Option(4, help="Максимум игроков в эпизоде"),
    total_steps: int = typer.Option(500_000, help="Количество шагов сбора опыта"),
    rollout_size: int = typer.Option(4096, help="Размер батча rollout"),
    mini_batch: int = typer.Option(512, help="Размер минибатча PPO"),
    epochs: int = typer.Option(4, help="PPO эпох на обновление"),
    gamma: float = typer.Option(0.99, help="Коэффициент дисконтирования"),
    gae_lambda: float = typer.Option(0.95, help="Лямбда GAE"),
    clip_coef: float = typer.Option(0.2, help="PPO clip"),
    entropy_coef: float = typer.Option(0.01, help="Коэффициент энтропии"),
    value_coef: float = typer.Option(0.5, help="Коэффициент value loss"),
    lr: float = typer.Option(3e-4, help="Learning rate"),
    device: str = typer.Option("cpu", help="Устройство PyTorch"),
    seed: Optional[int] = typer.Option(None, help="Seed генератора"),
    output: Optional[str] = typer.Option(None, help="Путь для сохранения весов"),
) -> None:
    names = _parse_players(players)
    if min_players < 3:
        raise typer.BadParameter("Базовая игра требует минимум 3 игроков.")
    if max_players < min_players:
        raise typer.BadParameter("Максимум игроков должен быть >= минимуму.")
    if max_players > len(names):
        raise typer.BadParameter("Максимум игроков превышает размер пула.")
    config = PPOConfig(
        total_steps=total_steps,
        rollout_size=rollout_size,
        mini_batch_size=mini_batch,
        epochs=epochs,
        gamma=gamma,
        gae_lambda=gae_lambda,
        clip_coef=clip_coef,
        entropy_coef=entropy_coef,
        value_coef=value_coef,
        lr=lr,
        device=device,
        seed=seed,
        player_pool=names,
        min_players=min_players,
        max_players=max_players,
    )
    trainer = SelfPlayPPOTrainer(GameConfig(player_names=names, seed=seed), config)
    trainer.train()
    if output:
        trainer.save(output)
        typer.echo(f"Self-play модель сохранена в {output}")


if __name__ == "__main__":
    app()
