use crate::geoparsers::v2ray::types::{Domain, GeoSite, GeoSiteList};
use prost::bytes::Buf;
use prost::Message;
use std::fs;

pub struct GeoSiteService {
    index: GeoSiteList,
}

impl GeoSiteService {
    // TODO: Make more smart memory mapping; geosite files can be > 70MB
    pub fn new(path: &str) -> Result<Self, Box<dyn std::error::Error>> {
        let bytes = fs::read(path)?;
        let geosite_list = decode_geosite_stream(&bytes)?;

        Ok(Self {
            index: geosite_list,
        })
    }

    // Idk but i think it can work
    pub fn lookup(&self, value: &str) -> Option<&GeoSite> {
        self.index
            .entry
            .iter()
            .find(|site| site.domain.iter().any(|d| d.value == value))
    }

    /// Returns the number of GeoSite entries in the list
    pub fn len(&self) -> usize {
        self.index.entry.len()
    }

    /// Returns true if the GeoSite list is empty
    pub fn is_empty(&self) -> bool {
        self.index.entry.is_empty()
    }
}

/// Decode a stream of length-delimited GeoSite messages
/// `geosite.dat` ts is not one protobuf-message, stream of length-delimited messages
/// so we need ts helper
fn decode_geosite_stream(bytes: &[u8]) -> Result<GeoSiteList, Box<dyn std::error::Error>> {
    let mut buf = bytes;
    let mut entries = Vec::new();

    while buf.has_remaining() {
        // Read tag (0x0a field 1, wire type 2)
        let tag = buf.get_u8();
        if tag != 0x0a {
            return Err(format!("Unexpected tag: {:#04x}", tag).into());
        }
        // varint
        let mut len = 0usize;
        let mut shift = 0;
        loop {
            if !buf.has_remaining() {
                return Err("Unexpected end of buffer while reading varint".into());
            }
            let b = buf.get_u8();
            len |= ((b & 0x7f) as usize) << shift;
            if b & 0x80 == 0 {
                break;
            }
            shift += 7;
            if shift >= 70 {
                return Err("Varint too long".into());
            }
        }

        let entry_bytes = &buf[..len];
        let site = GeoSite::decode(entry_bytes)?;
        entries.push(site);

        buf.advance(len);
    }

    Ok(GeoSiteList { entry: entries })
}