use crate::{Tape, TapeEntry, gpu::Gpu};
use anyhow::{Result, ensure};

pub const MAX_INPUT_BYTES: usize = 256;

pub struct Parser {
    gpu: Gpu,
}

impl Parser {
    pub fn try_new() -> Result<Self> {
        let gpu = Gpu::try_new()?;

        Ok(Self { gpu })
    }

    pub const fn max_input_bytes(&self) -> usize {
        MAX_INPUT_BYTES
    }

    pub fn parse(&self, json: impl AsRef<[u8]>) -> Result<Tape> {
        self.parse_bytes(json.as_ref())
    }

    pub fn parse_str(&self, json: &str) -> Result<Tape> {
        self.parse_bytes(json.as_bytes())
    }

    pub fn parse_bytes(&self, json: &[u8]) -> Result<Tape> {
        ensure_input_fits(json)?;

        let mut input = json.to_vec();
        input.resize(MAX_INPUT_BYTES, 0);

        let input_buf = self.gpu.storage_buffer("bytes", &input);

        let fsm_buf = self
            .gpu
            .storage_buffer_empty("fsm", buffer_size(std::mem::size_of::<[u32; 4]>()));

        self.gpu.dispatch(
            include_str!("shaders/scan_fsm.wgsl"),
            "main",
            &[(&input_buf, true), (&fsm_buf, false)],
            1,
        );

        let compact_buf = self
            .gpu
            .storage_buffer_empty("compact", buffer_size(std::mem::size_of::<u32>()));

        let count_buf = self.gpu.storage_buffer_empty("count", count_buffer_size());
        self.gpu.dispatch(
            include_str!("shaders/scan_compact.wgsl"),
            "main",
            &[
                (&input_buf, true),
                (&fsm_buf, true),
                (&compact_buf, false),
                (&count_buf, false),
            ],
            1,
        );

        let depth_buf = self
            .gpu
            .storage_buffer_empty("depth", buffer_size(std::mem::size_of::<i32>()));
        self.gpu.dispatch(
            include_str!("shaders/scan_depth.wgsl"),
            "main",
            &[
                (&input_buf, true),
                (&compact_buf, true),
                (&depth_buf, false),
            ],
            1,
        );

        let parent_buf = self
            .gpu
            .storage_buffer_empty("parents", buffer_size(std::mem::size_of::<i32>()));
        self.gpu.dispatch(
            include_str!("shaders/parent_link.wgsl"),
            "main",
            &[
                (&input_buf, true),
                (&compact_buf, true),
                (&depth_buf, true),
                (&parent_buf, false),
            ],
            1,
        );

        let tape_buf = self
            .gpu
            .storage_buffer_empty("tape", buffer_size(std::mem::size_of::<TapeEntry>()));
        self.gpu.dispatch(
            include_str!("shaders/assemble_tape.wgsl"),
            "main",
            &[
                (&input_buf, true),
                (&compact_buf, true),
                (&depth_buf, true),
                (&parent_buf, true),
                (&tape_buf, false),
                (&fsm_buf, true),
            ],
            1,
        );

        let structure_count = usize::try_from(self.gpu.read_buffer_as::<u32>(&count_buf)[0])
            .expect("structure count must fit in usize");
        let tape = self.gpu.read_buffer_as::<TapeEntry>(&tape_buf);

        Ok(Tape::new(tape[..structure_count].to_vec()))
    }
}

fn ensure_input_fits(json: &[u8]) -> Result<()> {
    ensure!(json.len() <= MAX_INPUT_BYTES, "input is too large");

    Ok(())
}

fn buffer_size(element_size: usize) -> u64 {
    u64::try_from(MAX_INPUT_BYTES * element_size).expect("buffer size must fit in u64")
}

fn count_buffer_size() -> u64 {
    u64::try_from(std::mem::size_of::<u32>()).expect("count buffer size must fit in u64")
}

#[cfg(test)]
mod tests {
    use crate::Parser;

    #[test]
    fn test_basic_json() {
        let Ok(parser) = Parser::try_new() else {
            return;
        };

        let tape = parser.parse_str(r#"{"foo":"bar"}"#).unwrap();
        let entries = tape.iter().copied().collect::<Vec<_>>();

        insta::assert_debug_snapshot!(&entries, @"
        [
            TapeEntry {
                byte_pos: 0,
                byte_end: 1,
                depth: 1,
                parent: -1,
                token_kind: LeftBrace,
            },
            TapeEntry {
                byte_pos: 1,
                byte_end: 6,
                depth: 1,
                parent: 0,
                token_kind: String,
            },
            TapeEntry {
                byte_pos: 6,
                byte_end: 7,
                depth: 1,
                parent: 0,
                token_kind: Colon,
            },
            TapeEntry {
                byte_pos: 7,
                byte_end: 12,
                depth: 1,
                parent: 0,
                token_kind: String,
            },
            TapeEntry {
                byte_pos: 12,
                byte_end: 13,
                depth: 0,
                parent: 0,
                token_kind: RightBrace,
            },
        ]
        ");
    }
}
