from __future__ import annotations

from typing import Any

import numpy as np

from blokus_gym.core.bots import (
    Bot,
    GreedyBot,
    GreedyCornersBot,
    MinimaxBot,
    RandomBot,
)
from blokus_gym.core.game import BlokusGame
from blokus_gym.core.pieces import STANDARD_PIECES, PieceSet


class GameSession:
    """Wraps BlokusGame for human-vs-bot play with move history and serialization.

    Tracks which players are human vs bot, records every move for replay,
    and provides methods to handle human input and auto-play bot turns.
    """

    BOT_TYPES: dict[str, type[Bot]] = {
        "random": RandomBot,
        "greedy": GreedyBot,
        "greedy_corners": GreedyCornersBot,
        "minimax": MinimaxBot,
    }

    PLAYER_NAMES = {0: "Red", 1: "Blue", 2: "Yellow", 3: "Green"}
    PLAYER_COLORS = {0: "red", 1: "blue", 2: "yellow", 3: "green"}

    def __init__(
        self,
        board_size: int = 20,
        num_players: int = 4,
        human_player: int = 0,
        bot_type: str = "greedy",
        bot_seed: int = 42,
        minimax_depth: int = 2,
    ):
        assert 2 <= num_players <= 4
        assert 0 <= human_player < num_players
        assert bot_type in self.BOT_TYPES

        self.board_size = board_size
        self.num_players = num_players
        self.human_player = human_player
        self.bot_type_name = bot_type
        self.bot_seed = bot_seed
        self.minimax_depth = minimax_depth

        self.game = BlokusGame(
            board_size=board_size,
            pieces=PieceSet(STANDARD_PIECES),
            num_players=num_players,
        )
        self.game.reset()

        self.move_history: list[dict[str, Any]] = []
        self.game_over = False
        self.winner: int | None = None
        self.last_bot_move: dict[str, Any] | None = None

        self._init_bots()
        self._play_bots_until_human()

    def _init_bots(self) -> None:
        bot_class = self.BOT_TYPES[self.bot_type_name]
        self.bots: dict[int, Bot] = {}
        for i in range(self.num_players):
            if i == self.human_player:
                continue
            if bot_class == MinimaxBot:
                self.bots[i] = bot_class(
                    player_idx=i, depth=self.minimax_depth, seed=self.bot_seed + i * 1000
                )
            else:
                self.bots[i] = bot_class(player_idx=i, seed=self.bot_seed + i * 1000)

    def _play_bots_until_human(self) -> None:
        max_turns = self.num_players * 20
        turns = 0
        while self.game.current_player != self.human_player and not self.game.is_game_over():
            if turns > max_turns:
                break
            self._play_bot_turn()
            turns += 1

    def _play_bot_turn(self) -> None:
        player_idx = self.game.current_player
        bot = self.bots.get(player_idx)
        if bot is None:
            self.game.next_player()
            return

        valid_actions = self.game.get_valid_actions(player_idx)
        if not np.any(valid_actions):
            self.game.players[player_idx].can_move = False
            self.game.next_player()
            return

        action = bot.select_action(self.game, valid_actions)
        if action is not None:
            self._record_move(player_idx, action)
            self.game.play_move(player_idx, action)
            self.last_bot_move = {
                "player_idx": player_idx,
                "piece_name": self.game.piece_set.get_piece(
                    self.game.get_move(action).piece_id
                ).name,
                "x": self.game.get_move(action).x,
                "y": self.game.get_move(action).y,
            }
        else:
            self.game.players[player_idx].can_move = False

        self.game.next_player()

    def _record_move(self, player_idx: int, action: int) -> None:
        move = self.game.get_move(action)
        piece = self.game.piece_set.get_piece(move.piece_id)
        self.move_history.append(
            {
                "player_idx": player_idx,
                "action": action,
                "piece_name": piece.name,
                "piece_id": move.piece_id,
                "orientation_id": move.orientation_id,
                "x": move.x,
                "y": move.y,
            }
        )

    def get_valid_placements(self, piece_id: int) -> list[dict[str, Any]]:
        if self.game.current_player != self.human_player or self.game_over:
            return []

        mask = self.game.get_valid_actions(self.human_player)
        placements: list[dict[str, Any]] = []

        for action_idx in np.where(mask)[0]:
            move = self.game.get_move(int(action_idx))
            if move.piece_id != piece_id:
                continue
            cells = self.game._get_placed_squares(move)
            placements.append(
                {
                    "action": int(action_idx),
                    "cells": [[x, y] for x, y in cells],
                    "x": move.x,
                    "y": move.y,
                }
            )

        return placements

    def apply_human_action(self, action: int) -> dict[str, Any]:
        if self.game_over or self.game.current_player != self.human_player:
            return {"success": False, "error": "Not human's turn"}

        valid = self.game.get_valid_actions(self.human_player)
        if action < 0 or action >= len(valid) or not valid[action]:
            return {"success": False, "error": "Invalid action"}

        self._record_move(self.human_player, action)
        self.game.play_move(self.human_player, action)
        self.game.next_player()

        self._play_bots_until_human()

        if self.game.is_game_over():
            self.game_over = True
            scores = self.game.get_scores()
            max_score = max(scores)
            winners = [i for i, s in enumerate(scores) if s == max_score]
            self.winner = winners[0] if len(winners) == 1 else None

        return {"success": True}

    def get_state(self) -> dict[str, Any]:
        board = self.game.board.grid.tolist()
        scores = self.game.get_scores()

        players = []
        for i in range(self.num_players):
            player = self.game.players[i]
            players.append(
                {
                    "idx": i,
                    "name": self.PLAYER_NAMES.get(i, f"Player {i + 1}"),
                    "color": self.PLAYER_COLORS.get(i, "grey"),
                    "is_human": i == self.human_player,
                    "pieces_remaining": len(player.available_pieces),
                    "total_pieces": len(self.game.piece_set.pieces),
                    "score": player.score,
                }
            )

        human_player_state = self.game.players[self.human_player]
        available_pieces = []
        for piece in self.game.piece_set.pieces:
            if piece.name in human_player_state.available_pieces:
                min_x = min(s[0] for s in piece.squares)
                min_y = min(s[1] for s in piece.squares)
                available_pieces.append(
                    {
                        "id": self.game.piece_set.get_piece_id(piece.name),
                        "name": piece.name,
                        "size": piece.size,
                        "squares": [
                            [x - min_x, y - min_y] for x, y in sorted(piece.squares)
                        ],
                    }
                )

        return {
            "board": board,
            "board_size": self.board_size,
            "current_player": self.game.current_player,
            "human_player": self.human_player,
            "game_over": self.game_over,
            "scores": scores,
            "players": players,
            "available_pieces": available_pieces,
            "move_history": list(self.move_history),
            "total_moves": len(self.move_history),
            "winner": self.winner,
        }

    def get_statistics(self) -> dict[str, Any]:
        board = self.game.board.grid
        total_cells = self.board_size * self.board_size
        occupied = int(np.count_nonzero(board))
        coverage = occupied / total_cells

        scores = self.game.get_scores()

        pieces_placed = []
        for i in range(self.num_players):
            total = len(self.game.piece_set.pieces)
            remaining = len(self.game.players[i].available_pieces)
            pieces_placed.append(total - remaining)

        largest_piece = []
        for i in range(self.num_players):
            max_size = 0
            for move in self.move_history:
                if move["player_idx"] == i:
                    piece = self.game.piece_set.get_piece(move["piece_id"])
                    max_size = max(max_size, piece.size)
            largest_piece.append(max_size)

        if self.game_over:
            sorted_scores = sorted(scores, reverse=True)
            margin = sorted_scores[0] - sorted_scores[1] if len(sorted_scores) > 1 else 0
        else:
            margin = 0

        monomino_bonus = []
        for i in range(self.num_players):
            bonus = 0
            player_moves = [m for m in self.move_history if m["player_idx"] == i]
            if player_moves:
                last_move = player_moves[-1]
                piece = self.game.piece_set.get_piece(last_move["piece_id"])
                if piece.size == 1:
                    bonus = 5
            monomino_bonus.append(bonus)

        return {
            "coverage": round(coverage * 100, 1),
            "pieces_placed": pieces_placed,
            "squares_placed": scores,
            "largest_piece": largest_piece,
            "winner_margin": margin,
            "monomino_bonus": monomino_bonus,
            "total_moves": len(self.move_history),
        }

    def to_dict(self) -> dict[str, Any]:
        return {
            "config": {
                "board_size": self.board_size,
                "num_players": self.num_players,
                "human_player": self.human_player,
                "bot_type": self.bot_type_name,
                "bot_seed": self.bot_seed,
                "minimax_depth": self.minimax_depth,
            },
            "move_history": self.move_history,
            "game_over": self.game_over,
            "winner": self.winner,
            "scores": self.game.get_scores(),
            "statistics": self.get_statistics(),
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> GameSession:
        config = data["config"]
        session = cls.__new__(cls)
        session.board_size = config["board_size"]
        session.num_players = config["num_players"]
        session.human_player = config["human_player"]
        session.bot_type_name = config["bot_type"]
        session.bot_seed = config["bot_seed"]
        session.minimax_depth = config.get("minimax_depth", 2)

        session.game = BlokusGame(
            board_size=session.board_size,
            pieces=PieceSet(STANDARD_PIECES),
            num_players=session.num_players,
        )
        session.game.reset()

        session.move_history = []
        session.game_over = False
        session.winner = None
        session.last_bot_move = None

        session._init_bots()

        for move in data["move_history"]:
            action = move["action"]
            session._record_move(move["player_idx"], action)
            session.game.play_move(move["player_idx"], action)
            session.game.next_player()

        if data["game_over"]:
            session.game_over = True
            session.winner = data["winner"]

        return session
