use enum_map::{Enum, EnumMap};
use fastrand::Rng;

use crate::stack::Stack;


#[derive(Debug, Copy, Clone, Eq, PartialEq, Default, Enum)]
pub enum Color {
    #[default] Red, Green, Blue, Yellow, Purple,
}

// const COLORS: Stack<Color, 5> = Stack::from_array([
//     Color::Red,
//     Color::Green,
//     Color::Blue,
//     Color::Yellow,
//     Color::Purple,
// ]);

const COLORS: [Color; 5] = [
    Color::Red,
    Color::Green,
    Color::Blue,
    Color::Yellow,
    Color::Purple,
];


type ColorStack = Stack<Color, 5>;


#[derive(Debug, Copy, Clone)]
pub enum Tile {
    Forward,
    Backward,
}


#[derive(Debug, Copy, Clone)]
pub enum Square {
    Camels(ColorStack),
    Tile(Tile),
}

impl Square {
    fn assume_stack(&self) -> &ColorStack {
        match self {
            Square::Camels(stack) => stack,
            _ => panic!("Attempted to use the stack from a non-stack square"),
        }
    }

    fn assume_stack_mut(&mut self) -> &mut ColorStack {
        match self {
            Square::Camels(stack) => stack,
            _ => panic!("Attempted to use the stack from a non-stack square"),
        }
    }
}

impl Default for Square {
    fn default() -> Self {
        Square::Camels(ColorStack::new())
    }
}


#[derive(Debug, Default, Copy, Clone)]
pub struct Game {
    squares: [Square; 16],
    dice: EnumMap<Color, bool>,
    camels: EnumMap<Color, usize>,
}

impl Game {
    pub fn new() -> Self {
        Self::default()
    }

    // new game with random starting positions
    pub fn new_random() -> Self {
        let mut game = Self::default();
        let rng = Rng::new();
        let mut dice = *&COLORS;
        rng.shuffle(&mut dice);
        for color in dice {
            let roll = rng.usize(1..=3);
            game.squares[roll - 1].assume_stack_mut().push(color);
            game.camels[color] = roll - 1;
        }
        game
    }

    pub fn set_state(&mut self, camels: &[(Color, usize); 5], dice: &EnumMap<Color, bool>) {
        for i in 0..16 {
            self.squares[i] = match self.squares[i] {
                Square::Camels(mut stack) => {
                    stack.clear();
                    Square::Camels(stack)
                },
                _ => Square::Camels(Stack::new())
            };
        }

        for square in self.squares {
            assert_eq!(square.assume_stack().len(), 0)
        }

        self.dice = *dice;
        for &(color, sq) in camels {
            self.squares[sq].assume_stack_mut().push(color);
            self.camels[color] = sq;
        }
    }

    pub fn get_state(&self) -> ([(Color, usize); 5], EnumMap<Color, bool>) {
        let mut state = [(Color::Red, 0); 5];

        let mut j = 0;
        for (sq_idx, square) in self.squares.iter().enumerate() {
            if let Square::Camels(stack) = square {
                for camel in stack.iter() {
                    state[j] = (*camel, sq_idx);
                    j += 1;
                }
            }
        }

        (state, self.dice)
    }

    // returns winner if there is one
    pub fn advance(&mut self, die: Color, roll: usize) -> Option<Color> {
        let src_sq = self.camels[die];
        let dst_sq = src_sq + roll;
        if dst_sq >= 16 {
            self.dice[die] = true;
            return self.squares[src_sq].assume_stack().last().copied();
        }

        // special case when the destination square is the same as the source square
        if let Square::Tile(Tile::Backward) = self.squares[dst_sq] {
            if roll == 1 {
                let src_stack = self.squares[src_sq].assume_stack_mut();
                let slice_start = src_stack.iter().position(|&c| c == die).unwrap();
                src_stack.shift_slice_under(slice_start);
            }
        }
        else {
            // we have to split self.squares into two slices using split_at_mut, otherwise
            // rustc complains that we're trying to use two mutable references to the same value
            let (left, right) = self.squares.split_at_mut(src_sq + 1);
            let src_stack = left[src_sq].assume_stack_mut();
            let slice_start = src_stack.iter().position(|&c| c == die).unwrap();

            // since `right` starts immediately after the source square, the index of the 
            // destination square will be roll - 1 (e.g. if roll is 1, dst will be right[0])
            let (dst_rel_idx, prepend) = match right[roll - 1] {
                Square::Tile(Tile::Forward) => (roll, false), // roll - 1 + 1
                Square::Tile(Tile::Backward) => (roll - 2, true), // roll is guaranteed to be >= 2 since we already handled roll == 1
                _ => (roll - 1, false),
            };
            let dst_stack = right[dst_rel_idx].assume_stack_mut();
            let dst_true_idx = src_sq + 1 + dst_rel_idx; // src_sq + 1 was the original split boundary, so add the relative index to that to get the true index

            if prepend {
                let slice_len = src_stack.len() - slice_start;
                src_stack.move_slice_under(dst_stack, slice_start);
                for i in 0..slice_len {
                    self.camels[dst_stack[i]] = dst_true_idx;
                }
            }
            else {
                let dst_prev_len = dst_stack.len();
                src_stack.move_slice(dst_stack, slice_start);
                for i in dst_prev_len..dst_stack.len() {
                    self.camels[dst_stack[i]] = dst_true_idx;
                }
            }
        }

        self.dice[die] = true;
        None
    }

    fn finish_leg_random(&mut self, rng: &Rng) -> Option<Color> {
        let mut leg_dice: Stack<Color, 5> = Stack::new();
        for (color, rolled) in self.dice {
            if !rolled {
                leg_dice.push(color);
            }
        }

        rng.shuffle(&mut leg_dice[..]);
        for color in leg_dice.iter() {
            let roll = rng.usize(1..=3);
            if let Some(winner) = self.advance(*color, roll) {
                return Some(winner);
            }
        }
        None
    }

    fn finish_game_random(&mut self, rng: &Rng) -> Color {
        if let Some(winner) = self.finish_leg_random(rng) {
            return winner;
        }

        let mut dice = COLORS; // makes a copy of the constant
        // we are now guaranteed to be at the start of a new leg,
        // so we don't need to check the dice state
        loop {
            // easiest if we shuffle at the start of the leg
            rng.shuffle(&mut dice);
            for i in 0..5 {
                let roll = rng.usize(1..=3);
                if let Some(winner) = self.advance(dice[i], roll) {
                    return winner;
                }
            }
        }
    }

    pub fn project_outcomes(&self, count: usize) -> EnumMap<Color, usize> {
        let (orig_camels, orig_dice) = self.get_state();
        let mut projection = *self;

        let mut scores: EnumMap<Color, usize> = EnumMap::default();
        let rng = Rng::new();
        for i in 0..count {
            let winner = projection.finish_game_random(&rng);
            scores[winner] += 1;
            projection.set_state(&orig_camels, &orig_dice);
        }

        scores
    }
}


#[cfg(test)]
mod test {
    use super::*;
    use Color::*;

    #[test]
    fn test_advance() {
        let mut game = Game::new();
        // all dice are false (not rolled) to start with
        assert_eq!(game.dice.values().any(|&v| v), false);

        let camel_state = [
            (Blue, 0),
            (Yellow, 0),
            (Red, 1),
            (Green, 2),
            (Purple, 2),
        ];
        game.set_state(&camel_state, &Default::default());
        assert_eq!(game.squares[0].assume_stack(), &Stack::from([Blue, Yellow]));
        assert_eq!(game.camels[Blue], 0);
        assert_eq!(game.camels[Yellow], 0);
        assert_eq!(game.squares[1].assume_stack(), &Stack::from([Red]));
        assert_eq!(game.camels[Red], 1);
        assert_eq!(game.squares[2].assume_stack(), &Stack::from([Green, Purple]));
        assert_eq!(game.camels[Green], 2);
        assert_eq!(game.camels[Purple], 2);
        //  BY, R, GP

        game.advance(Yellow, 2);
        assert_eq!(game.dice[Yellow], true);
        assert_eq!(game.camels[Yellow], 2);
        assert_eq!(game.squares[2].assume_stack(), &Stack::from([Green, Purple, Yellow]));
        // B, R, GPY

        game.advance(Red, 2);
        assert_eq!(game.dice[Red], true);
        assert_eq!(game.camels[Red], 3);
        // B, _, GPY, R
        
        game.advance(Purple, 1);
        assert_eq!(game.dice[Purple], true);
        assert_eq!(game.squares[3].assume_stack(), &Stack::from([Red, Purple, Yellow]));
        // B, _, G, RPY
    }

    #[test]
    fn test_new_random() {
        for _ in 0..100 {
            let game = Game::new_random();
            for (camel, i) in game.camels {
                assert!(i < 3); // since we've only rolled the die once for each camel
                let stack = game.squares[i].assume_stack();
                assert!(stack[..].contains(&camel));
            }
        }
    }

    #[test]
    fn test_finish_leg() {
        let mut game = Game::new();
        let camel_state = [
            (Purple, 0),
            (Blue, 0),
            (Green, 1),
            (Red, 1),
            (Yellow, 2),
        ];
        game.set_state(&camel_state, &Default::default());
        // PB, G, RY
        game.advance(Green, 2);
        // PB, _, RY, G
        game.advance(Purple, 1);
        // _, PB, RY, G

        // since this is randomized, we should do it a bunch of times to make sure
        for _ in 0..100 {
            let mut projection = game; // copy?
            assert_eq!(projection.squares[1].assume_stack(), &Stack::from([Purple, Blue]));
            let rng = Rng::new();
            projection.finish_leg_random(&rng);
            // since we already rolled Green, it can't have moved
            assert_eq!(projection.camels[Green], 3);
            // likewise purple
            assert_eq!(projection.camels[Purple], 1);
            // blue, red,and yellow, on the other hand, *must* have moved
            assert_ne!(projection.camels[Blue], 1);
            assert_ne!(projection.camels[Red], 2);
            assert_ne!(projection.camels[Yellow], 2);
        }
    }

    #[test]
    fn test_finish_leg_winner() {
        let mut game = Game::new();
        let camel_state = [
            (Green, 13),
            (Red, 14),
            (Purple, 14),
            (Blue, 15),
            (Yellow, 15),
        ];
        game.set_state(&camel_state, &Default::default());

        // since there are no tiles involved, and multiple camels are on 15, there must be a winner
        let rng = Rng::new();
        assert!(matches!(game.finish_leg_random(&rng), Some(_)));
    }

    #[test]
    fn test_project_outcomes() {
        let mut game = Game::new();
        let camel_state = [
            (Blue, 1),
            (Green, 2),
            (Yellow, 2),
            (Purple, 4),
            (Red, 10),
        ];
        game.set_state(&camel_state, &Default::default());
        // _, B, GY, _, P, _, _, _, _, _, R

        let scores = game.project_outcomes(10_000);
        let mut max = 0;
        let mut winner = Blue; // just "anything that's not red"
        for (color, score) in scores {
            if score > max {
                max = score;
                winner = color;
            }
        }

        assert_eq!(winner, Red);
    }
}