use std::collections::{HashMap, HashSet};

use advent_of_code_2024::{make_main, next, Pair, SResult};

make_main!();

fn solve(lines: Vec<String>) -> SResult<(usize, usize)> {
    let width = lines.len();
    let mut grid: Vec<Vec<char>> = Vec::with_capacity(lines.len());
    for line in lines {
        grid.push(line.trim().chars().collect());
    }

    let mut visited: HashSet<Pair> = HashSet::new();
    let mut size = 0;
    let mut size2 = 0;
    for (i, row) in grid.iter().enumerate() {
        for (j, _) in row.iter().enumerate() {
            if visited.contains(&(i, j)) {
                continue;
            }
            let section = fill(&grid, (i, j), &mut visited);
            size += perim(width, &section) * section.len();
            size2 += sides(width, &section) * section.len();
        }
    }

    // Find the regions
    Ok((size, size2))
}

fn fill(grid: &[Vec<char>], loc: Pair, global_visit: &mut HashSet<Pair>) -> HashSet<Pair> {
    let mut stack = vec![loc];
    let mut visited = HashSet::new();
    visited.insert(loc);
    global_visit.insert(loc);
    let v = grid[loc.0][loc.1];

    while let Some(loc) = stack.pop() {
        for vel in [(0, 1), (0, -1), (-1, 0), (1, 0)] {
            if let Some(loc2) = next(loc, vel, grid.len()) {
                if grid[loc2.0][loc2.1] == v && !visited.contains(&loc2) {
                    global_visit.insert(loc2);
                    visited.insert(loc2);
                    stack.push(loc2);
                }
            }
        }
    }

    visited
}

fn perim(width: usize, section: &HashSet<Pair>) -> usize {
    let mut perim = 0;
    for loc in section.iter() {
        if loc.0 == 0 || loc.0 + 1 == width {
            perim += 1;
        }
        if loc.1 == 0 || loc.1 + 1 == width {
            perim += 1;
        }
        for vel in [(0, 1), (0, -1), (-1, 0), (1, 0)] {
            if let Some(loc2) = next(*loc, vel, width) {
                if section.contains(&loc2) {
                    continue;
                }
                perim += 1;
            }
        }
    }
    perim
}

fn sides(width: usize, section: &HashSet<Pair>) -> usize {
    let mut count = 0;

    // Organize things by x, and y
    let mut by_x: HashMap<usize, Vec<Pair>> = HashMap::default();
    let mut by_y: HashMap<usize, Vec<Pair>> = HashMap::default();
    for pos in section.iter() {
        by_x.entry(pos.0).or_insert_with(|| Vec::new()).push(*pos);
        by_y.entry(pos.1).or_insert_with(|| Vec::new()).push(*pos);
    }

    for (_, v) in by_x.iter_mut() {
        v.sort_by(|a, b| a.1.cmp(&b.1));
    }
    for (_, v) in by_y.iter_mut() {
        v.sort_by(|a, b| a.0.cmp(&b.0));
    }

    for (x, v) in by_x {
        // Check if we need fence on left/right
        for vel in [(1, 0), (-1, 0)] {
            let mut in_fence = false;
            for pos in &v {
                let mut need_fence = true;
                // We don't need a fence if the next cell over is in the region
                if let Some(x) = next(*pos, vel, width) {
                    if section.contains(&x) {
                        need_fence = false;
                    }
                }
                if in_fence != need_fence {
                    if need_fence {
                        count += 1;
                    }
                    in_fence = need_fence;
                }
                // If the next cell is not adjaent to us, we are not in fence
                // when the next iteration starts
                in_fence = in_fence && match next(*pos, (0, 1), width) {
                    Some(x) => section.contains(&x),
                    None => false,
                };
            }
        }
    }

    for (y, v) in by_y {
        for vel in [(0, 1), (0, -1)] {
            let mut in_fence = false;
            for pos in &v {
                let mut need_fence = true;
                // We don't need a fence if the next cell over is in the region
                if let Some(x) = next(*pos, vel, width) {
                    if section.contains(&x) {
                        need_fence = false;
                    }
                }
                if in_fence != need_fence {
                    if need_fence {
                        count += 1;
                    }
                    in_fence = need_fence;
                }
                // If the next cell is not adjaent to us, we are not in fence
                // when the next iteration starts
                in_fence = in_fence && match next(*pos, (1, 0), width) {
                    // We are still in the fance
                    Some(x) => section.contains(&x),
                    None => false,
                };
            }
        }
    }
    count
}

#[cfg(test)]
mod tests {
    use advent_of_code_2024::input;

    use super::*;
    #[test]
    fn sample_input() {
        let strings: Vec<String> = input!("d12p1.txt");
        let got = solve(strings).unwrap();
        assert_eq!(got, (1930, 1206));
    }

    #[test]
    fn sample_input2() {
        let strings: Vec<String> = input!("d12i1.txt");
        let got = solve(strings).unwrap();
        assert_eq!(got, (140, 80));
    }

    #[test]
    fn sample_input3() {
        let strings: Vec<String> = input!("d12i2.txt");
        let got = solve(strings).unwrap();
        assert_eq!(got, (772, 436));
    }

    #[test]
    fn sample_input4() {
        let strings: Vec<String> = input!("d12i3.txt");
        let got = solve(strings).unwrap();
        assert_eq!(got, (692, 236));
    }
}