use std::{
    env,
    sync::atomic::{AtomicUsize, Ordering},
};

use axum::{
    Router,
    extract::{Json, State},
    response,
    routing::{get, post},
};
use battlesnake::types::{
    Direction,
    simulation::{Board, Game},
    wire::{Request, Response},
};
use float_ord::FloatOrd;
use futures_util::future::join_all;
use hashbrown::HashMap;
use log::{debug, error, info, trace, warn};
use rand::prelude::*;
use serde::Serialize;
use tokio::{
    fs::File,
    io::AsyncWriteExt,
    net::TcpListener,
    sync::mpsc::{UnboundedSender, unbounded_channel},
    time::{Duration, Instant},
};

static THREADS: AtomicUsize = AtomicUsize::new(1);

#[tokio::main]
async fn main() {
    env_logger::init();

    let record = env::var("RECORD")
        .ok()
        .and_then(|record| {
            record
                .parse()
                .inspect_err(|err| error!("Unable to parse record options: {err}"))
                .ok()
        })
        .unwrap_or(false);

    let (sender, receiver) = if record {
        let (sender, receiver) = unbounded_channel();
        (Some(sender), Some(receiver))
    } else {
        (None, None)
    };

    debug!("Creating routes");
    let app = Router::new()
        .route("/", get(info))
        .route("/start", post(start))
        .route("/move", post(get_move))
        .route("/end", post(end))
        .with_state(sender);

    let threads = env::var("THREADS")
        .ok()
        .and_then(|threads| {
            threads
                .parse()
                .inspect_err(|err| error!("Unable to parse number of threads: {err}"))
                .ok()
        })
        .unwrap_or(1);
    THREADS.store(threads, Ordering::Relaxed);

    debug!("Creating listener");
    let port = env::var("PORT").unwrap_or_else(|_| "8000".into());
    let listener = TcpListener::bind(format!("0.0.0.0:{port}")).await.unwrap();

    if let Some(mut receiver) = receiver {
        debug!("Starting observer");
        tokio::spawn(async move {
            let mut games = HashMap::new();
            while let Some((request_type, request)) = receiver.recv().await {
                match request_type {
                    RequestType::Start => {
                        let game_id = request.game.id.clone();
                        info!("Got start request {game_id}");
                        if let Some(old_requests) = games.insert(game_id, vec![request]) {
                            warn!("evicted duplicate game: {old_requests:?}");
                        }
                    }
                    RequestType::GetMove => {
                        let game_id = request.game.id.clone();
                        info!("Got move request {game_id}");
                        games.entry(game_id).or_default().push(request);
                    }
                    RequestType::End => {
                        let game_id = request.game.id.clone();
                        info!("Got end request {game_id}");
                        if let Some(mut requests) = games.remove(&game_id) {
                            requests.push(request);
                            let json = match serde_json::to_vec_pretty(&requests) {
                                Ok(json) => json,
                                Err(e) => {
                                    error!("Unable to serealize json: {e}");
                                    continue;
                                }
                            };
                            match File::create_new(format!("games/{game_id}.json")).await {
                                Ok(mut file) => {
                                    if let Err(e) = file.write_all(&json).await {
                                        error!("Unable to write jsone: {e}");
                                    }
                                }
                                Err(e) => error!("Unable to open file: {e}"),
                            }
                        } else {
                            warn!("end of game without game: {request:?}");
                        }
                    }
                }
            }
            warn!("Observer stopped");
        });
    }

    debug!("Starting server");
    axum::serve(listener, app).await.unwrap();
}

async fn info() -> response::Json<Info> {
    info!("got info request");
    response::Json(Info {
        apiversion: "1",
        author: "der-informatiker",
        color: "#00FFEE",
        head: "smart-caterpillar",
        tail: "mouse",
        version: env!("CARGO_PKG_VERSION"),
    })
}

#[derive(Debug, Clone, Serialize)]
struct Info {
    apiversion: &'static str,
    author: &'static str,
    color: &'static str,
    head: &'static str,
    tail: &'static str,
    version: &'static str,
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
enum RequestType {
    Start,
    GetMove,
    End,
}

async fn start(
    State(sender): State<Option<UnboundedSender<(RequestType, Request)>>>,
    Json(request): Json<Request>,
) {
    if let Some(sender) = sender
        && let Err(e) = sender.send((RequestType::Start, request.clone()))
    {
        warn!("Unable to observe request: {e}");
    }
    let board = Board::from(&request);
    info!("got start request: {board}");
}

#[allow(clippy::too_many_lines)]
async fn get_move(
    State(sender): State<Option<UnboundedSender<(RequestType, Request)>>>,
    Json(request): Json<Request>,
) -> response::Json<Response> {
    let start = Instant::now();
    if let Some(sender) = sender
        && let Err(e) = sender.send((RequestType::GetMove, request.clone()))
    {
        warn!("Unable to observe request: {e}");
    }
    let board = Game::from(&request);
    let timeout = Duration::from_millis(u64::from(request.game.timeout));
    let id = board.board.get_id(&request.you.id).unwrap_or_else(|| {
        error!("My id is not in the simulation board");
        0
    });
    debug!("got move request: {}", board.board);
    let actions = board.board.valid_actions(id).collect::<Vec<_>>();
    if actions.len() <= 1 {
        info!(
            "only one possible action. Fast forwarding {:?}",
            actions.first()
        );
        return response::Json(Response {
            direction: actions.first().copied().unwrap_or(Direction::Up),
            shout: None,
        });
    }
    info!("valid actions: {actions:?}");

    if start.elapsed() > Duration::from_millis(10) {
        error!(
            "The calculation started late ({}ms)",
            start.elapsed().as_millis()
        );
    }
    let score_fn: fn(&Board, u8) -> u32 = match &*request.game.ruleset.name {
        "solo" => score_solo,
        _ => score_standard,
    };

    let action_futures = (0..THREADS.load(Ordering::Relaxed)).map(|_| {
        let request = request.clone();
        let board = board.clone();
        let mut rng = SmallRng::from_os_rng();
        tokio::task::spawn_blocking(move || {
            let mut mcts_managers: Vec<_> = (0..request.board.snakes.len())
                .map(|id| MctsManager::new(u8::try_from(id).unwrap()))
                .collect();
            let c = f32::sqrt(2.0);
            let mut mcts_actions = Vec::new();
            while start.elapsed() < timeout * 4 / 5 {
                let mut board = board.clone();
                let mut game_over = false;
                while !game_over {
                    mcts_actions.clear();
                    mcts_actions.extend(mcts_managers.iter_mut().filter_map(|mcts_manager| {
                        mcts_manager
                            .next_action(&board.board, c, &mut rng)
                            .map(|action| (mcts_manager.snake, action))
                    }));
                    game_over = board.next_turn_random(&mcts_actions, &mut rng);
                    if mcts_actions.is_empty() {
                        break;
                    }
                }
                if !game_over {
                    board.simulate_random(&mut rng);
                }
                for mcts_manager in &mut mcts_managers {
                    let id = mcts_manager.snake;
                    let score = score_fn(&board.board, id);
                    mcts_manager.apply_score(score);
                }
            }
            let my_mcts_manager = mcts_managers.into_iter().nth(usize::from(id)).unwrap();

            let actions = my_mcts_manager.base.next.map(|action| {
                action.map_or(0.0, |action| action.score as f32 / action.played as f32)
            });

            (actions, my_mcts_manager.base.played)
        })
    });
    let (scores, played) = join_all(action_futures).await.into_iter().fold(
        ([0.0; 4], 0),
        |(mut total, mut games), actions| {
            if let Ok((actions, new_games)) = actions {
                for i in 0..total.len() {
                    total[i] += actions[i];
                }
                games += new_games;
            }
            (total, games)
        },
    );
    for action in actions {
        let score = scores[usize::from(action)];
        info!("{action:?} -> {score}");
    }
    let action = scores
        .into_iter()
        .enumerate()
        .max_by_key(|(_, score)| FloatOrd(*score))
        .map(|(index, _)| match index {
            0 => Direction::Up,
            1 => Direction::Down,
            2 => Direction::Left,
            3 => Direction::Right,
            _ => unreachable!(),
        });

    if let Some(action) = action {
        info!("found action {action:?} after {played} simulations.",);
    } else {
        warn!("unable to find a valid action");
    }
    info!("chose {action:?}");
    response::Json(Response {
        direction: action.unwrap_or(Direction::Up),
        shout: None,
    })
}

const fn score_solo(board: &Board, _id: u8) -> u32 {
    board.turn()
}

fn score_standard(board: &Board, id: u8) -> u32 {
    if board.alive(id) {
        1 + u32::from(board.max_length() == board.length(id))
    } else {
        0
    }
}

async fn end(
    State(sender): State<Option<UnboundedSender<(RequestType, Request)>>>,
    Json(request): Json<Request>,
) {
    if let Some(sender) = sender
        && let Err(e) = sender.send((RequestType::End, request.clone()))
    {
        warn!("Unable to observe request: {e}");
    }
    let board = Board::from(&request);
    info!("got end request: {board}");
}

#[derive(Debug)]
struct ActionInfo {
    score: u32,
    played: u32,
    next: Box<[Option<ActionInfo>; 4]>,
}

impl ActionInfo {
    fn new() -> Self {
        Self {
            score: 0,
            played: 0,
            next: Box::new([None, None, None, None]),
        }
    }

    fn uct(&self, c: f32) -> [Option<f32>; 4] {
        let mut ucts = [None; 4];
        for (action, uct) in self.next.iter().zip(ucts.iter_mut()) {
            if let Some(action) = action {
                let exploitation = action.score as f32 / action.played as f32;
                let exploration = f32::sqrt(f32::ln(self.played as f32) / action.played as f32);
                uct.replace(c.mul_add(exploration, exploitation));
            }
        }
        ucts
    }
}

#[derive(Debug)]
struct MctsManager {
    base: ActionInfo,
    actions: Vec<Direction>,
    expanded: bool,
    snake: u8,
}

impl MctsManager {
    fn new(snake: u8) -> Self {
        Self {
            base: ActionInfo::new(),
            actions: Vec::new(),
            expanded: false,
            snake,
        }
    }

    fn apply_score(&mut self, score: u32) {
        self.base.played += 1;
        self.base.score += score;
        let mut current = &mut self.base;
        for action in &self.actions {
            let Some(new_current) = &mut current.next[usize::from(*action)] else {
                error!("got action without actioninfo");
                break;
            };
            current = new_current;
            current.played += 1;
            current.score += score;
        }
        self.actions.clear();
        self.expanded = false;
    }

    fn next_action(&mut self, board: &Board, c: f32, rng: &mut impl RngCore) -> Option<Direction> {
        if self.expanded {
            return None;
        }

        let mut current = &mut self.base;
        for action in &self.actions {
            let Some(new_current) = &mut current.next[usize::from(*action)] else {
                error!("got action without actioninfo");
                return None;
            };
            current = new_current;
        }

        let ucts = current.uct(c);
        let valid_actions = board.valid_actions(self.snake);
        let ucts: Vec<_> = valid_actions
            .map(|action| (action, ucts[usize::from(action)]))
            .collect();
        trace!("got actions: {ucts:?}");
        if ucts.iter().any(|(_, uct)| uct.is_none()) {
            let action = ucts.iter().filter(|(_, uct)| uct.is_none()).choose(rng)?.0;
            self.expanded = true;
            current.next[usize::from(action)].replace(ActionInfo::new());
            self.actions.push(action);
            return Some(action);
        }

        let action = ucts
            .iter()
            .max_by_key(|(_, uct)| FloatOrd(uct.unwrap_or(f32::NEG_INFINITY)))
            .map(|(action, _)| *action);
        if let Some(action) = action {
            self.actions.push(action);
        }
        action
    }
}
