use palette::{IntoColor, OklabHue, Oklch, Srgb};
use std::{
    collections::{BinaryHeap, HashMap, HashSet, VecDeque},
    fmt::Display,
    hash::{DefaultHasher, Hash, Hasher},
    io::{self, Write, stdin, stdout},
    num::{self, NonZero},
    ops::ControlFlow,
};

const CLEAR_SCREEN: &str = "\u{1b}c";
const RESET: &str = "\u{1b}[0m";
const BLOCK: char = '█';

fn crate_cli_color((r, g, b): (u8, u8, u8)) -> String {
    format!("\u{1b}[38;2;{r};{g};{b}m")
}

#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub struct Ring(NonZero<usize>);

impl Display for Ring {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut initial_hue = OklabHue::new(0.0);
        let mut initial_chroma = 0.5;

        let num = self.0.get() - 1;
        for _ in 0..num {
            initial_hue += 80.0;
            initial_chroma += 0.8;
            if initial_chroma > 1.0 {
                initial_chroma -= 1.0;
            }
        }

        let color = Oklch::new(0.7, initial_chroma, initial_hue);
        let color: Srgb = color.into_color();
        let color: Srgb<u8> = color.into_format();

        if f.alternate() {
            write!(
                f,
                "{}{BLOCK}{:^5}{BLOCK}{RESET}",
                crate_cli_color((color.red, color.green, color.blue)),
                self.0
            )
        } else {
            write!(
                f,
                "{}{BLOCK}{BLOCK}{BLOCK}{BLOCK}{BLOCK}{BLOCK}{BLOCK}{RESET}",
                crate_cli_color((color.red, color.green, color.blue))
            )
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Tower<const TOWER_HEIGHT: usize = 4> {
    pub rings: [Option<Ring>; TOWER_HEIGHT],
}

macro_rules! tower {
    ($($rings: literal),*) => {
        Tower::new(&[$(Ring(std::num::NonZero::new($rings).unwrap())),*])
    };
}

impl<const TOWER_HEIGHT: usize> Tower<TOWER_HEIGHT> {
    const EMPTY: Self = Self {
        rings: [None; TOWER_HEIGHT],
    };

    pub fn new(rings: &[Ring]) -> Self {
        let mut tower = Self::EMPTY;

        for (src, dst) in rings.into_iter().zip(&mut tower.rings) {
            *dst = Some(*src);
        }

        tower
    }

    fn burried_score(&self) -> usize {
        let mut ring_type = None;
        let mut same_ring_type_score = 0;
        let mut max_same_ring_type_score = 0;
        let mut burried_score = 0;

        for i in self.rings.iter().rev() {
            if let Some(ring) = i {
                if let Some(existing_ring_type) = ring_type {
                    if existing_ring_type == *ring {
                        same_ring_type_score += 2;
                    } else {
                        max_same_ring_type_score =
                            max_same_ring_type_score.max(same_ring_type_score);
                        burried_score += 2usize.pow(max_same_ring_type_score);
                        same_ring_type_score = 0;
                        ring_type = Some(*ring);
                    }
                } else {
                    same_ring_type_score = 0;
                    ring_type = Some(*ring);
                }
            }
        }

        burried_score
    }

    fn get_movable_top(&self) -> Result<(Ring, usize), MoveError> {
        let mut topmost_ring = None;
        let mut count = 0;

        for i in self.rings.iter() {
            match (i, &mut topmost_ring) {
                (None, _) => break,
                (Some(i), topmost @ None) => {
                    count = 1;
                    *topmost = Some(*i);
                }
                (Some(i), Some(topmost)) if i == topmost => count += 1,
                (Some(i), Some(topmost)) => {
                    count = 1;
                    *topmost = *i;
                }
            }
        }

        topmost_ring
            .map(|ring| {
                assert_ne!(count, 0);
                (ring, count)
            })
            .ok_or(MoveError::SourceEmpty)
    }

    fn remove_movable_top(&mut self, (ring, mut num): (Ring, usize)) {
        for i in self.rings.iter_mut().rev() {
            if num == 0 {
                return;
            }

            if let Some(i) = i.take() {
                assert_eq!(i, ring);
                num -= 1;
            }
        }

        assert_eq!(num, 0);
    }

    fn add_movable_top(
        &mut self,
        (ring, num): (Ring, usize),
        dummy: bool,
    ) -> Result<usize, MoveError> {
        if self.top_ring().is_some_and(|top| top != ring) {
            return Err(MoveError::DestinationWrongTop);
        }

        let height_left = TOWER_HEIGHT - self.height();
        let num_rings_moved = num.min(height_left);

        if dummy {
            return Ok(num_rings_moved);
        }

        self.rings
            .iter_mut()
            .filter(|i| i.is_none())
            .zip(std::iter::repeat_n(ring, num_rings_moved))
            .for_each(|(dst, src)| *dst = Some(src));

        self.check();

        Ok(num_rings_moved)
    }

    fn top_ring(&self) -> Option<Ring> {
        self.rings.iter().rev().filter_map(|i| *i).next()
    }

    fn height(&self) -> usize {
        self.rings
            .iter()
            .enumerate()
            .find(|(_, r)| r.is_none())
            .map(|(pos, _)| pos)
            .unwrap_or(TOWER_HEIGHT)
    }

    fn check(&self) {
        let mut has_noned = false;
        for i in &self.rings {
            if i.is_none() {
                has_noned = true;
            } else if has_noned {
                panic!("tower has some after none");
            }
        }
    }

    fn count_if_single_ring_type(&self) -> Option<(Option<Ring>, usize)> {
        let mut ring_type = None;
        let mut count = 0;

        for i in &self.rings {
            if let Some(ring) = i {
                if let Some(existing_ring_type) = ring_type {
                    if existing_ring_type == *ring {
                        count += 1;
                    } else {
                        return None;
                    }
                } else {
                    count = 1;
                    ring_type = Some(*ring);
                }
            }
        }

        Some((ring_type, count))
    }
}

pub struct Game<const TOWER_HEIGHT: usize = 4> {
    pub towers: Vec<Tower<TOWER_HEIGHT>>,
    pub ring_sets: HashMap<Ring, usize>,
}

#[derive(Clone, Debug)]
pub enum MoveError {
    DestinationWrongTop,
    SourceEmpty,
}

#[derive(Debug, Clone, Copy)]
pub struct Move {
    pub from_tower: usize,
    pub to_tower: usize,
}

impl Move {
    pub fn new(from_tower: usize, to_tower: usize) -> Self {
        Self {
            from_tower,
            to_tower,
        }
    }
}

impl<const TOWER_HEIGHT: usize> Game<TOWER_HEIGHT> {
    pub fn new(towers: Vec<Tower<TOWER_HEIGHT>>) -> Self {
        let mut ring_sets = HashMap::new();
        for t in &towers {
            t.check();
            for r in &t.rings {
                if let Some(r) = r {
                    *ring_sets.entry(*r).or_insert(0) += 1;
                }
            }
        }

        for (ring, _num) in &ring_sets {
            // ok as long as the goal isn't sets of 4 but single-ringtype towers
            // if num != TOWER_HEIGHT {
            // panic!("incomplete ring set: expected {TOWER_HEIGHT} of {ring} but found {num}");
            // }

            if ring.0.get() > towers.len() {
                panic!(
                    "ring set identifier too high: found ring with id {} but there are only {} towers so the largest expected ring id is {1}",
                    ring.0,
                    towers.len()
                )
            }
        }

        Self { ring_sets, towers }
    }

    pub fn possible_moves(&mut self) -> Vec<Move> {
        let mut res = Vec::new();

        for from_tower in 0..self.towers.len() {
            for to_tower in 0..self.towers.len() {
                if from_tower == to_tower {
                    continue;
                }

                let m = Move::new(from_tower, to_tower);
                if self.try_make_move(m, true).is_ok() {
                    res.push(m);
                }
            }
        }

        res
    }

    fn burried_ring_score(&self) -> usize {
        self.towers.iter().map(|i| i.burried_score()).sum()
    }

    pub fn solve(&self) -> Option<Vec<Move>> {
        let mut working_instance = Self {
            towers: Vec::new(),
            ring_sets: self.ring_sets.clone(),
        };

        #[derive(Clone)]
        struct State<const TOWER_HEIGHT: usize> {
            towers: Vec<Tower<TOWER_HEIGHT>>,
            m: Move,
            parent_index: usize,
        }
        impl<const TOWER_HEIGHT: usize> Hash for State<TOWER_HEIGHT> {
            fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
                self.towers.hash(state);
            }
        }

        let hash_state = |s: &[Tower<TOWER_HEIGHT>]| {
            let mut hasher = DefaultHasher::default();

            for x in s.iter().filter(|i| i.height() != 0).enumerate() {
                x.hash(&mut hasher);
            }

            hasher.finish()
        };

        #[derive(PartialEq, PartialOrd)]
        struct Ordf32(f32);
        impl Eq for Ordf32 {}
        impl Ord for Ordf32 {
            fn cmp(&self, other: &Self) -> std::cmp::Ordering {
                self.0.total_cmp(&other.0)
            }
        }

        let mut states = Vec::new();
        let mut todo = BinaryHeap::<(Ordf32, usize, usize)>::new();
        let mut had = HashSet::new();
        let mut ctr = 0;

        states.push(State {
            towers: self.towers.clone(),
            parent_index: 0,
            m: Move::new(0, 0),
        });
        todo.push((Ordf32(0.0), 1000000000, 0));

        while let Some((score, depth, state_index)) = todo.pop() {
            ctr += 1;
            let state = states[state_index].clone();

            let hash = hash_state(&state.towers);

            // already had
            if had.contains(&hash) {
                continue;
            }
            had.insert(hash);

            working_instance.towers = state.towers.clone();

            if ctr % 100000 == 0 {
                println!("{working_instance}");
                println!("{} {depth}", score.0);
                stdout().flush().unwrap();
                stdin().read_line(&mut String::new()).unwrap();
            }
            if working_instance.check_done() {
                println!("done!");
                let mut res = Vec::new();

                let mut curr = state_index;
                while curr != 0 {
                    res.push(states[curr].m);
                    curr = states[curr].parent_index;
                }

                res.reverse();

                return Some(res);
            }

            for m in working_instance.possible_moves() {
                working_instance.towers = state.towers.clone();
                working_instance.make_move(m).unwrap();

                let num_solved = working_instance.num_solved();
                let brs = working_instance.burried_ring_score();

                let new_state_hash = hash_state(&working_instance.towers);
                if had.contains(&new_state_hash) {
                    continue;
                }

                let new_state_index = states.len();
                states.push(State {
                    towers: working_instance.towers.clone(),
                    parent_index: state_index,
                    m,
                });

                let score = (num_solved as f32 * 10.0) + (-(brs as f32) * 100.0);

                todo.push((Ordf32(score), depth - 1, new_state_index));
            }
        }

        None
    }

    pub fn try_make_move(&mut self, m: Move, dummy: bool) -> Result<usize, MoveError> {
        let Move {
            from_tower,
            to_tower,
        } = m;

        let (ring, mut num) = self.towers[from_tower].get_movable_top()?;
        num = self.towers[to_tower].add_movable_top((ring, num), dummy)?;

        if dummy {
            return Ok(num);
        }

        self.towers[from_tower].remove_movable_top((ring, num));

        Ok(num)
    }

    pub fn make_move(&mut self, m: Move) -> Result<usize, MoveError> {
        self.try_make_move(m, false)
    }

    pub fn check_done(&self) -> bool {
        self.num_solved() == self.towers.len()
    }

    pub fn num_solved(&self) -> usize {
        let mut num_solved = 0;
        for t in &self.towers {
            t.check();

            // If the tower itself didn't have matching rings, return false
            let Some((ring_type, c)) = t.count_if_single_ring_type() else {
                continue;
            };

            // If the tower had no rings, trivially solved
            if c == 0 {
                num_solved += 1;
                continue;
            }

            // If the tower had only rings of one type,
            // see how many rings of that type it is supposed to have to be done.
            // if this tower has only a subset of the total, we're not done.
            let ring_type = ring_type.unwrap();
            if *self.ring_sets.get(&ring_type).unwrap() != c {
                continue;
            }

            num_solved += 1;
        }

        num_solved
    }

    fn ask_tower(&self, msg: &str) -> io::Result<Input> {
        let mut res = String::new();
        loop {
            print!("{msg} (enter a tower number, exit, clear or undo):");
            stdout().flush()?;
            stdin().read_line(&mut res)?;

            match res.trim() {
                "exit" => return Ok(Input::Exit),
                "clear" => return Ok(Input::Clear),
                "undo" => return Ok(Input::Undo),

                "e" => return Ok(Input::Exit),
                "c" => return Ok(Input::Clear),
                "u" => return Ok(Input::Undo),
                _ => {}
            }

            let mut parts = res.trim().split(' ').collect::<Vec<_>>();

            if parts.len() > 2 {
                println!("expected either one number or two numbers separated by a space");
                println!("ignoring all but the first");
                parts.truncate(1);
            }

            macro_rules! parse_part {
                ($e: expr) => {
                    match $e.parse::<usize>() {
                        Ok(i) if i > self.towers.len() => {
                            println!("tower index out of bounds");
                            continue;
                        }
                        Ok(i) if i == 0 => {
                            println!("(towers are 1-indexed)");
                            continue;
                        }
                        Ok(i) => i - 1,
                        Err(e) => {
                            println!(
                                "couldn't parse your input as an integer: {e}. please try again."
                            );
                            continue;
                        }
                    }
                };
            }

            match parts.len() {
                0 => {
                    println!("couldln't parse input");
                    continue;
                }
                1 => {
                    return Ok(Input::Tower(parse_part!(parts[0])));
                }
                2 => {
                    return Ok(Input::Done(parse_part!(parts[0]), parse_part!(parts[1])));
                }
                _ => unreachable!(),
            }
        }
    }

    fn report_move_error(&self, error: MoveError) {
        match error {
            MoveError::DestinationWrongTop => println!("can't place stack at destination"),
            MoveError::SourceEmpty => println!("can't move from this tower: it is empty"),
        }
    }

    fn ask_move(&self) -> io::Result<ControlFlow<(), Action>> {
        loop {
            println!("{self}");

            macro_rules! validate {
                ($e: expr) => {
                    match $e? {
                        Input::Exit => return Ok(ControlFlow::Break(())),
                        Input::Clear => {
                            println!("{CLEAR_SCREEN}");
                            continue;
                        }
                        Input::Tower(t) => t,
                        Input::Undo => return Ok(ControlFlow::Continue(Action::Undo)),
                        Input::Done(from_tower, to_tower) => {
                            return Ok(ControlFlow::Continue(Action::Move(Move {
                                from_tower,
                                to_tower,
                            })))
                        }
                    }
                };
            }

            let from_tower = validate!(self.ask_tower("move from"));

            let i = match self.towers[from_tower].get_movable_top() {
                Ok((ring, num)) => format!("{num} {ring}s"),
                Err(e) => {
                    self.report_move_error(e);
                    continue;
                }
            };

            let to_tower = validate!(self.ask_tower(&format!("move {i} to")));

            return Ok(ControlFlow::Continue(Action::Move(Move {
                from_tower,
                to_tower,
            })));
        }
    }

    pub fn cli_move(&mut self, undo: &mut UndoStack<TOWER_HEIGHT>) -> io::Result<ControlFlow<()>> {
        let ControlFlow::Continue(action) = self.ask_move()? else {
            return Ok(ControlFlow::Break(()));
        };
        println!("{CLEAR_SCREEN}");

        let m = match action {
            Action::Move(m) => m,
            Action::Undo => {
                if !undo.undo(self) {
                    println!("no more moves to undo");
                }
                return Ok(ControlFlow::Continue(()));
            }
        };

        let expected_num_moves = self.towers[m.from_tower]
            .get_movable_top()
            .map(|i| i.1)
            .unwrap_or(0);

        undo.save(self);

        match self.make_move(m) {
            Ok(actual_moves) => {
                if actual_moves != expected_num_moves {
                    println!("only moved {actual_moves} out of {expected_num_moves}");
                }
            }
            Err(e) => {
                undo.pop();
                self.report_move_error(e);
            }
        }

        if self.check_done() {
            println!("{CLEAR_SCREEN}{self}");
            println!("yayy! you did it :3");
            Ok(ControlFlow::Break(()))
        } else {
            Ok(ControlFlow::Continue(()))
        }
    }
}

enum Action {
    Move(Move),
    Undo,
}

enum Input {
    Done(usize, usize),
    Tower(usize),
    Exit,
    Clear,
    Undo,
}

pub struct GameGenerator {
    pub num_extra_towers: usize,
    pub num_ring_types: usize,
}

impl GameGenerator {
    pub fn generate<const TOWER_HEIGHT: usize>(self) -> Game<TOWER_HEIGHT> {
        let mut towers = Vec::new();
        let mut rings = HashSet::new();

        for ring_type in 0..self.num_ring_types {
            for idx in 0..TOWER_HEIGHT {
                rings.insert((Ring(NonZero::new(ring_type + 1).unwrap()), idx));
            }
        }

        let rings = rings.into_iter().map(|(i, _)| i).collect::<Vec<_>>();

        for i in rings.chunks(TOWER_HEIGHT) {
            towers.push(Tower::new(i));
        }

        for _ in 0..self.num_extra_towers {
            towers.push(Tower::EMPTY);
        }

        Game::new(towers)
    }
}

impl<const TOWER_HEIGHT: usize> Display for Game<TOWER_HEIGHT> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for i in (0..TOWER_HEIGHT).rev() {
            for m in 0..3 {
                for t in &self.towers {
                    if let Some(r) = t.rings[i] {
                        if m == 1 {
                            write!(f, "   {r:#}   ")?;
                        } else {
                            write!(f, "   {r:}   ")?;
                        }
                    } else {
                        write!(f, "             ")?;
                    }
                }
                writeln!(f)?;
            }
            if i != 0 {
                for _ in 0..self.towers.len() {
                    write!(f, "{:━^13}", "")?;
                }
                writeln!(f)?;
            }
        }

        for t in 0..self.towers.len() {
            let num = format!(" {} ", t + 1);
            write!(f, "{num:━^13}")?;
        }
        writeln!(f)?;
        for t in 0..self.towers.len() {
            let num = format!(" {} ", self.towers[t].burried_score());
            write!(f, "{num:^13}")?;
        }

        Ok(())
    }
}

pub struct UndoStack<const TOWER_HEIGHT: usize> {
    states: Vec<Vec<Tower<TOWER_HEIGHT>>>,
}

impl<const TOWER_HEIGHT: usize> UndoStack<TOWER_HEIGHT> {
    pub fn new() -> Self {
        Self { states: Vec::new() }
    }

    pub fn save(&mut self, game: &Game<TOWER_HEIGHT>) {
        self.states.push(game.towers.clone())
    }

    pub fn pop(&mut self) {
        self.states.pop();
    }

    pub fn undo(&mut self, game: &mut Game<TOWER_HEIGHT>) -> bool {
        if let Some(s) = self.states.pop() {
            game.towers = s;
            true
        } else {
            false
        }
    }
}

fn main() -> io::Result<()> {
    let mut g = GameGenerator {
        num_extra_towers: 2,
        num_ring_types: 12,
    }
    .generate::<4>();

    if let Some(solution) = g.solve() {
        let mut g = Game {
            towers: g.towers.clone(),
            ring_sets: g.ring_sets.clone(),
        };
        for i in solution {
            println!("{CLEAR_SCREEN}");
            g.make_move(i).unwrap();
            println!("{g}");
            stdout().flush().unwrap();
            stdin().read_line(&mut String::new()).unwrap();
        }
    }

    // println!("{CLEAR_SCREEN}");
    // let mut u = UndoStack::new();
    // while let ControlFlow::Continue(()) = g.cli_move(&mut u)? {}

    Ok(())
}

#[cfg(test)]
mod tests {
    use crate::{Game, Move, Ring, Tower};

    #[test]
    fn test_moves() {
        let mut g = Game::<4>::new(vec![tower![1, 2, 1, 2], Tower::EMPTY, Tower::EMPTY]);

        assert!(!g.check_done());

        // [1 2 1 2] [] [] -> [1 2 1 _] [2 _ _ _] []
        g.make_move(Move::new(0, 1)).unwrap();

        assert!(g.make_move(Move::new(0, 1)).is_err());

        // [1 2 1 _] [2 _ _ _] [] -> [1 2 _ _] [2 _ _ _] [1 _ _ _]
        g.make_move(Move::new(0, 2)).unwrap();

        assert!(!g.check_done());
        assert!(g.make_move(Move::new(0, 2)).is_err());

        // [1 2 _ _] [2 _ _ _] [1 _ _ _] -> [1 _ _ _] [2 2 _ _] [1 _ _ _]
        g.make_move(Move::new(0, 1)).unwrap();

        // still not done! 1s are split up
        assert!(!g.check_done());

        // [1 _ _ _] [2 2 _ _] [1 _ _ _] -> [] [2 2 _ _] [1 1 _ _]
        g.make_move(Move::new(0, 2)).unwrap();

        // all sorted
        assert!(g.check_done());
    }
}