//! 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,
}

impl StealthDnsForwarder {
    pub fn new(listen: SocketAddr, upstream: SocketAddr) -> Self {
        Self {
            listen_addr: listen,
            upstream,
        }
    }

    /// 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];

            // 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);
                }
            }
        }
    }

    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
}