use anyhow::{Context as _, Result};
use std::collections::HashMap;
use std::path::Path;

pub fn compile_ruleset(in_path: impl AsRef<Path>, out_path: impl AsRef<Path>) -> Result<()> {
    let mut in_file = std::fs::File::open(in_path.as_ref())
        .with_context(|| format!("could not open input file {}", in_path.as_ref().display()))?;
    let mut out_file = std::fs::File::create(out_path.as_ref())
        .with_context(|| format!("could not open output file {}", out_path.as_ref().display()))?;
    let rules_data: RulesData =
        serde_json::from_reader(&mut in_file).context("error parsing input file")?;
    compile(&rules_data)
        .emit(&mut out_file)
        .context("error writing to output file")
}

/// Input data (JSON file)
#[derive(serde::Deserialize)]
struct RulesData {
    shapes: HashMap<String, Vec<(i16, i16)>>,
    kicks: HashMap<String, [Vec<(i16, i16)>; 4]>,
    spawn: (i16, i16),
}

/// Output data (Rust code)
struct RulesCode {
    shape_indices: HashMap<String, usize>,
    extents: Extents<Vec<i16>>,
    cells_bits: Vec<u16>,
    cells_indices: Vec<usize>,
    kicks_tests: Vec<(i16, i16)>,
    kicks_counts: Vec<usize>,
    kicks_indices: Vec<usize>,
    spawn: (i16, i16),
}

struct Extents<E> {
    x0: E,
    x1: E,
    y0: E,
    y1: E,
}

fn compile(rules: &RulesData) -> RulesCode {
    let mut shape_names: Vec<&str> = rules.shapes.keys().map(String::as_str).collect();
    // ensure deterministic order
    shape_names.sort();

    // `kicks` is keyed by strings indicating multiple shape names, e.g. "LJSTZ" means
    // that this entry applies to each of those types. split the shape names into entries
    // per shape for easier lookup.
    let mut shapes_kicks: HashMap<String, [&[(i16, i16)]; 4]> = HashMap::new();
    for (shapes_string, kicks) in rules.kicks.iter() {
        for ch in shapes_string.chars() {
            let shape = std::iter::once(ch).collect::<String>();
            let kicks = std::array::from_fn(|i| &kicks[i][..]);
            shapes_kicks.insert(shape, kicks);
        }
    }

    let len_per_shape_rot = shape_names.len() * 4;
    let len_per_shape_kicks = shape_names.len() * 16;

    let mut extents = Extents {
        x0: vec![i16::MAX; len_per_shape_rot],
        x1: vec![i16::MIN; len_per_shape_rot],
        y0: vec![i16::MAX; len_per_shape_rot],
        y1: vec![i16::MIN; len_per_shape_rot],
    };

    let shape_indices: HashMap<String, usize> = shape_names
        .iter()
        .enumerate()
        .map(|(i, &name)| (name.to_owned(), i))
        .collect();

    let mut cells_bits = vec![];
    let mut cells_indices = vec![usize::MAX; len_per_shape_rot];

    let mut kicks_tests = vec![];
    let mut kicks_counts = vec![0; len_per_shape_kicks];
    let mut kicks_indices = vec![usize::MAX; len_per_shape_kicks];

    for &shape in shape_names.iter() {
        let i0 = shape_indices.get(shape).unwrap();
        let coords = rules.shapes.get(shape).unwrap();
        for r in 0u8..4 {
            let i = i0 * 4 + r as usize;
            let (exts, data) = process_rotated_coords(coords, r);
            extents.x0[i] = exts.x0;
            extents.x1[i] = exts.x1;
            extents.y0[i] = exts.y0;
            extents.y1[i] = exts.y1;
            cells_indices[i] = find_or_insert(&mut cells_bits, &data);
        }

        let kicks = shapes_kicks.get(shape).unwrap();
        for r0 in 0u8..4 {
            for r1 in 0u8..4 {
                let i = i0 * 16 + r0 as usize * 4 + r1 as usize;
                let tests = process_kicks(kicks, r0, r1);
                kicks_counts[i] = tests.len();
                kicks_indices[i] = find_or_insert(&mut kicks_tests, &tests);
            }
        }
    }

    RulesCode {
        shape_indices,
        extents,
        cells_bits,
        cells_indices,
        kicks_tests,
        kicks_counts,
        kicks_indices,
        spawn: rules.spawn,
    }
}

fn process_rotated_coords(coords: &[(i16, i16)], rot: u8) -> (Extents<i16>, Vec<u16>) {
    let mut exts = Extents {
        x0: i16::MAX,
        x1: i16::MIN,
        y0: i16::MAX,
        y1: i16::MIN,
    };
    for &(x0, y0) in coords.iter() {
        let (x, y) = rotate(x0, y0, rot);
        exts.x0 = std::cmp::min(exts.x0, x);
        exts.x1 = std::cmp::max(exts.x1, x + 1);
        exts.y0 = std::cmp::min(exts.y0, y);
        exts.y1 = std::cmp::max(exts.y1, y + 1);
    }

    let mut data = vec![0u16; (exts.y1 - exts.y0) as usize];
    for &(x0, y0) in coords.iter() {
        let (x, y) = rotate(x0, y0, rot);
        data[(y - exts.y0) as usize] |= 1 << (x - exts.x0);
    }
    (exts, data)
}

fn rotate(x: i16, y: i16, r: u8) -> (i16, i16) {
    let mut xy = (x, y);
    for _ in 0..r {
        let (x, y) = xy;
        xy = (y, -x);
    }
    xy
}

fn process_kicks(kicks: &[&[(i16, i16)]], r0: u8, r1: u8) -> Vec<(i16, i16)> {
    if r0 == r1 {
        return vec![];
    }
    let kicks0 = kicks[r0 as usize].iter();
    let kicks1 = kicks[r1 as usize].iter();
    kicks0
        .zip(kicks1)
        .map(|(&(x0, y0), &(x1, y1))| (x0 - x1, y0 - y1))
        .collect()
}

/// Finds the first occurence of `needle` in `haystack`, or appends a suffix of it (up to
/// the entire length), returning the index within `haystack` that now contains `needle`.
fn find_or_insert<T: Eq + Copy>(haystack: &mut Vec<T>, needle: &[T]) -> usize {
    for i in 0.. {
        let len = std::cmp::min(haystack.len() - i, needle.len());
        if haystack[i..i + len] == needle[..len] {
            haystack.extend_from_slice(&needle[len..]);
            return i;
        }
    }
    unreachable!()
}

impl RulesCode {
    fn emit(&self, out: &mut impl std::io::Write) -> std::io::Result<()> {
        writeln!(out, "pub mod shape_indices {{")?;
        for (name, &index) in self.shape_indices.iter() {
            writeln!(out, "pub const {name}: usize = {index:?};")?;
        }
        writeln!(out, "}}")?;

        writeln!(out, "pub mod extents {{")?;
        for (name, arr) in [
            ("X0", &self.extents.x0),
            ("X1", &self.extents.x1),
            ("Y0", &self.extents.y0),
            ("Y1", &self.extents.y1),
        ] {
            let len = arr.len();
            writeln!(out, "pub const {name}: [i16; {len}] = {arr:?};")?;
        }
        writeln!(out, "}}")?;

        let arr = &self.cells_bits;
        let len = self.cells_bits.len();
        writeln!(out, "pub const CELLS_BITS: [u16; {len}] = {arr:#x?};")?;

        let arr = &self.cells_indices;
        let len = self.cells_indices.len();
        writeln!(out, "pub const CELLS_INDEX: [usize; {len}] = {arr:?};")?;

        let arr = &self.kicks_tests;
        let len = self.kicks_tests.len();
        writeln!(out, "pub const KICKS_TESTS: [(i16, i16); {len}] = {arr:?};")?;

        let arr = &self.kicks_counts;
        let len = self.kicks_counts.len();
        writeln!(out, "pub const KICKS_COUNT: [usize; {len}] = {arr:?};")?;

        let arr = &self.kicks_indices;
        let len = self.kicks_indices.len();
        writeln!(out, "pub const KICKS_INDEX: [usize; {len}] = {arr:?};")?;

        let spawn = self.spawn;
        writeln!(out, "pub const SPAWN: (i16, i16) = {spawn:?};")?;

        Ok(())
    }
}