pub mod body;
pub mod headers;
pub mod websocket;
use crate::{
    config::request_host,
    runtime::App,
    server::{ConnectionLease, LocalAddr, PeerAddr, ServerShutdown, Transport},
};
use axum::{
    body::Body,
    extract::{Request, State},
    middleware::{self, Next},
    response::{IntoResponse, Response},
    Extension, Router,
};
use body::{Direction, Exchange, MeteredBody};
use http::{
    header::{HeaderValue, HOST, USER_AGENT},
    HeaderMap, Method, StatusCode, Uri, Version,
};
use http_body::Body as HttpBody;
use http_body_util::BodyExt;
#[derive(Clone)]
struct ResponseRulesApplied;
#[derive(Clone)]
struct LocalGeneratedResponse;
#[derive(Clone)]
struct StaticSourceResponse;
#[derive(Clone, Copy)]
struct RedirectResponse {
    http3_enabled: bool,
}
#[derive(Clone, Copy)]
struct LocalHostResponse {
    http3_enabled: bool,
}
#[derive(Clone, Copy)]
struct OriginTiming(Duration);
#[derive(Clone)]
struct OriginBackend(String);
use std::{
    sync::{
        atomic::{AtomicBool, AtomicU16, AtomicU64, Ordering},
        Arc, Mutex,
    },
    time::{Duration, Instant},
};

pub fn router(state: Arc<App>) -> Router {
    Router::new()
        .fallback(forward)
        .layer(middleware::from_fn_with_state(state.clone(), observe))
        .layer(middleware::from_fn_with_state(
            state.clone(),
            crate::acme::challenge,
        ))
        .with_state(state)
}
const MAX_UPSTREAM_ERROR_DETAIL: usize = 2048;

fn truncate_utf8(value: &mut String, max_bytes: usize) {
    if value.len() <= max_bytes {
        return;
    }
    let mut end = max_bytes;
    while end > 0 && !value.is_char_boundary(end) {
        end -= 1;
    }
    value.truncate(end);
}

fn format_error_chain(error: &(dyn std::error::Error + 'static)) -> String {
    let mut parts = Vec::new();
    let mut current = Some(error);
    for _ in 0..12 {
        let Some(item) = current else { break };
        let message = item.to_string();
        if !message.is_empty() && parts.last() != Some(&message) {
            parts.push(message);
        }
        current = item.source();
    }
    if current.is_some() {
        parts.push("<error chain truncated>".into());
    }
    parts.join(" -> ")
}

fn classify_stale_error(
    error: &(dyn std::error::Error + 'static),
) -> crate::cache::StaleErrorClass {
    let mut current = Some(error);
    let mut text = String::new();
    for _ in 0..12 {
        let Some(item) = current else { break };
        if let Some(io) = item.downcast_ref::<std::io::Error>() {
            if io.kind() == std::io::ErrorKind::TimedOut {
                return crate::cache::StaleErrorClass::Timeout;
            }
        }
        if !text.is_empty() {
            text.push(' ');
        }
        text.push_str(&item.to_string().to_ascii_lowercase());
        current = item.source();
    }
    if text.contains("tls")
        || text.contains("certificate")
        || text.contains("cert verify")
        || text.contains("handshake")
    {
        crate::cache::StaleErrorClass::Tls
    } else if text.contains("dns")
        || text.contains("name resolution")
        || text.contains("resolve host")
        || text.contains("lookup address")
    {
        crate::cache::StaleErrorClass::Dns
    } else if text.contains("timed out") || text.contains("timeout") {
        crate::cache::StaleErrorClass::Timeout
    } else {
        crate::cache::StaleErrorClass::Connection
    }
}

fn upstream_error_class(
    error: &(dyn std::error::Error + 'static),
    connect_failure: bool,
) -> &'static str {
    match classify_stale_error(error) {
        crate::cache::StaleErrorClass::Dns => "dns",
        crate::cache::StaleErrorClass::Tls => "tls",
        crate::cache::StaleErrorClass::Timeout => "response_header_timeout",
        crate::cache::StaleErrorClass::Connection if connect_failure => "connect",
        crate::cache::StaleErrorClass::Connection => "io_error",
    }
}

fn retry_safe_method(method: &Method) -> bool {
    method == Method::GET || method == Method::HEAD || method == Method::OPTIONS
}

fn upstream_uri(backend: &crate::config::CompiledBackend, suffix: &str) -> Option<Uri> {
    format!("{}{suffix}", backend.request_origin).parse().ok()
}

fn legacy_websocket_origin(
    request_origin: &str,
    upstream_host: &HeaderValue,
) -> Option<HeaderValue> {
    let scheme = request_origin.split_once("://")?.0;
    let authority = upstream_host.to_str().ok()?.trim();
    if authority.is_empty() {
        return None;
    }
    HeaderValue::from_str(&format!("{scheme}://{authority}")).ok()
}

fn upstream_request_version(
    backend: &crate::config::CompiledBackend,
    websocket_h2: bool,
) -> Version {
    if websocket_h2 || backend.protocol == crate::config::UpstreamProtocol::Http2 {
        Version::HTTP_2
    } else {
        Version::HTTP_11
    }
}

fn retry_request(
    method: &Method,
    version: Version,
    headers: &HeaderMap,
    backend: &crate::config::CompiledBackend,
    original_host: &HeaderValue,
    origin_suffix: &str,
    host_follows_backend: bool,
) -> Option<Request> {
    let uri = upstream_uri(backend, origin_suffix)?;
    let mut request = Request::builder()
        .method(method.clone())
        .version(version)
        .uri(uri)
        .body(Body::empty())
        .ok()?;
    *request.headers_mut() = headers.clone();
    if host_follows_backend {
        request.headers_mut().insert(
            HOST,
            if backend.preserve_host {
                original_host.clone()
            } else {
                backend.host_header.clone()
            },
        );
    }
    Some(request)
}

fn upstream_failure_detail<E>(
    cause: &E,
    backend: &crate::config::CompiledBackend,
    elapsed: Duration,
    connect_failure: bool,
) -> String
where
    E: std::error::Error + std::fmt::Debug + 'static,
{
    let phase = if connect_failure { "connect" } else { "http" };
    let chain = format_error_chain(cause);
    let debug = format!("{cause:?}");
    let mut detail = format!(
        "phase={phase}; origin={}; request_origin={}; tls_verify={}; elapsed_ms={}; error_chain={chain}",
        backend.origin,
        backend.request_origin,
        backend.tls_verify,
        elapsed.as_millis(),
    );
    if !debug.is_empty() && debug != cause.to_string() {
        detail.push_str("; debug=");
        detail.push_str(&debug);
    }
    detail = detail
        .chars()
        .map(|ch| if ch.is_control() { ' ' } else { ch })
        .collect();
    truncate_utf8(&mut detail, MAX_UPSTREAM_ERROR_DETAIL);
    detail
}

#[derive(Clone)]
struct CacheRefreshContext {
    state: Arc<App>,
    snapshot: Arc<crate::config::Snapshot>,
    backend: Arc<crate::config::CompiledBackend>,
    cache_settings: crate::cache::Settings,
    transform: Option<crate::transform::Settings>,
    uri: Uri,
    public_suffix: String,
    host: String,
    path: String,
    scope: String,
    timeout: Duration,
    local: std::net::SocketAddr,
    tls: bool,
    base_headers: http::HeaderMap,
}

async fn refresh_cache_once(
    context: &CacheRefreshContext,
    ticket: crate::cache::Ticket,
    cached: crate::cache::Revalidation,
) -> bool {
    if !context.state.cache.ticket_current(&ticket) {
        return false;
    }
    let mut refresh_headers = context.base_headers.clone();
    for name in [
        "if-none-match",
        "if-modified-since",
        "if-match",
        "if-unmodified-since",
        "if-range",
        "range",
    ] {
        refresh_headers.remove(name);
    }
    if !cached.apply_validators(&mut refresh_headers) {
        return false;
    }
    if context
        .transform
        .as_ref()
        .is_some_and(|settings| settings.enabled)
    {
        refresh_headers.remove("accept-encoding");
    }
    let mut refresh = Request::new(Body::empty());
    *refresh.method_mut() = Method::GET;
    *refresh.uri_mut() = context.uri.clone();
    *refresh.version_mut() = Version::HTTP_11;
    *refresh.headers_mut() = refresh_headers;
    let started = Instant::now();
    let Ok(Ok(origin)) =
        tokio::time::timeout(context.timeout, context.backend.client.request(refresh)).await
    else {
        return false;
    };
    let origin_validation_headers = origin.headers().clone();
    let Some(route) =
        context
            .snapshot
            .select_on(&context.host, &context.path, context.local, context.tls)
    else {
        return false;
    };
    if origin.status() == StatusCode::NOT_MODIFIED {
        let mut validation_headers = origin_validation_headers.clone();
        headers::strip_hop_by_hop(&mut validation_headers);
        apply_response_rules(
            &mut validation_headers,
            &context.snapshot,
            route.frontend.as_deref(),
            Some(route),
            false,
        );
        if let Some(pool) = &route.pool {
            pool.strip_affinity_response(&mut validation_headers);
        }
        let merged = cached.merge_validation_headers(&validation_headers);
        return context
            .state
            .cache
            .publish_revalidated(
                ticket,
                &cached,
                merged,
                &origin_validation_headers,
                &context.cache_settings,
                started.elapsed(),
            )
            .is_some();
    }
    if origin.status() != StatusCode::OK {
        return false;
    }
    let Some(origin_plan) = crate::cache::plan(
        origin.status(),
        origin.headers(),
        &context.cache_settings,
        started.elapsed(),
    ) else {
        return false;
    };
    let origin_no_transform = crate::transform::no_transform(origin.headers());
    let skip_compression = crate::compression::private_or_no_transform(origin.headers());
    let origin_validators = crate::transform::OriginValidators {
        etag: origin.headers().get("etag").cloned(),
        last_modified: origin.headers().get("last-modified").cloned(),
    };
    let (mut parts, body) = origin.into_parts();
    parts.extensions.insert(origin_validators);
    if skip_compression {
        parts.extensions.insert(crate::compression::Skip);
    }
    headers::strip_hop_by_hop(&mut parts.headers);
    apply_response_rules(
        &mut parts.headers,
        &context.snapshot,
        route.frontend.as_deref(),
        Some(route),
        false,
    );
    if let Some(pool) = &route.pool {
        pool.strip_affinity_response(&mut parts.headers);
    }
    let mut refreshed = Response::from_parts(parts, Body::new(body));
    refreshed.extensions_mut().insert(ResponseRulesApplied);
    if let Some(settings) = context
        .transform
        .as_ref()
        .filter(|settings| settings.enabled)
    {
        let transform_started = Instant::now();
        let Ok(transformed) = crate::transform::apply(
            refreshed,
            &context.public_suffix,
            settings,
            origin_no_transform,
        )
        .await
        else {
            context
                .state
                .transform_metrics
                .record_error(&context.scope, transform_started.elapsed());
            context
                .state
                .cache
                .record_reason(&context.scope, &context.host, "transform_error");
            return false;
        };
        context.state.transform_metrics.record_response(
            &context.scope,
            &transformed,
            transform_started.elapsed(),
        );
        refreshed = transformed;
    }
    let canonical_identity = refreshed
        .extensions()
        .get::<crate::transform::CanonicalIdentity>()
        .is_some();
    let Some(transformed_plan) = crate::cache::plan(
        refreshed.status(),
        refreshed.headers(),
        &context.cache_settings,
        started.elapsed(),
    ) else {
        return false;
    };
    let origin_plan = if canonical_identity {
        origin_plan.without_vary("accept-encoding")
    } else {
        origin_plan
    };
    let verify_ticket = ticket.clone();
    let refreshed = context.state.cache.wrap_with_fill(
        refreshed,
        ticket,
        origin_plan.restrict(transformed_plan),
        &context.cache_settings,
        None,
    );
    if http_body_util::BodyExt::collect(refreshed.into_body())
        .await
        .is_err()
    {
        return false;
    }
    context
        .state
        .cache
        .revalidation_tiered(&verify_ticket)
        .await
        .is_some()
}