use std::collections::HashMap;
use std::fs::File;
use std::io::{BufWriter, Read, Write};
use std::path::{Path, PathBuf};
use std::thread;

use crate::graph::Node;
use crate::index::MAX_ENTRY_POINTS;
use crate::{AnnConfig, AnnError, AnnIndex, Result};

const MAGIC: &[u8; 4] = b"FXA1";
const VERSION: u16 = 1;
const NO_ENTRY: u64 = u64::MAX;

impl AnnIndex {
    pub fn save_to(&self, path: impl AsRef<Path>) -> Result<()> {
        let file = File::create(path)?;
        let mut writer = BufWriter::new(file);
        self.write_to(&mut writer)?;
        writer.flush()?;
        writer.get_ref().sync_all()?;
        Ok(())
    }

    pub fn load_from(path: impl AsRef<Path>) -> Result<Self> {
        Self::read_from(File::open(path)?)
    }

    pub fn save_in_background(&self, path: PathBuf) -> thread::JoinHandle<Result<()>> {
        let snapshot = self.clone();
        thread::spawn(move || snapshot.save_to(path))
    }

    pub fn write_to<W: Write>(&self, writer: W) -> Result<()> {
        let mut writer = CrcWriter::new(writer);

        writer.write_all(MAGIC)?;
        write_u16(&mut writer, VERSION)?;
        write_u16(&mut writer, 0)?;
        write_u32(&mut writer, self.config.dimensions as u32)?;
        write_u32(&mut writer, self.config.max_neighbors as u32)?;
        write_u32(&mut writer, self.config.max_base_neighbors as u32)?;
        write_u32(&mut writer, self.config.ef_construction as u32)?;
        write_u32(&mut writer, self.config.ef_search as u32)?;
        write_u32(&mut writer, self.config.max_level as u32)?;
        write_u64(
            &mut writer,
            self.entry.map(|e| e as u64).unwrap_or(NO_ENTRY),
        )?;
        write_u32(&mut writer, self.max_level as u32)?;
        write_u64(&mut writer, self.active as u64)?;
        write_u64(&mut writer, self.nodes.len() as u64)?;
        write_u32(&mut writer, self.entry_points.len() as u32)?;
        for &entry_point in &self.entry_points {
            write_u64(&mut writer, entry_point as u64)?;
        }

        for node in &self.nodes {
            let id_bytes = node.id.as_bytes();
            if id_bytes.len() > u16::MAX as usize {
                return Err(AnnError::InvalidVector(
                    "record id exceeds 65535 bytes and cannot be serialized".to_string(),
                ));
            }
            write_u16(&mut writer, id_bytes.len() as u16)?;
            writer.write_all(id_bytes)?;
            writer.write_all(&[u8::from(node.deleted)])?;
            write_u16(&mut writer, node.neighbors.len() as u16)?;
            for value in &node.vector {
                writer.write_all(&value.to_le_bytes())?;
            }
            for layer in &node.neighbors {
                write_u32(&mut writer, layer.len() as u32)?;
                for &neighbor in layer {
                    write_u64(&mut writer, neighbor as u64)?;
                }
            }
        }

        let crc = writer.crc();
        writer.into_inner().write_all(&crc.to_le_bytes())?;
        Ok(())
    }

    pub fn read_from<R: Read>(mut reader: R) -> Result<Self> {
        let mut bytes = Vec::new();
        reader.read_to_end(&mut bytes)?;
        if bytes.len() < MAGIC.len() + 4 {
            return Err(corrupted("snapshot is too short"));
        }

        let (payload, crc_bytes) = bytes.split_at(bytes.len() - 4);
        let stored_crc = u32::from_le_bytes(crc_bytes.try_into().expect("crc is four bytes"));
        if crc32(payload) != stored_crc {
            return Err(corrupted("checksum mismatch"));
        }

        let mut cursor = Cursor::new(payload);
        if cursor.take_bytes(4)? != MAGIC {
            return Err(corrupted("bad magic"));
        }
        let version = cursor.take_u16()?;
        if version != VERSION {
            return Err(corrupted(&format!("unsupported version {version}")));
        }
        cursor.take_u16()?;

        let config = AnnConfig {
            dimensions: cursor.take_u32()? as usize,
            max_neighbors: cursor.take_u32()? as usize,
            max_base_neighbors: cursor.take_u32()? as usize,
            ef_construction: cursor.take_u32()? as usize,
            ef_search: cursor.take_u32()? as usize,
            max_level: cursor.take_u32()? as usize,
        };
        let entry_raw = cursor.take_u64()?;
        let max_level = cursor.take_u32()? as usize;
        let active = cursor.take_u64()? as usize;
        let node_count = cursor.take_u64()? as usize;
        let entry_points_len = cursor.take_u32()? as usize;
        if entry_points_len > MAX_ENTRY_POINTS {
            return Err(corrupted("entry point count exceeds maximum"));
        }
        let mut entry_points = Vec::with_capacity(entry_points_len);
        for _ in 0..entry_points_len {
            let entry_point = cursor.take_u64()? as usize;
            if entry_point >= node_count {
                return Err(corrupted("entry point out of bounds"));
            }
            entry_points.push(entry_point);
        }

        let mut index = AnnIndex::new(config)?;
        index.reserve(node_count);
        let mut ids: HashMap<String, usize> = HashMap::with_capacity(node_count);
        let mut visible = 0_usize;

        for position in 0..node_count {
            let id_len = cursor.take_u16()? as usize;
            let id = String::from_utf8(cursor.take_bytes(id_len)?.to_vec())
                .map_err(|_| corrupted("record id is not utf-8"))?;
            let deleted = match cursor.take_bytes(1)?[0] {
                0 => false,
                1 => true,
                _ => return Err(corrupted("invalid deleted flag")),
            };
            let levels = cursor.take_u16()? as usize;
            if levels == 0 || levels > index.config.max_level + 1 {
                return Err(corrupted("invalid level count"));
            }
            let mut vector = Vec::with_capacity(index.config.dimensions);
            for _ in 0..index.config.dimensions {
                vector.push(f32::from_le_bytes(
                    cursor.take_bytes(4)?.try_into().expect("f32 is four bytes"),
                ));
            }
            let mut neighbors = Vec::with_capacity(levels);
            for _ in 0..levels {
                let count = cursor.take_u32()? as usize;
                let mut layer = Vec::with_capacity(count);
                for _ in 0..count {
                    let neighbor = cursor.take_u64()? as usize;
                    if neighbor >= node_count {
                        return Err(corrupted("neighbor index out of bounds"));
                    }
                    layer.push(neighbor);
                }
                neighbors.push(layer);
            }

            if !deleted {
                if ids.insert(id.clone(), position).is_some() {
                    return Err(corrupted("duplicate visible record id"));
                }
                visible += 1;
            }
            index.nodes.push(Node {
                id,
                vector,
                deleted,
                neighbors,
            });
        }

        if !cursor.is_empty() {
            return Err(corrupted("trailing bytes after last node"));
        }
        if visible != active {
            return Err(corrupted("active count does not match visible records"));
        }

        index.ids = ids;
        index.active = active;
        index.max_level = max_level;
        index.entry_points = entry_points;
        index.entry = if entry_raw == NO_ENTRY {
            None
        } else {
            let entry = entry_raw as usize;
            if entry >= node_count {
                return Err(corrupted("entry index out of bounds"));
            }
            Some(entry)
        };
        Ok(index)
    }
}

fn corrupted(message: &str) -> AnnError {
    AnnError::Corrupted(message.to_string())
}

struct Cursor<'a> {
    bytes: &'a [u8],
}

impl<'a> Cursor<'a> {
    fn new(bytes: &'a [u8]) -> Self {
        Self { bytes }
    }

    fn take_bytes(&mut self, len: usize) -> Result<&'a [u8]> {
        if self.bytes.len() < len {
            return Err(corrupted("unexpected end of snapshot"));
        }
        let (taken, rest) = self.bytes.split_at(len);
        self.bytes = rest;
        Ok(taken)
    }

    fn take_u16(&mut self) -> Result<u16> {
        Ok(u16::from_le_bytes(
            self.take_bytes(2)?.try_into().expect("u16 is two bytes"),
        ))
    }

    fn take_u32(&mut self) -> Result<u32> {
        Ok(u32::from_le_bytes(
            self.take_bytes(4)?.try_into().expect("u32 is four bytes"),
        ))
    }

    fn take_u64(&mut self) -> Result<u64> {
        Ok(u64::from_le_bytes(
            self.take_bytes(8)?.try_into().expect("u64 is eight bytes"),
        ))
    }

    fn is_empty(&self) -> bool {
        self.bytes.is_empty()
    }
}

struct CrcWriter<W: Write> {
    inner: W,
    crc: u32,
}

impl<W: Write> CrcWriter<W> {
    fn new(inner: W) -> Self {
        Self {
            inner,
            crc: 0xffff_ffff,
        }
    }

    fn crc(&self) -> u32 {
        self.crc ^ 0xffff_ffff
    }

    fn into_inner(self) -> W {
        self.inner
    }
}

impl<W: Write> Write for CrcWriter<W> {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let written = self.inner.write(buf)?;
        self.crc = crc32_update(self.crc, &buf[..written]);
        Ok(written)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.inner.flush()
    }
}

fn write_u16<W: Write>(writer: &mut W, value: u16) -> Result<()> {
    writer.write_all(&value.to_le_bytes())?;
    Ok(())
}

fn write_u32<W: Write>(writer: &mut W, value: u32) -> Result<()> {
    writer.write_all(&value.to_le_bytes())?;
    Ok(())
}

fn write_u64<W: Write>(writer: &mut W, value: u64) -> Result<()> {
    writer.write_all(&value.to_le_bytes())?;
    Ok(())
}

fn crc32(bytes: &[u8]) -> u32 {
    crc32_update(0xffff_ffff, bytes) ^ 0xffff_ffff
}

fn crc32_update(mut crc: u32, bytes: &[u8]) -> u32 {
    static TABLE: std::sync::OnceLock<[u32; 256]> = std::sync::OnceLock::new();
    let table = TABLE.get_or_init(|| {
        let mut table = [0_u32; 256];
        let mut i = 0;
        while i < 256 {
            let mut value = i as u32;
            let mut bit = 0;
            while bit < 8 {
                value = if value & 1 == 1 {
                    0xedb8_8320 ^ (value >> 1)
                } else {
                    value >> 1
                };
                bit += 1;
            }
            table[i] = value;
            i += 1;
        }
        table
    });

    for &byte in bytes {
        crc = table[((crc ^ u32::from(byte)) & 0xff) as usize] ^ (crc >> 8);
    }
    crc
}
