import argparse
from pathlib import Path

import numpy as np
import torch

from game import UltimateTicTacToe
from mcts import MCTS
from model import UltimateTicTacToeModel
from trainer import Trainer


DEFAULT_ARGS = {
    "num_simulations": 100,
    "numIters": 50,
    "numEps": 20,
    "epochs": 5,
    "batch_size": 64,
    "lr": 5e-4,
    "weight_decay": 1e-4,
    "replay_buffer_size": 50000,
    "value_loss_weight": 1.0,
    "grad_clip_norm": 5.0,
    "checkpoint_path": "latest.pth",
    "temperature_threshold": 10,
    "root_dirichlet_alpha": 0.3,
    "root_exploration_fraction": 0.25,
    "arena_compare_games": 6,
    "arena_accept_threshold": 0.55,
    "arena_compare_simulations": 8,
}


def get_device(device_arg):
    if device_arg:
        return device_arg
    return "cuda" if torch.cuda.is_available() else "cpu"


def build_model(game, device):
    return UltimateTicTacToeModel(
        game.get_board_size(),
        game.get_action_size(),
        device,
    )


def load_checkpoint(model, checkpoint_path, device, optimizer=None, required=True):
    checkpoint = Path(checkpoint_path)
    if not checkpoint.exists():
        if required:
            raise FileNotFoundError(f"Checkpoint not found: {checkpoint}")
        return False

    state = torch.load(checkpoint, map_location=device)
    model.load_state_dict(state["state_dict"])
    if optimizer is not None and "optimizer_state_dict" in state:
        optimizer.load_state_dict(state["optimizer_state_dict"])
    model.eval()
    return True


def canonical_state(game, state, player):
    board_data, active_board = state
    return (game.get_canonical_board_data(board_data, player), active_board)


def apply_moves(game, moves):
    state = game.get_init_board()
    player = 1
    for action in moves:
        next_state = game.get_next_state(state, player, action, verify_move=True)
        if next_state is False:
            raise ValueError(f"Illegal move in sequence: {action}")
        state, player = next_state
    return state, player


def format_board(board_data):
    symbols = {1: "X", -1: "O", 0: "."}
    rows = []
    for row in range(9):
        cells = [symbols[int(board_data[row * 9 + col])] for col in range(9)]
        groups = [" ".join(cells[idx:idx + 3]) for idx in (0, 3, 6)]
        rows.append(" | ".join(groups))
        if row in (2, 5):
            rows.append("-" * 23)
    return "\n".join(rows)


def top_policy_moves(policy, limit):
    ranked = np.argsort(policy)[::-1][:limit]
    return [(int(action), float(policy[action])) for action in ranked]


def parse_moves(text):
    if not text:
        return []
    return [int(part.strip()) for part in text.split(",") if part.strip()]


def parse_action(text):
    raw = text.strip().replace(",", " ").split()
    if len(raw) == 1:
        action = int(raw[0])
    elif len(raw) == 2:
        row, col = (int(value) for value in raw)
        if not (0 <= row < 9 and 0 <= col < 9):
            raise ValueError("Row and column must be in [0, 8].")
        action = row * 9 + col
    else:
        raise ValueError("Enter either a flat move index or 'row col'.")
    if not (0 <= action < 81):
        raise ValueError("Move index must be in [0, 80].")
    return action


def scalar_value(value):
    return float(np.asarray(value).reshape(-1)[0])


def train_command(args):
    device = get_device(args.device)
    game = UltimateTicTacToe()
    model = build_model(game, device)

    train_args = dict(DEFAULT_ARGS)
    train_args.update(
        {
            "num_simulations": args.num_simulations,
            "numIters": args.num_iters,
            "numEps": args.num_eps,
            "epochs": args.epochs,
            "batch_size": args.batch_size,
            "lr": args.lr,
            "weight_decay": args.weight_decay,
            "replay_buffer_size": args.replay_buffer_size,
            "value_loss_weight": args.value_loss_weight,
            "grad_clip_norm": args.grad_clip_norm,
            "checkpoint_path": args.checkpoint,
            "temperature_threshold": args.temperature_threshold,
            "root_dirichlet_alpha": args.root_dirichlet_alpha,
            "root_exploration_fraction": args.root_exploration_fraction,
            "arena_compare_games": args.arena_compare_games,
            "arena_accept_threshold": args.arena_accept_threshold,
            "arena_compare_simulations": args.arena_compare_simulations,
        }
    )

    trainer = Trainer(game, model, train_args)
    if args.resume:
        load_checkpoint(model, args.checkpoint, device, optimizer=trainer.optimizer)
    trainer.learn()


def eval_command(args):
    device = get_device(args.device)
    game = UltimateTicTacToe()
    model = build_model(game, device)
    load_checkpoint(model, args.checkpoint, device)

    moves = parse_moves(args.moves)
    state, player = apply_moves(game, moves)
    current_state = canonical_state(game, state, player)
    encoded = game.encode_state(current_state)
    policy, value = model.predict(encoded)
    legal_mask = np.array(game.get_valid_moves(state), dtype=np.float32)
    policy = policy * legal_mask
    if policy.sum() > 0:
        policy = policy / policy.sum()

    print("Board:")
    print(format_board(state[0]))
    print()
    print(f"Side to move: {'X' if player == 1 else 'O'}")
    print(f"Active small board: {state[1]}")
    print(f"Model value: {scalar_value(value):.4f}")
    print("Top policy moves:")
    for action, prob in top_policy_moves(policy, args.top_k):
        print(f"  {action:2d} -> {prob:.4f}")

    if args.with_mcts:
        mcts_args = dict(DEFAULT_ARGS)
        mcts_args.update(
            {
                "num_simulations": args.num_simulations,
                "root_dirichlet_alpha": None,
                "root_exploration_fraction": None,
            }
        )
        root = MCTS(game, model, mcts_args).run(model, current_state, to_play=1)
        action = root.select_action(temperature=0)
        print(f"MCTS best move: {action}")


def ai_action(game, model, state, player, num_simulations):
    current_state = canonical_state(game, state, player)
    mcts_args = dict(DEFAULT_ARGS)
    mcts_args.update(
        {
            "num_simulations": num_simulations,
            "root_dirichlet_alpha": None,
            "root_exploration_fraction": None,
        }
    )
    root = MCTS(game, model, mcts_args).run(model, current_state, to_play=1)
    return root.select_action(temperature=0)


def random_action(game, state):
    legal_actions = [index for index, allowed in enumerate(game.get_valid_moves(state)) if allowed]
    if not legal_actions:
        raise ValueError("No legal actions available.")
    return int(np.random.choice(legal_actions))


def load_player_model(game, checkpoint, device):
    model = build_model(game, device)
    load_checkpoint(model, checkpoint, device)
    return model


def choose_action(game, player_kind, model, state, player, num_simulations):
    if player_kind == "random":
        return random_action(game, state)
    return ai_action(game, model, state, player, num_simulations)


def play_match(game, x_kind, x_model, o_kind, o_model, num_simulations):
    state = game.get_init_board()
    player = 1

    while True:
        reward = game.get_reward_for_player(state, player)
        if reward is not None:
            if reward == 0:
                return 0
            return player if reward == 1 else -player

        if player == 1:
            action = choose_action(game, x_kind, x_model, state, player, num_simulations)
        else:
            action = choose_action(game, o_kind, o_model, state, player, num_simulations)
        state, player = game.get_next_state(state, player, action)


def arena_command(args):
    device = get_device(args.device)
    game = UltimateTicTacToe()

    x_model = None
    o_model = None
    if args.x_player == "checkpoint":
        x_model = load_player_model(game, args.x_checkpoint, device)
    if args.o_player == "checkpoint":
        o_model = load_player_model(game, args.o_checkpoint, device)

    results = {1: 0, -1: 0, 0: 0}
    for _ in range(args.games):
        winner = play_match(
            game,
            args.x_player,
            x_model,
            args.o_player,
            o_model,
            args.num_simulations,
        )
        results[winner] += 1

    print(f"Games: {args.games}")
    print(f"X ({args.x_player}) wins: {results[1]}")
    print(f"O ({args.o_player}) wins: {results[-1]}")
    print(f"Draws: {results[0]}")


def play_command(args):
    device = get_device(args.device)
    game = UltimateTicTacToe()
    model = build_model(game, device)
    load_checkpoint(model, args.checkpoint, device)

    state = game.get_init_board()
    player = 1
    human_player = args.human_player

    while True:
        print()
        print(format_board(state[0]))
        print(f"Turn: {'X' if player == 1 else 'O'}")
        print(f"Active small board: {state[1]}")

        reward = game.get_reward_for_player(state, player)
        if reward is not None:
            if reward == 0:
                print("Result: draw")
            else:
                winner = player if reward == 1 else -player
                print(f"Winner: {'X' if winner == 1 else 'O'}")
            return

        valid_moves = game.get_valid_moves(state)
        legal_actions = [index for index, allowed in enumerate(valid_moves) if allowed]
        print(f"Legal moves: {legal_actions}")

        if player == human_player:
            while True:
                try:
                    action = parse_action(input("Your move (index or 'row col'): "))
                    next_state = game.get_next_state(state, player, action, verify_move=True)
                    if next_state is False:
                        raise ValueError(f"Illegal move: {action}")
                    state, player = next_state
                    break
                except ValueError as exc:
                    print(exc)
        else:
            action = ai_action(game, model, state, player, args.num_simulations)
            print(f"AI move: {action}")
            state, player = game.get_next_state(state, player, action)


def build_parser():
    parser = argparse.ArgumentParser(description="Ultimate Tic-Tac-Toe Runner")
    subparsers = parser.add_subparsers(dest="command", required=True)

    train_parser = subparsers.add_parser("train", help="Train the model with self-play")
    train_parser.add_argument("--device")
    train_parser.add_argument("--checkpoint", default=DEFAULT_ARGS["checkpoint_path"])
    train_parser.add_argument("--resume", action="store_true")
    train_parser.add_argument("--num-simulations", type=int, default=DEFAULT_ARGS["num_simulations"])
    train_parser.add_argument("--num-iters", type=int, default=DEFAULT_ARGS["numIters"])
    train_parser.add_argument("--num-eps", type=int, default=DEFAULT_ARGS["numEps"])
    train_parser.add_argument("--epochs", type=int, default=DEFAULT_ARGS["epochs"])
    train_parser.add_argument("--batch-size", type=int, default=DEFAULT_ARGS["batch_size"])
    train_parser.add_argument("--lr", type=float, default=DEFAULT_ARGS["lr"])
    train_parser.add_argument("--weight-decay", type=float, default=DEFAULT_ARGS["weight_decay"])
    train_parser.add_argument("--replay-buffer-size", type=int, default=DEFAULT_ARGS["replay_buffer_size"])
    train_parser.add_argument("--value-loss-weight", type=float, default=DEFAULT_ARGS["value_loss_weight"])
    train_parser.add_argument("--grad-clip-norm", type=float, default=DEFAULT_ARGS["grad_clip_norm"])
    train_parser.add_argument("--temperature-threshold", type=int, default=DEFAULT_ARGS["temperature_threshold"])
    train_parser.add_argument("--root-dirichlet-alpha", type=float, default=DEFAULT_ARGS["root_dirichlet_alpha"])
    train_parser.add_argument("--root-exploration-fraction", type=float, default=DEFAULT_ARGS["root_exploration_fraction"])
    train_parser.add_argument("--arena-compare-games", type=int, default=DEFAULT_ARGS["arena_compare_games"])
    train_parser.add_argument("--arena-accept-threshold", type=float, default=DEFAULT_ARGS["arena_accept_threshold"])
    train_parser.add_argument("--arena-compare-simulations", type=int, default=DEFAULT_ARGS["arena_compare_simulations"])
    train_parser.set_defaults(func=train_command)

    eval_parser = subparsers.add_parser("eval", help="Inspect a checkpoint on a position")
    eval_parser.add_argument("--device")
    eval_parser.add_argument("--checkpoint", default=DEFAULT_ARGS["checkpoint_path"])
    eval_parser.add_argument("--moves", default="", help="Comma-separated move sequence")
    eval_parser.add_argument("--top-k", type=int, default=10)
    eval_parser.add_argument("--with-mcts", action="store_true")
    eval_parser.add_argument("--num-simulations", type=int, default=DEFAULT_ARGS["num_simulations"])
    eval_parser.set_defaults(func=eval_command)

    play_parser = subparsers.add_parser("play", help="Play against the checkpoint")
    play_parser.add_argument("--device")
    play_parser.add_argument("--checkpoint", default=DEFAULT_ARGS["checkpoint_path"])
    play_parser.add_argument("--human-player", type=int, choices=[1, -1], default=1)
    play_parser.add_argument("--num-simulations", type=int, default=DEFAULT_ARGS["num_simulations"])
    play_parser.set_defaults(func=play_command)

    arena_parser = subparsers.add_parser("arena", help="Run repeated matches between agents")
    arena_parser.add_argument("--device")
    arena_parser.add_argument("--games", type=int, default=20)
    arena_parser.add_argument("--num-simulations", type=int, default=DEFAULT_ARGS["num_simulations"])
    arena_parser.add_argument("--x-player", choices=["checkpoint", "random"], default="checkpoint")
    arena_parser.add_argument("--o-player", choices=["checkpoint", "random"], default="random")
    arena_parser.add_argument("--x-checkpoint", default=DEFAULT_ARGS["checkpoint_path"])
    arena_parser.add_argument("--o-checkpoint", default=DEFAULT_ARGS["checkpoint_path"])
    arena_parser.set_defaults(func=arena_command)

    return parser


def main():
    parser = build_parser()
    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()