use std::iter::repeat_n;
use std::{collections::BinaryHeap, f64::consts, fmt::Display};

use aoc_runner_derive::{aoc, aoc_generator};
use grid::{Coord2d, Grid, MirrorAxis};
use itertools::{Format, Itertools};

#[derive(Debug, Clone)]
struct PresentShape {
    idx: u8,
    shape: Grid<u8>,
    variants: Vec<Grid<u8>>,
}

impl From<&str> for PresentShape {
    fn from(value: &str) -> Self {
        let (idx, rest) = value.split_once(":\n").unwrap();

        let shape: Grid<u8> = rest.parse().unwrap();

        Self {
            idx: idx.parse().unwrap(),
            variants: (0..4)
                .map(|rot| shape.rotated(rot))
                .chain([MirrorAxis::X, MirrorAxis::Y].map(|axis| shape.mirrored(axis)))
                .unique()
                .collect(),
            shape,
        }
    }
}

impl Display for PresentShape {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("{}:\n", self.idx))?;
        f.write_fmt(format_args!("{}", self.shape))
    }
}

impl PresentShape {
    fn make(&self) -> Present {
        Present {
            shape: self.clone(),
            location: None,
            variant: 0,
        }
    }
    fn popcount(&self) -> u64 {
        self.shape.data.iter().filter(|c| **c != b'.').count() as u64
    }
}

#[derive(Debug, Clone)]
struct Present {
    shape: PresentShape,
    location: Option<Coord2d>,
    variant: usize,
}

#[derive(Debug, Clone)]
struct ChristmasTree {
    area: Grid<u8>,
    presents: Vec<Present>,
}

impl ChristmasTree {}

impl Display for ChristmasTree {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("{}", self.area))?;
        for (i, p) in self.presents.iter().enumerate() {
            f.write_fmt(format_args!("  {i}: @{:?} rot {}\n", p.location, p.variant))?;
            f.write_fmt(format_args!("{}", p.shape.variants[p.variant]))?;
        }
        std::fmt::Result::Ok(())
    }
}

#[derive(Debug, Clone)]
struct PresentsProblem {
    shapes: Vec<PresentShape>,
    trees: Vec<ChristmasTree>,
}

impl PresentsProblem {
    fn present(&self, idx: usize) -> Present {
        Present {
            shape: self.shapes[idx].clone(),
            location: None,
            variant: 0,
        }
    }
    fn add_tree(&mut self, w: usize, h: usize, presents: Vec<Present>) {
        self.trees.push(ChristmasTree {
            area: Grid::with_shape(w, h, b'.'),
            presents,
        })
    }
}

#[aoc_generator(day12)]
fn parse(input: &str) -> PresentsProblem {
    let mut parts = input.split("\n\n").collect_vec();
    let trees_s = parts.pop().unwrap();

    let shapes: Vec<PresentShape> = parts.into_iter().map(|p| p.into()).collect();
    let mut trees = Vec::new();

    for l in trees_s.lines() {
        let (d, el) = l.split_once(": ").unwrap();
        let (w, h) = d.split_once('x').unwrap();
        let present_counts: Vec<usize> = el
            .split_ascii_whitespace()
            .map(|count| count.parse().unwrap())
            .collect_vec();
        let presents = present_counts
            .iter()
            .enumerate()
            .flat_map(|(i, count)| repeat_n(shapes[i].make(), *count))
            .collect();
        trees.push(ChristmasTree {
            area: Grid::with_shape(w.parse().unwrap(), h.parse().unwrap(), b'.'),
            presents,
        });
    }

    PresentsProblem { shapes, trees }
}

// place b on a and test if any # overlap any non-.
fn overlaps(a: &Grid<u8>, b: &Grid<u8>, p: &Coord2d) -> bool {
    b.find_all(&b'#')
        .any(|c| a.get(&(c + p)).is_some_and(|c| *c != b'.'))
}

fn place_shape(
    mut t: ChristmasTree,
    idx: usize,
    pos: Coord2d,
    variant: usize,
) -> Option<ChristmasTree> {
    let variant = &t.presents[idx].shape.variants[variant];

    if !overlaps(&t.area, variant, &pos) {
        for xy in variant.find_all(&b'#') {
            t.area.set(&(pos + &(xy)), b'0' + idx as u8).unwrap();
        }

        Some(t)
    } else {
        None
    }
}

fn place_next(mut t: ChristmasTree, cur: usize) -> Option<ChristmasTree> {
    if cur == t.presents.len() {
        return Some(t);
    }
    let variants = &t.presents[cur].shape.variants;
    'variant: for (i, _v) in variants.iter().enumerate() {
        //make sure the longest 'bar' can fit
        for y in 0..&t.area.height() - 2 {
            for x in 0..&t.area.width() - 2 {
                if let Some(new_t) = place_shape(
                    t.clone(),
                    cur,
                    Coord2d {
                        x: x as i64,
                        y: y as i64,
                    },
                    i,
                ) {
                    if let Some(solution) = place_next(new_t, cur + 1) {
                        return Some(solution);
                    }
                }
            }
        }
    }
    None
}

#[aoc(day12, part1)]
fn part1(input: &PresentsProblem) -> u64 {
    let input = input.clone();
    let mut count = 0;
    // for t in input.trees[0..2].iter() {
    //     println!("trying:\n{t}");
    //     if let Some(r) = place_next(t.clone(), 0) {
    //         count += 1;
    //         println!("ok!\n{r}");
    //     } else {
    //         println!("failed :(");
    //     }
    // }
    for t in input.trees.iter() {
        let area = (t.area.width() * t.area.height()) as u64;
        let occupied_area = t.presents.iter().map(|p| p.shape.popcount()).sum::<u64>();
        if occupied_area <= area {
            count += 1
        }
    }

    count
}

#[aoc(day12, part2)]
fn part2(input: &PresentsProblem) -> u64 {
    0
}

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

    const EXAMPLE: &str = "0:
###
##.
##.

1:
###
##.
.##

2:
.##
###
##.

3:
##.
###
##.

4:
###
#..
###

5:
###
.#.
###

4x4: 0 0 0 0 2 0
12x5: 1 0 1 0 2 2
12x5: 1 0 1 0 3 2";

    #[rstest]
    #[case(EXAMPLE, 0)]
    fn part1_example(#[case] input: &str, #[case] expected: u64) {
        assert_eq!(part1(&parse(EXAMPLE)), 2);
    }

    #[rstest]
    #[case(EXAMPLE, 0)]
    fn part2_example(#[case] input: &str, #[case] expected: u64) {
        assert_eq!(part2(&parse(EXAMPLE)), 0);
    }
}