use axum::extract::Query;
use axum::http::{header, StatusCode};
use axum::response::{IntoResponse, Response};
use axum::Extension;
use base64::Engine;
use std::sync::Arc;
use tracing::{error, info, warn};

use crate::dns::query::parse_query;
use crate::dns::response::build_servfail;
use crate::upstream::doh_client::{DohClient, Strategy};

/// Shared application state passed to route handlers.
#[derive(Clone)]
pub struct AppState {
    /// The DoH upstream client.
    pub doh_client: DohClient,
    /// Counter for upstream selection (round-robin index / failover start).
    pub rr_counter: Arc<std::sync::atomic::AtomicUsize>,
    /// Upstream selection strategy.
    pub strategy: Strategy,
}

/// Query parameters for the GET variant of `/dns-query`.
#[derive(Debug, serde::Deserialize)]
pub struct DnsQueryParams {
    /// Base64url-encoded DNS wire-format query (RFC 8484 §4.1).
    pub dns: String,
}

/// POST handler for `/dns-query` (RFC 8484 §4.2).
///
/// Expects `Content-Type: application/dns-message` and a raw DNS
/// wire-format body.
pub async fn dns_query_post(Extension(state): Extension<AppState>, body: bytes::Bytes) -> Response {
    let query_bytes = body.to_vec();

    // Validate that it's a parseable DNS message
    let query = match parse_query(&query_bytes) {
        Ok(q) => q,
        Err(e) => {
            warn!(error = %e, "invalid DNS query received");
            return (StatusCode::BAD_REQUEST, "invalid DNS query").into_response();
        }
    };

    let qname = query
        .queries()
        .first()
        .map(|q| q.name().to_string())
        .unwrap_or_default();
    let qtype = query
        .queries()
        .first()
        .map(|q| q.query_type().to_string())
        .unwrap_or_default();

    info!(query = %qname, qtype = %qtype, "received DoH POST query");

    forward_and_respond(&state, &query_bytes, &query).await
}

/// GET handler for `/dns-query` (RFC 8484 §4.1).
///
/// Expects a `?dns=<base64url>` query parameter.
pub async fn dns_query_get(
    Extension(state): Extension<AppState>,
    Query(params): Query<DnsQueryParams>,
) -> Response {
    use base64::engine::general_purpose::URL_SAFE_NO_PAD;

    let query_bytes = match URL_SAFE_NO_PAD.decode(&params.dns) {
        Ok(b) => b,
        Err(e) => {
            warn!(error = %e, "invalid base64url in dns query param");
            return (StatusCode::BAD_REQUEST, "invalid base64url dns param").into_response();
        }
    };

    let query = match parse_query(&query_bytes) {
        Ok(q) => q,
        Err(e) => {
            warn!(error = %e, "invalid DNS query from base64 param");
            return (StatusCode::BAD_REQUEST, "invalid DNS query").into_response();
        }
    };

    let qname = query
        .queries()
        .first()
        .map(|q| q.name().to_string())
        .unwrap_or_default();
    let qtype = query
        .queries()
        .first()
        .map(|q| q.query_type().to_string())
        .unwrap_or_default();

    info!(query = %qname, qtype = %qtype, "received DoH GET query");

    forward_and_respond(&state, &query_bytes, &query).await
}

/// Forward a query to the upstream and build the response.
async fn forward_and_respond(
    state: &AppState,
    query_bytes: &[u8],
    query: &hickory_proto::op::Message,
) -> Response {
    let idx = state
        .rr_counter
        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);

    match state
        .doh_client
        .forward(query_bytes, state.strategy, idx)
        .await
    {
        Ok(response_bytes) => {
            info!(bytes = response_bytes.len(), "forwarding upstream response");
            (
                StatusCode::OK,
                [(header::CONTENT_TYPE, "application/dns-message")],
                response_bytes,
            )
                .into_response()
        }
        Err(e) => {
            error!(error = %e, "upstream forward failed, returning SERVFAIL");
            match build_servfail(query) {
                Ok(servfail) => (
                    StatusCode::OK,
                    [(header::CONTENT_TYPE, "application/dns-message")],
                    servfail,
                )
                    .into_response(),
                Err(e) => {
                    error!(error = %e, "failed to build SERVFAIL response");
                    (StatusCode::INTERNAL_SERVER_ERROR, "internal error").into_response()
                }
            }
        }
    }
}
