//! Stealth DNS resolver - DoH/DoT with anti-hijacking

use bytes::{BufMut, BytesMut};
use std::net::SocketAddr;
use thiserror::Error;
use tokio::net::UdpSocket;

#[derive(Error, Debug)]
pub enum DnsError {
    #[error("io error: {0}")]
    Io(#[from] std::io::Error),
    #[error("invalid query")]
    InvalidQuery,
    #[error("upstream timeout")]
    Timeout,
    #[error("hijacking detected")]
    HijackDetected,
}

/// DNS query types
#[derive(Debug, Clone, Copy)]
pub enum QueryType {
    A = 1,
    AAAA = 28,
    CNAME = 5,
    TXT = 16,
}

/// Minimal DNS query builder
pub fn build_dns_query(domain: &str, qtype: QueryType) -> BytesMut {
    let mut buf = BytesMut::with_capacity(512);

    // Transaction ID (random)
    let txid: u16 = rand::random();
    buf.put_u16(txid);

    // Flags: standard query, recursion desired
    buf.put_u16(0x0100);

    // Questions: 1, Answers: 0, Authority: 0, Additional: 0
    buf.put_u16(1);
    buf.put_u16(0);
    buf.put_u16(0);
    buf.put_u16(0);

    // QNAME
    for label in domain.split('.') {
        buf.put_u8(label.len() as u8);
        buf.put_slice(label.as_bytes());
    }
    buf.put_u8(0); // null terminator

    // QTYPE and QCLASS (IN)
    buf.put_u16(qtype as u16);
    buf.put_u16(1); // IN class

    buf
}

/// DNS forwarder that tunnels queries through OSTP
pub struct StealthDnsForwarder {
    listen_addr: SocketAddr,
    /// Upstream resolver (will be tunneled through OSTP)
    upstream: SocketAddr,
    /// Master node internal IP for *.ospab.internal queries
    master_node_ip: std::net::IpAddr,
}

impl StealthDnsForwarder {
    pub fn new(listen: SocketAddr, upstream: SocketAddr) -> Self {
        Self {
            listen_addr: listen,
            upstream,
            master_node_ip: std::net::IpAddr::V4(std::net::Ipv4Addr::new(10, 8, 0, 1)),
        }
    }

    pub fn with_master_ip(mut self, master_ip: std::net::IpAddr) -> Self {
        self.master_node_ip = master_ip;
        self
    }

    /// Start DNS listener (intercepts local queries)
    pub async fn run(&self) -> Result<(), DnsError> {
        let socket = UdpSocket::bind(self.listen_addr).await?;
        tracing::info!("DNS forwarder listening on {}", self.listen_addr);

        let mut buf = [0u8; 512];
        loop {
            let (len, src) = socket.recv_from(&mut buf).await?;
            let query = &buf[..len];

            // Check if query is for *.ospab.internal
            if let Some(response) = self.handle_internal_query(query) {
                let _ = socket.send_to(&response, src).await;
                continue;
            }

            // TODO: Encrypt and forward through OSTP tunnel instead of direct
            // For now, direct forward (to be replaced with tunnel)
            match self.forward_query(query).await {
                Ok(response) => {
                    let _ = socket.send_to(&response, src).await;
                }
                Err(e) => {
                    tracing::warn!("DNS forward failed: {}", e);
                }
            }
        }
    }

    /// Handle *.ospab.internal queries - resolve to master node IP
    fn handle_internal_query(&self, query: &[u8]) -> Option<Vec<u8>> {
        if query.len() < 12 {
            return None;
        }

        // Parse domain from query (simplified)
        let qname_start = 12;
        let mut pos = qname_start;
        let mut domain = String::new();

        while pos < query.len() && query[pos] != 0 {
            let len = query[pos] as usize;
            if len == 0 {
                break;
            }
            if !domain.is_empty() {
                domain.push('.');
            }
            pos += 1;
            if pos + len > query.len() {
                return None;
            }
            domain.push_str(&String::from_utf8_lossy(&query[pos..pos + len]));
            pos += len;
        }

        // Check if domain ends with .ospab.internal
        if !domain.ends_with(".ospab.internal") {
            return None;
        }

        tracing::debug!("Resolving internal domain: {}", domain);

        // Build DNS response
        let mut response = Vec::from(query);

        // Set response flags (QR=1, AA=1, RCODE=0)
        response[2] = 0x84; // QR + AA
        response[3] = 0x00; // No error

        // Set answer count to 1
        response[6] = 0x00;
        response[7] = 0x01;

        // Add answer section (pointer to question + A record)
        response.push(0xC0); // Pointer to question name
        response.push(0x0C); // Offset 12

        // Type A
        response.push(0x00);
        response.push(0x01);

        // Class IN
        response.push(0x00);
        response.push(0x01);

        // TTL (300 seconds)
        response.push(0x00);
        response.push(0x00);
        response.push(0x01);
        response.push(0x2C);

        // Data length (4 bytes for IPv4)
        response.push(0x00);
        response.push(0x04);

        // IP address
        match self.master_node_ip {
            std::net::IpAddr::V4(ip) => {
                response.extend_from_slice(&ip.octets());
            }
            std::net::IpAddr::V6(_) => {
                // For IPv6, would need AAAA record - fallback to IPv4 loopback
                response.extend_from_slice(&[127, 0, 0, 1]);
            }
        }

        Some(response)
    }

    async fn forward_query(&self, query: &[u8]) -> Result<Vec<u8>, DnsError> {
        let socket = UdpSocket::bind("0.0.0.0:0").await?;
        socket.send_to(query, self.upstream).await?;

        let mut response = vec![0u8; 512];
        let timeout = tokio::time::timeout(
            std::time::Duration::from_secs(5),
            socket.recv_from(&mut response),
        )
        .await;

        match timeout {
            Ok(Ok((len, _))) => {
                response.truncate(len);
                Ok(response)
            }
            Ok(Err(e)) => Err(DnsError::Io(e)),
            Err(_) => Err(DnsError::Timeout),
        }
    }
}

/// Anti-hijacking: verify response matches known-good resolver fingerprint
pub fn detect_hijack(response: &[u8], expected_patterns: &[[u8; 4]]) -> bool {
    // Check if response contains known hijack IPs (ISP redirect pages, etc.)
    // Simplified: check for common hijack patterns
    if response.len() < 12 {
        return true; // Suspicious short response
    }

    // Check for NXDOMAIN being rewritten (common hijack)
    let rcode = response[3] & 0x0F;
    if rcode == 0 {
        // NOERROR - check if IP is in known hijack list
        for pattern in expected_patterns {
            if response.windows(4).any(|w| w == pattern) {
                return true;
            }
        }
    }
    false
}
