use anyhow::{bail, ensure, Context, Result};
use std::{fs, net::IpAddr, path::Path};

const METADATA_MARKER: &[u8] = b"\xab\xcd\xefMaxMind.com";
const MAX_METADATA_BYTES: usize = 128 * 1024;
const MAX_VALUES: usize = 65_536;
const MAX_DEPTH: usize = 512;
const MAX_STRING_BYTES: usize = 4096;

#[derive(Debug)]
pub(super) struct Database {
    bytes: Vec<u8>,
    node_count: u32,
    record_size: u16,
    ip_version: u16,
    search_tree_size: usize,
    data_section_start: usize,
    metadata_offset: usize,
    database_type: String,
}

#[derive(Clone, Copy)]
struct Header {
    kind: u8,
    size: usize,
    payload: usize,
    next: usize,
    pointer: Option<usize>,
}

#[derive(Default)]
struct Budget {
    values: usize,
}

impl Budget {
    fn charge(&mut self) -> Result<()> {
        self.values = self.values.saturating_add(1);
        ensure!(
            self.values <= MAX_VALUES,
            "MMDB value traversal limit exceeded"
        );
        Ok(())
    }
}

impl Database {
    pub(super) fn load(path: &Path, max_bytes: u64) -> Result<Self> {
        let metadata = fs::symlink_metadata(path)
            .with_context(|| format!("open GeoIP MMDB {}", path.display()))?;
        ensure!(
            metadata.is_file() && !metadata.file_type().is_symlink(),
            "GeoIP MMDB must be a regular file"
        );
        ensure!(
            metadata.len() <= max_bytes,
            "GeoIP MMDB exceeds configured maximum size"
        );
        let bytes =
            fs::read(path).with_context(|| format!("read GeoIP MMDB {}", path.display()))?;
        let marker = metadata_marker(&bytes).context("MMDB metadata marker not found")?;
        let metadata_offset = marker + METADATA_MARKER.len();
        let base = metadata_offset;
        let node_count = u32::try_from(read_uint_path(
            &bytes,
            metadata_offset,
            base,
            &["node_count"],
        )?)
        .context("MMDB node_count is too large")?;
        let record_size = u16::try_from(read_uint_path(
            &bytes,
            metadata_offset,
            base,
            &["record_size"],
        )?)
        .context("MMDB record_size is too large")?;
        let ip_version = u16::try_from(read_uint_path(
            &bytes,
            metadata_offset,
            base,
            &["ip_version"],
        )?)
        .context("MMDB ip_version is too large")?;
        let major = read_uint_path(
            &bytes,
            metadata_offset,
            base,
            &["binary_format_major_version"],
        )?;
        let database_type = read_string_path(&bytes, metadata_offset, base, &["database_type"])?;
        ensure!(node_count > 0, "MMDB node_count must be positive");
        ensure!(
            matches!(record_size, 24 | 28 | 32),
            "unsupported MMDB record_size {record_size}"
        );
        ensure!(
            matches!(ip_version, 4 | 6),
            "invalid MMDB ip_version {ip_version}"
        );
        ensure!(major == 2, "unsupported MMDB major format version {major}");
        let node_bytes = usize::from(record_size)
            .checked_mul(2)
            .and_then(|bits| bits.checked_div(8))
            .context("MMDB node size overflow")?;
        let search_tree_size = (node_count as usize)
            .checked_mul(node_bytes)
            .context("MMDB search tree size overflow")?;
        let data_section_start = search_tree_size
            .checked_add(16)
            .context("MMDB data section overflow")?;
        ensure!(
            data_section_start <= marker,
            "MMDB search tree overlaps metadata"
        );
        ensure!(
            bytes.get(search_tree_size..data_section_start) == Some(&[0u8; 16][..]),
            "MMDB data-section separator is invalid"
        );
        Ok(Self {
            bytes,
            node_count,
            record_size,
            ip_version,
            search_tree_size,
            data_section_start,
            metadata_offset,
            database_type,
        })
    }

    pub(super) fn database_type(&self) -> &str {
        &self.database_type
    }

    pub(super) fn country(&self, ip: IpAddr) -> Option<String> {
        let offset = self.lookup_offset(ip).ok().flatten()?;
        for path in [
            &["country", "iso_code"][..],
            &["registered_country", "iso_code"][..],
            &["country_code"][..],
        ] {
            if let Ok(code) = read_string_path(&self.bytes, offset, self.data_section_start, path) {
                if let Some(code) = super::normalize_country_code(&code) {
                    return Some(String::from_utf8_lossy(&code).into_owned());
                }
            }
        }
        let mut budget = Budget::default();
        if let Ok(code) = read_string(&self.bytes, offset, self.data_section_start, 0, &mut budget)
        {
            if let Some(code) = super::normalize_country_code(&code) {
                return Some(String::from_utf8_lossy(&code).into_owned());
            }
        }
        None
    }

    fn lookup_offset(&self, ip: IpAddr) -> Result<Option<usize>> {
        let (bytes, bits): ([u8; 16], usize) = match (self.ip_version, ip) {
            (4, IpAddr::V4(value)) => {
                let mut out = [0u8; 16];
                out[..4].copy_from_slice(&value.octets());
                (out, 32)
            }
            (4, IpAddr::V6(_)) => return Ok(None),
            (6, IpAddr::V4(value)) => {
                let mut out = [0u8; 16];
                out[12..].copy_from_slice(&value.octets());
                (out, 128)
            }
            (6, IpAddr::V6(value)) => (value.octets(), 128),
            _ => bail!("invalid MMDB IP version"),
        };
        let mut node = 0u32;
        for bit_index in 0..bits {
            if node >= self.node_count {
                break;
            }
            let bit = (bytes[bit_index / 8] >> (7 - (bit_index % 8))) & 1;
            node = self.read_node(node, bit as usize)?;
        }
        if node == self.node_count {
            return Ok(None);
        }
        ensure!(
            node >= self.node_count.saturating_add(16),
            "invalid MMDB data pointer"
        );
        let offset = (node as usize)
            .checked_sub(self.node_count as usize)
            .and_then(|value| value.checked_add(self.search_tree_size))
            .context("MMDB data pointer overflow")?;
        ensure!(
            offset >= self.data_section_start
                && offset < self.metadata_offset - METADATA_MARKER.len(),
            "MMDB data pointer is out of bounds"
        );
        Ok(Some(offset))
    }

    fn read_node(&self, node: u32, side: usize) -> Result<u32> {
        ensure!(
            (node as usize) < self.node_count as usize,
            "MMDB node is out of range"
        );
        let base = match self.record_size {
            24 => node as usize * 6,
            28 => node as usize * 7,
            32 => node as usize * 8,
            _ => unreachable!(),
        };
        let value = match self.record_size {
            24 => read_be(&self.bytes, base + side * 3, 3)? as u32,
            28 => {
                let middle = *self.bytes.get(base + 3).context("MMDB node is truncated")?;
                if side == 0 {
                    (((middle & 0xf0) as u32) << 20) | read_be(&self.bytes, base, 3)? as u32
                } else {
                    (((middle & 0x0f) as u32) << 24) | read_be(&self.bytes, base + 4, 3)? as u32
                }
            }
            32 => read_be(&self.bytes, base + side * 4, 4)? as u32,
            _ => unreachable!(),
        };
        Ok(value)
    }
}

fn metadata_marker(bytes: &[u8]) -> Option<usize> {
    let start = bytes
        .len()
        .saturating_sub(MAX_METADATA_BYTES + METADATA_MARKER.len());
    bytes[start..]
        .windows(METADATA_MARKER.len())
        .rposition(|window| window == METADATA_MARKER)
        .map(|offset| start + offset)
}

fn read_uint_path(bytes: &[u8], offset: usize, base: usize, path: &[&str]) -> Result<u64> {
    let mut budget = Budget::default();
    let value = find_path(bytes, offset, base, path, 0, &mut budget)?
        .context("MMDB metadata field missing")?;
    read_uint(bytes, value, base, 0, &mut budget)
}

fn read_string_path(bytes: &[u8], offset: usize, base: usize, path: &[&str]) -> Result<String> {
    let mut budget = Budget::default();
    let value = find_path(bytes, offset, base, path, 0, &mut budget)?
        .context("MMDB string field missing")?;
    read_string(bytes, value, base, 0, &mut budget)
}

fn find_path(
    bytes: &[u8],
    offset: usize,
    base: usize,
    path: &[&str],
    depth: usize,
    budget: &mut Budget,
) -> Result<Option<usize>> {
    ensure!(depth <= MAX_DEPTH, "MMDB nesting limit exceeded");
    budget.charge()?;
    let header = header(bytes, offset, base)?;
    if let Some(pointer) = header.pointer {
        return find_path(bytes, pointer, base, path, depth + 1, budget);
    }
    if path.is_empty() {
        return Ok(Some(offset));
    }
    ensure!(header.kind == 7, "MMDB selected value is not a map");
    let mut cursor = header.payload;
    for _ in 0..header.size {
        let (key, next) = inline_string(bytes, cursor, base)?;
        cursor = next;
        let value_offset = cursor;
        if key == path[0] {
            if path.len() == 1 {
                return Ok(Some(value_offset));
            }
            if let Some(found) =
                find_path(bytes, value_offset, base, &path[1..], depth + 1, budget)?
            {
                return Ok(Some(found));
            }
        }
        cursor = skip(bytes, value_offset, base, depth + 1, budget)?;
    }
    Ok(None)
}

fn read_string(
    bytes: &[u8],
    offset: usize,
    base: usize,
    depth: usize,
    budget: &mut Budget,
) -> Result<String> {
    ensure!(depth <= MAX_DEPTH, "MMDB nesting limit exceeded");
    budget.charge()?;
    let h = header(bytes, offset, base)?;
    if let Some(pointer) = h.pointer {
        return read_string(bytes, pointer, base, depth + 1, budget);
    }
    ensure!(h.kind == 2, "MMDB value is not a string");
    ensure!(
        h.size <= MAX_STRING_BYTES,
        "MMDB string exceeds the decode limit"
    );
    let data = bytes
        .get(h.payload..h.next)
        .context("MMDB string is truncated")?;
    Ok(std::str::from_utf8(data)
        .context("MMDB string is not UTF-8")?
        .to_owned())
}

fn read_uint(
    bytes: &[u8],
    offset: usize,
    base: usize,
    depth: usize,
    budget: &mut Budget,
) -> Result<u64> {
    ensure!(depth <= MAX_DEPTH, "MMDB nesting limit exceeded");
    budget.charge()?;
    let h = header(bytes, offset, base)?;
    if let Some(pointer) = h.pointer {
        return read_uint(bytes, pointer, base, depth + 1, budget);
    }
    ensure!(
        matches!(h.kind, 5 | 6 | 9),
        "MMDB value is not an unsigned integer"
    );
    ensure!(h.size <= 8, "MMDB integer is too wide");
    read_be(bytes, h.payload, h.size)
}

fn inline_string(bytes: &[u8], offset: usize, base: usize) -> Result<(&str, usize)> {
    let h = header(bytes, offset, base)?;
    ensure!(
        h.pointer.is_none() && h.kind == 2,
        "MMDB map key is not an inline string"
    );
    let value = bytes
        .get(h.payload..h.next)
        .context("MMDB map key is truncated")?;
    Ok((
        std::str::from_utf8(value).context("MMDB map key is not UTF-8")?,
        h.next,
    ))
}

fn skip(
    bytes: &[u8],
    offset: usize,
    base: usize,
    depth: usize,
    budget: &mut Budget,
) -> Result<usize> {
    ensure!(depth <= MAX_DEPTH, "MMDB nesting limit exceeded");
    budget.charge()?;
    let h = header(bytes, offset, base)?;
    if h.pointer.is_some() {
        return Ok(h.next);
    }
    let mut cursor = h.payload;
    match h.kind {
        7 => {
            for _ in 0..h.size {
                cursor = skip(bytes, cursor, base, depth + 1, budget)?;
                cursor = skip(bytes, cursor, base, depth + 1, budget)?;
            }
            Ok(cursor)
        }
        11 => {
            for _ in 0..h.size {
                cursor = skip(bytes, cursor, base, depth + 1, budget)?;
            }
            Ok(cursor)
        }
        14 => Ok(h.payload),
        12 | 13 => bail!("unsupported deprecated MMDB data type {}", h.kind),
        _ => Ok(h.next),
    }
}

fn header(bytes: &[u8], offset: usize, base: usize) -> Result<Header> {
    let control = *bytes.get(offset).context("MMDB value is truncated")?;
    let primary = control >> 5;
    let size_bits = (control & 0x1f) as usize;
    if primary == 1 {
        let class = (size_bits >> 3) & 0x03;
        let low = size_bits & 0x07;
        let count = class + 1;
        let payload = offset + 1;
        let raw = read_be(bytes, payload, count)? as usize;
        let pointer = match class {
            0 => (low << 8) | raw,
            1 => ((low << 16) | raw)
                .checked_add(2048)
                .context("MMDB pointer overflow")?,
            2 => ((low << 24) | raw)
                .checked_add(526_336)
                .context("MMDB pointer overflow")?,
            3 => raw,
            _ => unreachable!(),
        };
        let target = base.checked_add(pointer).context("MMDB pointer overflow")?;
        ensure!(target < bytes.len(), "MMDB pointer is out of bounds");
        return Ok(Header {
            kind: 1,
            size: 0,
            payload,
            next: payload + count,
            pointer: Some(target),
        });
    }
    let mut cursor = offset + 1;
    let kind = if primary == 0 {
        let extended = *bytes
            .get(cursor)
            .context("MMDB extended type is truncated")?;
        cursor += 1;
        extended.checked_add(7).context("MMDB type overflow")?
    } else {
        primary
    };
    let size = match size_bits {
        0..=28 => size_bits,
        29 => 29 + *bytes.get(cursor).context("MMDB size is truncated")? as usize,
        30 => 285 + read_be(bytes, cursor, 2)? as usize,
        31 => 65_821 + read_be(bytes, cursor, 3)? as usize,
        _ => unreachable!(),
    };
    cursor += match size_bits {
        29 => 1,
        30 => 2,
        31 => 3,
        _ => 0,
    };
    let next = match kind {
        7 | 11 | 14 => cursor,
        _ => cursor
            .checked_add(size)
            .context("MMDB field size overflow")?,
    };
    ensure!(next <= bytes.len(), "MMDB field is out of bounds");
    Ok(Header {
        kind,
        size,
        payload: cursor,
        next,
        pointer: None,
    })
}

fn read_be(bytes: &[u8], offset: usize, count: usize) -> Result<u64> {
    ensure!(count <= 8, "MMDB integer is too wide");
    let end = offset.checked_add(count).context("MMDB read overflow")?;
    let data = bytes.get(offset..end).context("MMDB value is truncated")?;
    Ok(data
        .iter()
        .fold(0u64, |value, byte| (value << 8) | u64::from(*byte)))
}

#[cfg(test)]
include!("mmdb_tests.rs");