use reqwest::Client;
use tracing::{debug, warn};

use crate::config::ResolverEntry;

/// Errors from the upstream DoH client.
#[derive(Debug, thiserror::Error)]
pub enum UpstreamError {
    /// An HTTP-level error occurred.
    #[error("HTTP error from upstream {name}: {source}")]
    HttpError {
        name: String,
        source: reqwest::Error,
    },
    /// The upstream returned a non-2xx status.
    #[error("upstream {name} returned status {status}")]
    BadStatus { name: String, status: u16 },
}

/// Upstream selection strategy.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Strategy {
    /// Distribute queries across resolvers in order.
    RoundRobin,
    /// Try resolvers in order until one succeeds.
    Failover,
}

/// A DoH client that forwards DNS wire-format queries to upstream resolvers.
#[derive(Clone)]
pub struct DohClient {
    http: Client,
    resolvers: Vec<ResolverEntry>,
}

impl DohClient {
    /// Create a new DoH client with the given upstream resolvers.
    pub fn new(resolvers: Vec<ResolverEntry>) -> Self {
        let http = Client::builder()
            .connect_timeout(std::time::Duration::from_secs(3))
            .timeout(std::time::Duration::from_secs(5))
            .build()
            .expect("failed to build HTTP client");
        Self { http, resolvers }
    }

    /// Forward a DNS wire-format query using the given strategy.
    ///
    /// Both strategies retry with the next resolver on failure.
    /// For round-robin, `index` picks the starting resolver.
    /// For failover, resolvers are tried in order starting from `index`.
    pub async fn forward(
        &self,
        query_bytes: &[u8],
        strategy: Strategy,
        index: usize,
    ) -> Result<Vec<u8>, UpstreamError> {
        match strategy {
            Strategy::RoundRobin => self.forward_with_retry(query_bytes, index).await,
            Strategy::Failover => self.forward_with_retry(query_bytes, index).await,
        }
    }

    /// Try resolvers starting from `start_index`, retrying with the next
    /// on failure. This provides resilience for both strategies.
    async fn forward_with_retry(
        &self,
        query_bytes: &[u8],
        start_index: usize,
    ) -> Result<Vec<u8>, UpstreamError> {
        let n = self.resolvers.len();
        let mut last_err = None;

        for i in 0..n {
            let idx = (start_index + i) % n;
            let resolver = &self.resolvers[idx];
            match self.send_to(query_bytes, resolver).await {
                Ok(bytes) => return Ok(bytes),
                Err(e) => {
                    warn!(resolver = %resolver.name, error = %e, "resolver failed, trying next");
                    last_err = Some(e);
                }
            }
        }

        Err(last_err.expect("no resolvers configured"))
    }

    /// Send a query to a single resolver.
    async fn send_to(
        &self,
        query_bytes: &[u8],
        resolver: &ResolverEntry,
    ) -> Result<Vec<u8>, UpstreamError> {
        debug!(resolver = %resolver.name, url = %resolver.url, "forwarding query to upstream");

        let response = self
            .http
            .post(&resolver.url)
            .header("content-type", "application/dns-message")
            .header("accept", "application/dns-message")
            .body(query_bytes.to_vec())
            .send()
            .await
            .map_err(|e| UpstreamError::HttpError {
                name: resolver.name.clone(),
                source: e,
            })?;

        let status = response.status();
        if !status.is_success() {
            warn!(resolver = %resolver.name, status = %status, "upstream returned error status");
            return Err(UpstreamError::BadStatus {
                name: resolver.name.clone(),
                status: status.as_u16(),
            });
        }

        let body = response
            .bytes()
            .await
            .map_err(|e| UpstreamError::HttpError {
                name: resolver.name.clone(),
                source: e,
            })?;

        debug!(resolver = %resolver.name, bytes = body.len(), "received upstream response");
        Ok(body.to_vec())
    }

    /// Return the number of configured resolvers.
    pub fn resolver_count(&self) -> usize {
        self.resolvers.len()
    }
}

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

    #[test]
    fn test_client_creation() {
        let resolvers = vec![
            ResolverEntry {
                name: "cloudflare".into(),
                url: "https://cloudflare-dns.com/dns-query".into(),
            },
            ResolverEntry {
                name: "google".into(),
                url: "https://dns.google/dns-query".into(),
            },
        ];
        let client = DohClient::new(resolvers);
        assert_eq!(client.resolver_count(), 2);
    }
}
