use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};

use crate::{IpRange, Record, Registry, Status};

pub struct Entry4 {
    pub start: u32,
    pub end: u32,
    pub cc: [u8; 2],
    pub registry: Registry,
    pub status: Status,
    pub date: u32,
}

pub struct Entry6 {
    pub start: u128,
    pub end: u128,
    pub cc: [u8; 2],
    pub registry: Registry,
    pub status: Status,
    pub date: u32,
}

/// Sorted table of IPv4 ranges for binary-search lookup.
pub struct Table4 {
    entries: Vec<Entry4>,
}

impl Table4 {
    pub fn build(mut entries: Vec<Entry4>) -> Self {
        entries.sort_unstable_by_key(|e| e.start);
        Self { entries }
    }

    pub fn lookup(&self, ip: u32) -> Option<Record> {
        let idx = self.entries.partition_point(|e| e.start <= ip);
        if idx == 0 {
            return None;
        }
        let e = &self.entries[idx - 1];
        if ip <= e.end {
            Some(Record::new(
                e.cc,
                e.registry,
                e.status,
                e.date,
                IpRange {
                    start: IpAddr::V4(Ipv4Addr::from(e.start)),
                    end: IpAddr::V4(Ipv4Addr::from(e.end)),
                },
            ))
        } else {
            None
        }
    }

    pub fn len(&self) -> usize {
        self.entries.len()
    }
}

/// Sorted table of IPv6 ranges for binary-search lookup.
pub struct Table6 {
    entries: Vec<Entry6>,
}

impl Table6 {
    pub fn build(mut entries: Vec<Entry6>) -> Self {
        entries.sort_unstable_by_key(|e| e.start);
        Self { entries }
    }

    pub fn lookup(&self, ip: u128) -> Option<Record> {
        let idx = self.entries.partition_point(|e| e.start <= ip);
        if idx == 0 {
            return None;
        }
        let e = &self.entries[idx - 1];
        if ip <= e.end {
            Some(Record::new(
                e.cc,
                e.registry,
                e.status,
                e.date,
                IpRange {
                    start: IpAddr::V6(Ipv6Addr::from(e.start)),
                    end: IpAddr::V6(Ipv6Addr::from(e.end)),
                },
            ))
        } else {
            None
        }
    }

    pub fn len(&self) -> usize {
        self.entries.len()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn t4(start: &str, count: u32, cc: &str) -> Entry4 {
        let s = u32::from(start.parse::<std::net::Ipv4Addr>().unwrap());
        let mut cc_bytes = [0u8; 2];
        cc_bytes.copy_from_slice(cc.as_bytes());
        Entry4 {
            start: s,
            end: s + count - 1,
            cc: cc_bytes,
            registry: Registry::Apnic,
            status: Status::Allocated,
            date: 20110811,
        }
    }

    #[test]
    fn lookup_hit() {
        let table = Table4::build(vec![
            t4("1.0.0.0", 256, "JP"),
            t4("1.0.1.0", 256, "CN"),
        ]);
        let r = table.lookup(u32::from("1.0.0.1".parse::<std::net::Ipv4Addr>().unwrap())).unwrap();
        assert_eq!(r.cc(), "JP");
        assert_eq!(r.registry, Registry::Apnic);
        assert_eq!(r.status, Status::Allocated);
        assert_eq!(r.date, 20110811);

        let r2 = table.lookup(u32::from("1.0.1.255".parse::<std::net::Ipv4Addr>().unwrap())).unwrap();
        assert_eq!(r2.cc(), "CN");
    }

    #[test]
    fn lookup_miss_between_ranges() {
        let table = Table4::build(vec![
            t4("1.0.0.0", 256, "JP"),
            t4("1.0.2.0", 256, "AU"),
        ]);
        assert!(table.lookup(u32::from("1.0.1.0".parse::<std::net::Ipv4Addr>().unwrap())).is_none());
    }

    #[test]
    fn lookup_miss_before_first() {
        let table = Table4::build(vec![t4("10.0.0.0", 256, "JP")]);
        assert!(table.lookup(u32::from("1.0.0.0".parse::<std::net::Ipv4Addr>().unwrap())).is_none());
    }

    #[test]
    fn lookup_exact_boundary() {
        let table = Table4::build(vec![t4("192.0.2.0", 256, "US")]);
        let start = u32::from("192.0.2.0".parse::<std::net::Ipv4Addr>().unwrap());
        assert_eq!(table.lookup(start).unwrap().cc(), "US");
        assert_eq!(table.lookup(start + 255).unwrap().cc(), "US");
        assert!(table.lookup(start + 256).is_none());
    }

    #[test]
    fn lookup_range_correct() {
        let table = Table4::build(vec![t4("1.0.0.0", 256, "JP")]);
        let rec = table.lookup(u32::from("1.0.0.50".parse::<std::net::Ipv4Addr>().unwrap())).unwrap();
        assert_eq!(rec.range.start, "1.0.0.0".parse::<IpAddr>().unwrap());
        assert_eq!(rec.range.end,   "1.0.0.255".parse::<IpAddr>().unwrap());
    }
}
