//! DNS query handler

use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;

use tokio::net::UdpSocket;
use tracing::{debug, error, info, warn};
use trust_dns_proto::op::{Message, MessageType, OpCode, ResponseCode};
use trust_dns_proto::rr::rdata::{A, AAAA, CNAME, MX, NS, TXT};
use trust_dns_proto::rr::{Name, RData, Record as DnsRecord, RecordType as DnsRecordType};
use trust_dns_proto::serialize::binary::{BinDecodable, BinEncodable};
use ipnet::IpNet;

use crate::dns::ptr_patterns::{parse_ptr_query_to_ip, apply_pattern};
use crate::metadata::DnsMetadataStore;
use flashdns_types::{Record, RecordData, RecordType, ReverseZone, Zone, ZoneId};

/// DNS query handler
pub struct DnsHandler {
    /// UDP socket for DNS queries
    socket: Arc<UdpSocket>,
    /// Metadata store for zones and records
    metadata: Arc<DnsMetadataStore>,
}

/// DNS handler error
#[derive(Debug, thiserror::Error)]
pub enum DnsError {
    #[error("Parse error: {0}")]
    Parse(String),
    #[error("Metadata error: {0}")]
    Metadata(String),
    #[error("Not found: {0}")]
    NotFound(String),
    #[error("Encode error: {0}")]
    Encode(String),
}

impl DnsHandler {
    /// Create a new DNS handler bound to the given address
    pub async fn bind(addr: SocketAddr, metadata: Arc<DnsMetadataStore>) -> std::io::Result<Self> {
        let socket = UdpSocket::bind(addr).await?;
        info!("DNS handler listening on UDP {}", addr);
        Ok(Self {
            socket: Arc::new(socket),
            metadata,
        })
    }

    /// Run the DNS handler (process queries)
    pub async fn run(self) {
        let mut buf = [0u8; 512]; // Standard DNS UDP max size

        loop {
            match self.socket.recv_from(&mut buf).await {
                Ok((len, src)) => {
                    let query = buf[..len].to_vec();
                    let socket = Arc::clone(&self.socket);
                    let metadata = Arc::clone(&self.metadata);

                    // Spawn task to handle query
                    tokio::spawn(async move {
                        let handler = DnsQueryHandler::new(metadata);
                        match handler.handle_query(&query).await {
                            Ok(response) => {
                                if let Err(e) = socket.send_to(&response, src).await {
                                    error!("Failed to send DNS response to {}: {}", src, e);
                                }
                            }
                            Err(e) => {
                                warn!("Failed to handle DNS query from {}: {}", src, e);
                                // Send SERVFAIL response
                                if let Ok(response) = Self::servfail_response(&query) {
                                    let _ = socket.send_to(&response, src).await;
                                }
                            }
                        }
                    });
                }
                Err(e) => {
                    error!("Error receiving DNS query: {}", e);
                }
            }
        }
    }

    /// Build SERVFAIL response
    fn servfail_response(query: &[u8]) -> Result<Vec<u8>, &'static str> {
        Self::error_response(query, 2) // RCODE 2 = SERVFAIL
    }

    /// Build error response with given RCODE
    fn error_response(query: &[u8], rcode: u8) -> Result<Vec<u8>, &'static str> {
        if query.len() < 12 {
            return Err("Query too short");
        }

        let mut response = query.to_vec();

        // Set QR bit (response) and preserve opcode
        response[2] = (response[2] & 0x78) | 0x80;

        // Set RCODE
        response[3] = (response[3] & 0xF0) | (rcode & 0x0F);

        // Zero out counts except QDCOUNT
        response[6] = 0;
        response[7] = 0;
        response[8] = 0;
        response[9] = 0;
        response[10] = 0;
        response[11] = 0;

        Ok(response)
    }
}

/// Internal query handler with metadata access
struct DnsQueryHandler {
    metadata: Arc<DnsMetadataStore>,
}

impl DnsQueryHandler {
    fn new(metadata: Arc<DnsMetadataStore>) -> Self {
        Self { metadata }
    }

    /// Handle a DNS query and return response
    async fn handle_query(&self, query_bytes: &[u8]) -> Result<Vec<u8>, DnsError> {
        // Parse the query
        let query = Message::from_bytes(query_bytes)
            .map_err(|e| DnsError::Parse(format!("Failed to parse DNS message: {}", e)))?;

        debug!(
            "DNS query: id={} opcode={:?} questions={}",
            query.id(),
            query.op_code(),
            query.query_count()
        );

        // Only handle standard queries
        if query.op_code() != OpCode::Query {
            return self.build_error_response(&query, ResponseCode::NotImp);
        }

        // Get the first question
        let question = query
            .queries()
            .first()
            .ok_or_else(|| DnsError::Parse("No question in query".to_string()))?;

        let qname = question.name();
        let qtype = question.query_type();

        debug!("Question: {} {:?}", qname, qtype);

        // Find the zone that matches this query
        let zone = match self.find_zone_for_name(qname).await? {
            Some(z) => z,
            None => {
                debug!("No zone found for {}", qname);
                return self.build_error_response(&query, ResponseCode::Refused);
            }
        };

        debug!("Found zone: {} (id={})", zone.name, zone.id);

        // Handle PTR queries with reverse zone pattern matching
        if qtype == DnsRecordType::PTR {
            if let Some((ptr_value, ttl)) = self.handle_ptr_query(qname).await {
                // Build PTR response from pattern
                return self.build_ptr_response(&query, qname, &ptr_value, ttl);
            }
            // If no reverse zone match, fall through to normal record lookup
        }

        // Get the record name relative to the zone
        let record_name = self.get_relative_name(qname, &zone);

        // Look up records
        let records = self.lookup_records(&zone.id, &record_name, qtype).await?;

        if records.is_empty() {
            debug!("No records found for {} {:?}", record_name, qtype);
            return self.build_nxdomain_response(&query, &zone);
        }

        // Build success response
        self.build_success_response(&query, &zone, &records, qname)
    }

    /// Find the zone that matches a query name (longest suffix match)
    async fn find_zone_for_name(&self, qname: &Name) -> Result<Option<Zone>, DnsError> {
        // Get all zones (we list for all orgs - in production, this would be cached)
        // For now, we do a simple scan. In production, use a proper zone cache.
        let qname_str = qname.to_lowercase().to_string();

        // List zones for all orgs (this is a simplification - production would have proper indexing)
        // We'll try common org prefixes or use a special "global" org
        let _zones = self
            .metadata
            .list_zones("*", None)
            .await
            .unwrap_or_default();

        // Also try listing without org filter by scanning known prefixes
        let all_zones = self.get_all_zones().await;

        // Find longest suffix match
        let mut best_match: Option<Zone> = None;
        let mut best_len = 0;

        for zone in all_zones {
            let zone_name = zone.name.as_str().to_lowercase();
            if (qname_str.ends_with(&zone_name) || qname_str == zone_name)
                && zone_name.len() > best_len {
                    best_len = zone_name.len();
                    best_match = Some(zone);
                }
        }

        Ok(best_match)
    }

    /// Get all zones from metadata (simplified - in production use proper caching/indexing)
    async fn get_all_zones(&self) -> Vec<Zone> {
        // Try to list zones with wildcard - if that doesn't work, fall back to empty
        // In production, this would be a proper zone cache
        self.metadata
            .list_zones("*", None)
            .await
            .unwrap_or_default()
    }

    /// Get the record name relative to the zone
    fn get_relative_name(&self, qname: &Name, zone: &Zone) -> String {
        let qname_str = qname.to_lowercase().to_string();
        let zone_name = zone.name.as_str().to_lowercase();

        if qname_str == zone_name {
            "@".to_string()
        } else if let Some(prefix) = qname_str.strip_suffix(&zone_name) {
            // Remove trailing dot from prefix
            prefix.trim_end_matches('.').to_string()
        } else {
            qname_str
        }
    }

    /// Look up records for a name and type
    async fn lookup_records(
        &self,
        zone_id: &ZoneId,
        record_name: &str,
        qtype: DnsRecordType,
    ) -> Result<Vec<Record>, DnsError> {
        let records = self
            .metadata
            .list_records_by_name(zone_id, record_name)
            .await
            .map_err(|e| DnsError::Metadata(e.to_string()))?;

        // Filter by type
        let filtered: Vec<_> = records
            .into_iter()
            .filter(|r| {
                r.enabled && (qtype == DnsRecordType::ANY || self.matches_type(r.record_type, qtype))
            })
            .collect();

        Ok(filtered)
    }

    /// Check if a record type matches the query type
    fn matches_type(&self, record_type: RecordType, qtype: DnsRecordType) -> bool {
        match (record_type, qtype) {
            (RecordType::A, DnsRecordType::A) => true,
            (RecordType::Aaaa, DnsRecordType::AAAA) => true,
            (RecordType::Cname, DnsRecordType::CNAME) => true,
            (RecordType::Mx, DnsRecordType::MX) => true,
            (RecordType::Txt, DnsRecordType::TXT) => true,
            (RecordType::Ns, DnsRecordType::NS) => true,
            (RecordType::Srv, DnsRecordType::SRV) => true,
            (RecordType::Ptr, DnsRecordType::PTR) => true,
            (RecordType::Caa, DnsRecordType::CAA) => true,
            (RecordType::Soa, DnsRecordType::SOA) => true,
            _ => false,
        }
    }

    /// Convert flashdns record to DNS record
    fn convert_record(&self, record: &Record, qname: &Name) -> Option<DnsRecord> {
        let rdata = self.convert_record_data(&record.data)?;
        let dns_record_type = self.to_dns_record_type(record.record_type)?;

        let mut dns_record = DnsRecord::new();
        dns_record.set_name(qname.clone());
        dns_record.set_rr_type(dns_record_type);
        dns_record.set_ttl(record.ttl.as_secs());
        dns_record.set_data(Some(rdata));

        Some(dns_record)
    }

    /// Convert flashdns RecordData to DNS RData
    fn convert_record_data(&self, data: &RecordData) -> Option<RData> {
        match data {
            RecordData::A { address } => {
                let addr = std::net::Ipv4Addr::new(address[0], address[1], address[2], address[3]);
                Some(RData::A(A(addr)))
            }
            RecordData::Aaaa { address } => {
                let addr = std::net::Ipv6Addr::from(*address);
                Some(RData::AAAA(AAAA(addr)))
            }
            RecordData::Cname { target } => {
                let name = Name::from_ascii(target).ok()?;
                Some(RData::CNAME(CNAME(name)))
            }
            RecordData::Mx {
                preference,
                exchange,
            } => {
                let name = Name::from_ascii(exchange).ok()?;
                Some(RData::MX(MX::new(*preference, name)))
            }
            RecordData::Txt { text } => Some(RData::TXT(TXT::new(vec![text.clone()]))),
            RecordData::Ns { nameserver } => {
                let name = Name::from_ascii(nameserver).ok()?;
                Some(RData::NS(NS(name)))
            }
            RecordData::Srv {
                priority,
                weight,
                port,
                target,
            } => {
                let name = Name::from_ascii(target).ok()?;
                Some(RData::SRV(trust_dns_proto::rr::rdata::SRV::new(
                    *priority, *weight, *port, name,
                )))
            }
            RecordData::Ptr { target } => {
                let name = Name::from_ascii(target).ok()?;
                Some(RData::PTR(trust_dns_proto::rr::rdata::PTR(name)))
            }
            RecordData::Caa { flags, tag, value } => {
                // CAA record construction - use the appropriate variant based on tag
                let issuer_critical = (*flags & 0x80) != 0;
                match tag.as_str() {
                    "issue" | "issuewild" => {
                        // For issue/issuewild, value is domain name
                        let name = Name::from_ascii(value).ok();
                        if tag == "issue" {
                            Some(RData::CAA(trust_dns_proto::rr::rdata::CAA::new_issue(
                                issuer_critical,
                                name,
                                vec![],
                            )))
                        } else {
                            Some(RData::CAA(trust_dns_proto::rr::rdata::CAA::new_issuewild(
                                issuer_critical,
                                name,
                                vec![],
                            )))
                        }
                    }
                    "iodef" => {
                        // For iodef, value is URL
                        if let Ok(url) = value.parse() {
                            Some(RData::CAA(trust_dns_proto::rr::rdata::CAA::new_iodef(
                                issuer_critical,
                                url,
                            )))
                        } else {
                            None
                        }
                    }
                    _ => {
                        // Unknown tag - skip for now
                        None
                    }
                }
            }
        }
    }

    /// Convert RecordType to DNS RecordType
    fn to_dns_record_type(&self, rt: RecordType) -> Option<DnsRecordType> {
        match rt {
            RecordType::A => Some(DnsRecordType::A),
            RecordType::Aaaa => Some(DnsRecordType::AAAA),
            RecordType::Cname => Some(DnsRecordType::CNAME),
            RecordType::Mx => Some(DnsRecordType::MX),
            RecordType::Txt => Some(DnsRecordType::TXT),
            RecordType::Ns => Some(DnsRecordType::NS),
            RecordType::Srv => Some(DnsRecordType::SRV),
            RecordType::Ptr => Some(DnsRecordType::PTR),
            RecordType::Caa => Some(DnsRecordType::CAA),
            RecordType::Soa => Some(DnsRecordType::SOA),
        }
    }

    /// Build success response with records
    fn build_success_response(
        &self,
        query: &Message,
        _zone: &Zone,
        records: &[Record],
        qname: &Name,
    ) -> Result<Vec<u8>, DnsError> {
        let mut response = Message::new();
        response.set_id(query.id());
        response.set_message_type(MessageType::Response);
        response.set_op_code(OpCode::Query);
        response.set_authoritative(true);
        response.set_response_code(ResponseCode::NoError);

        // Copy question
        for q in query.queries() {
            response.add_query(q.clone());
        }

        // Add answers
        for record in records {
            if let Some(dns_record) = self.convert_record(record, qname) {
                response.add_answer(dns_record);
            }
        }

        response
            .to_bytes()
            .map_err(|e| DnsError::Encode(format!("Failed to encode response: {}", e)))
    }

    /// Build NXDOMAIN response
    fn build_nxdomain_response(&self, query: &Message, _zone: &Zone) -> Result<Vec<u8>, DnsError> {
        let mut response = Message::new();
        response.set_id(query.id());
        response.set_message_type(MessageType::Response);
        response.set_op_code(OpCode::Query);
        response.set_authoritative(true);
        response.set_response_code(ResponseCode::NXDomain);

        // Copy question
        for q in query.queries() {
            response.add_query(q.clone());
        }

        response
            .to_bytes()
            .map_err(|e| DnsError::Encode(format!("Failed to encode response: {}", e)))
    }

    /// Build error response
    fn build_error_response(
        &self,
        query: &Message,
        rcode: ResponseCode,
    ) -> Result<Vec<u8>, DnsError> {
        let mut response = Message::new();
        response.set_id(query.id());
        response.set_message_type(MessageType::Response);
        response.set_op_code(query.op_code());
        response.set_response_code(rcode);

        // Copy question
        for q in query.queries() {
            response.add_query(q.clone());
        }

        response
            .to_bytes()
            .map_err(|e| DnsError::Encode(format!("Failed to encode response: {}", e)))
    }

    /// Handle PTR query with reverse zone pattern matching
    async fn handle_ptr_query(&self, qname: &Name) -> Option<(String, u32)> {
        // Parse PTR query to IP
        let ip = parse_ptr_query_to_ip(&qname.to_string())?;

        // Find matching reverse zone
        let reverse_zone = self.find_reverse_zone_for_ip(ip).await?;

        // Apply pattern substitution
        let ptr_value = apply_pattern(&reverse_zone.ptr_pattern, ip);

        Some((ptr_value, reverse_zone.ttl))
    }

    /// Find reverse zone that contains the given IP
    async fn find_reverse_zone_for_ip(&self, ip: IpAddr) -> Option<ReverseZone> {
        // List all reverse zones (in production, this would be cached/indexed)
        // For now, scan all orgs
        let zones = self.metadata.list_reverse_zones("*", None).await.ok()?;

        // Find the most specific (longest prefix) match
        let mut best_match: Option<ReverseZone> = None;
        let mut best_prefix_len = 0;

        for zone in zones {
            if let Ok(cidr) = zone.cidr.parse::<IpNet>() {
                if cidr.contains(&ip) {
                    let prefix_len = cidr.prefix_len();
                    if prefix_len > best_prefix_len {
                        best_prefix_len = prefix_len;
                        best_match = Some(zone);
                    }
                }
            }
        }

        best_match
    }

    /// Build PTR response from pattern-generated value
    fn build_ptr_response(
        &self,
        query: &Message,
        qname: &Name,
        ptr_value: &str,
        ttl: u32,
    ) -> Result<Vec<u8>, DnsError> {
        let ptr_name = Name::from_ascii(ptr_value)
            .map_err(|e| DnsError::Parse(format!("Invalid PTR value: {}", e)))?;

        let mut response = Message::new();
        response.set_id(query.id());
        response.set_message_type(MessageType::Response);
        response.set_op_code(OpCode::Query);
        response.set_authoritative(true);
        response.set_response_code(ResponseCode::NoError);

        // Copy question
        for q in query.queries() {
            response.add_query(q.clone());
        }

        // Add PTR answer
        let mut record = DnsRecord::new();
        record.set_name(qname.clone());
        record.set_rr_type(DnsRecordType::PTR);
        record.set_ttl(ttl);
        record.set_data(Some(RData::PTR(trust_dns_proto::rr::rdata::PTR(ptr_name))));
        response.add_answer(record);

        response
            .to_bytes()
            .map_err(|e| DnsError::Encode(format!("Failed to encode PTR response: {}", e)))
    }
}

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

    #[test]
    fn test_error_response() {
        // Minimal DNS query header
        let query = vec![
            0x00, 0x01, // ID
            0x01, 0x00, // Flags: RD
            0x00, 0x01, // QDCOUNT
            0x00, 0x00, // ANCOUNT
            0x00, 0x00, // NSCOUNT
            0x00, 0x00, // ARCOUNT
        ];

        let response = DnsHandler::servfail_response(&query).unwrap();
        assert_eq!(response.len(), 12);
        // Check QR bit is set
        assert_eq!(response[2] & 0x80, 0x80);
        // Check RCODE is SERVFAIL (2)
        assert_eq!(response[3] & 0x0F, 2);
    }
}
