pub mod command;
pub mod consts;
pub mod global_json;
pub mod json;
pub mod state;

use command::*;
use consts::*;
use pathfinding::prelude::*;
use state::*;
use std::cmp::Ordering;

pub fn choose_command(round: usize, state: &GameState) -> Command {
    if round <= 100 {
        choose_command_with_looking_forward_heuristic(state)
    } else {
        choose_command_with_astar(state)
    }
}

#[derive(Debug, Clone)]
struct HeuristicLookaheadState {
    moves: Vec<Command>,
    events: GameStateUpdateEvents,
    current_state: GameState,
}

fn choose_command_with_looking_forward_heuristic(state: &GameState) -> Command {
    let depth = 3;
    let mut states = vec![HeuristicLookaheadState {
        moves: Vec::new(),
        events: GameStateUpdateEvents::default(),
        current_state: state.clone(),
    }];

    for _ in 0..depth {
        states = states
            .into_iter()
            .flat_map(move |state| {
                state
                    .current_state
                    .good_moves(0)
                    .into_iter()
                    .map(move |player_move| {
                        let mut state = state.clone();
                        state.moves.push(player_move);
                        state
                            .current_state
                            .update([player_move, Command::Accelerate], &mut state.events);
                        state
                    })
            })
            .collect();
    }

    states
        .into_iter()
        .max_by(|a, b| compare_events(&a.events, &b.events))
        .unwrap()
        .moves
        .into_iter()
        .next()
        .unwrap()
}

fn compare_states(a: &GameState, b: &GameState) -> Ordering {
    if a.status == GameStatus::PlayerOneWon && b.status == GameStatus::PlayerOneWon {
        a.players[0].speed.cmp(&b.players[0].speed)
    } else if a.status == GameStatus::PlayerOneWon {
        Ordering::Greater
    } else if b.status == GameStatus::PlayerOneWon {
        Ordering::Less
    } else {
        let weighted_position_a = a.players[0].position.x + a.players[0].boosts * 2;
        let weighted_position_b = b.players[0].position.x + b.players[0].boosts * 2;

        weighted_position_a
            .cmp(&weighted_position_b)
            .then(a.players[0].speed.cmp(&b.players[0].speed))
            .then(a.players[0].position.x.cmp(&b.players[0].position.x))
            .then(a.players[0].boosts.cmp(&b.players[0].boosts))
    }
}

fn compare_events(a: &GameStateUpdateEvents, b: &GameStateUpdateEvents) -> Ordering {
    let a = &a.players[0];
    let b = &b.players[0];
    a.boosts_collected
        .cmp(&b.boosts_collected)
        .then(a.boosts_maintained.cmp(&b.boosts_maintained))
        .then(a.distance_travelled.cmp(&b.distance_travelled))
        .then(a.mud_hit.cmp(&b.mud_hit).reverse())
        .then(a.boosts_used.cmp(&b.boosts_used).reverse())
}

fn choose_command_with_astar(state: &GameState) -> Command {
    // TODO: Find all shortest paths, choose the one that has the
    // highest end speed, or otherwise wins
    shortest_path_first_command(state).unwrap_or(Command::Accelerate)
}

fn shortest_path_first_command(initial_state: &GameState) -> Option<Command> {
    let shortest_path_states = astar(
        initial_state,
        |state| {
            state
                .good_moves(0)
                .into_iter()
                .filter(|player_move| *player_move != Command::UseOil)
                .map(|player_move| {
                    let mut state = state.clone();
                    state.update(
                        [player_move, Command::Decelerate],
                        &mut GameStateUpdateEvents::default(),
                    );
                    (state, 1)
                })
                .collect::<Vec<_>>()
        },
        |state| (WIDTH - state.players[0].position.x) / SPEED_BOOST,
        |state| state.status != GameStatus::Continue,
    )
    .unwrap();

    shortest_path_states
        .0
        .iter()
        .zip(shortest_path_states.0.iter().skip(1))
        .map(|(state, next)| {
            let player_move = state
                .good_moves(0)
                .into_iter()
                .filter(|player_move| *player_move != Command::UseOil)
                .find(|player_move| {
                    let mut state = state.clone();
                    state.update(
                        [*player_move, Command::Decelerate],
                        &mut GameStateUpdateEvents::default(),
                    );
                    state == *next
                })
                .unwrap();

            player_move
        })
        .next()
}