use super::{
    public_roots, roots, AddressPolicy, BackendCertificateVerificationDisabled, TimedConnector,
    CONNECT_TIMEOUT,
};
use crate::config::{Backend, UpstreamProtocol};
use anyhow::{Context, Result};
use arc_swap::ArcSwap;
use axum::body::Body;
use http::{Request, Response};
use hyper::body::Incoming;
use hyper_rustls::HttpsConnectorBuilder;
use hyper_util::{
    client::legacy::{Client, Error as ClientError},
    rt::{TokioExecutor, TokioTimer},
};
use rustls::ClientConfig;
use std::{
    sync::{
        atomic::{AtomicBool, AtomicU64, Ordering},
        Arc, OnceLock,
    },
    time::{Duration, Instant},
};

static MONOTONIC_START: OnceLock<Instant> = OnceLock::new();
type RawProxyClient = Client<super::TimedConnector, Body>;

fn monotonic_millis() -> u64 {
    MONOTONIC_START
        .get_or_init(Instant::now)
        .elapsed()
        .as_millis()
        .min(u64::MAX as u128) as u64
}

#[derive(Clone, PartialEq, Eq)]
struct ClientFactory {
    backend: Backend,
    policy: AddressPolicy,
    websocket: bool,
    keepalive: crate::balance::Keepalive,
    timeout: Duration,
}

struct RotatingClient {
    current: ArcSwap<RawProxyClient>,
    factory: ClientFactory,
    generation: AtomicU64,
    generation_requests: AtomicU64,
    generation_started_ms: AtomicU64,
    rotations: AtomicU64,
    rotating: AtomicBool,
}

/// Backend HTTP client with bounded generation lifetime.
///
/// Hyper reuses a client/pool for efficiency. Under sustained traffic that also
/// makes connection/TLS/buffer high-water allocations live as long as the
/// compiled vhost snapshot. Recycling only this client generation bounds that
/// lifetime without reloading the vhost or touching cache/configuration state.
#[derive(Clone)]
pub struct ProxyClient {
    inner: Arc<RotatingClient>,
}

#[derive(Debug, Clone, Copy)]
pub struct ClientRotationStats {
    pub generation: u64,
    pub generation_requests: u64,
    pub generation_age_seconds: u64,
    pub rotations: u64,
}

impl std::fmt::Debug for ProxyClient {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("ProxyClient")
            .field("rotation", &self.rotation_stats())
            .finish_non_exhaustive()
    }
}

impl ProxyClient {
    fn new(current: RawProxyClient, factory: ClientFactory) -> Self {
        Self {
            inner: Arc::new(RotatingClient {
                current: ArcSwap::from_pointee(current),
                factory,
                generation: AtomicU64::new(0),
                generation_requests: AtomicU64::new(0),
                generation_started_ms: AtomicU64::new(monotonic_millis()),
                rotations: AtomicU64::new(0),
                rotating: AtomicBool::new(false),
            }),
        }
    }

    pub async fn request(&self, request: Request<Body>) -> Result<Response<Incoming>, ClientError> {
        let client = self.inner.client_for_request();
        client.request(request).await
    }

    pub fn rotation_stats(&self) -> ClientRotationStats {
        self.inner.stats()
    }

    /// Whether a newly compiled backend client has exactly the same connection
    /// policy and lifecycle settings as this one. Compatible clients can be
    /// retained across unrelated snapshot rebuilds so generation age and
    /// rotation counters are not restarted from zero.
    pub(crate) fn reusable_with(&self, other: &Self) -> bool {
        self.inner.factory == other.inner.factory
    }

    /// Recycle an aged connection pool even when no user request arrives.
    pub(crate) fn maintain_rotation(&self) {
        self.inner.maintain_rotation();
    }
}

impl RotatingClient {
    fn stats(&self) -> ClientRotationStats {
        let now = monotonic_millis();
        ClientRotationStats {
            generation: self.generation.load(Ordering::Relaxed),
            generation_requests: self.generation_requests.load(Ordering::Relaxed),
            generation_age_seconds: now
                .saturating_sub(self.generation_started_ms.load(Ordering::Relaxed))
                / 1_000,
            rotations: self.rotations.load(Ordering::Relaxed),
        }
    }

    fn client_for_request(&self) -> Arc<RawProxyClient> {
        let observed_generation = self.generation.load(Ordering::Acquire);
        let requests = self
            .generation_requests
            .fetch_add(1, Ordering::Relaxed)
            .saturating_add(1);
        let now = monotonic_millis();
        let age_ms = now.saturating_sub(self.generation_started_ms.load(Ordering::Relaxed));
        let request_due = self.factory.keepalive.pool_recycle_requests > 0
            && requests > self.factory.keepalive.pool_recycle_requests;
        let age_due = self.factory.keepalive.pool_recycle_seconds > 0
            && age_ms
                >= self
                    .factory
                    .keepalive
                    .pool_recycle_seconds
                    .saturating_mul(1_000);
        if request_due || age_due {
            self.rotate(
                observed_generation,
                if request_due { "requests" } else { "age" },
                true,
            );
        }
        self.current.load_full()
    }

    fn maintain_rotation(&self) {
        self.maintain_rotation_at(monotonic_millis());
    }

    fn maintain_rotation_at(&self, now_ms: u64) {
        let recycle_seconds = self.factory.keepalive.pool_recycle_seconds;
        if recycle_seconds == 0 {
            return;
        }
        let age_ms = now_ms.saturating_sub(self.generation_started_ms.load(Ordering::Relaxed));
        if age_ms >= recycle_seconds.saturating_mul(1_000) {
            self.rotate(self.generation.load(Ordering::Acquire), "age", false);
        }
    }

    fn rotate(
        &self,
        observed_generation: u64,
        reason: &'static str,
        triggering_request: bool,
    ) {
        if self.generation.load(Ordering::Acquire) != observed_generation
            || self
                .rotating
                .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
                .is_err()
        {
            return;
        }
        if self.generation.load(Ordering::Acquire) != observed_generation {
            self.rotating.store(false, Ordering::Release);
            return;
        }

        match build_raw_client(&self.factory) {
            Ok(client) => {
                let next_generation = observed_generation.saturating_add(1);
                self.current.store(Arc::new(client));
                // Request-driven rotation sends the triggering request through
                // the new generation. Maintenance-driven rotation happens idle.
                self.generation_requests
                    .store(if triggering_request { 1 } else { 0 }, Ordering::Relaxed);
                self.generation_started_ms
                    .store(monotonic_millis(), Ordering::Relaxed);
                self.generation.store(next_generation, Ordering::Release);
                self.rotations.fetch_add(1, Ordering::Relaxed);
                tracing::debug!(
                    backend = %self.factory.backend.id,
                    generation = next_generation,
                    reason,
                    "recycled upstream client pool"
                );
            }
            Err(error) => {
                // The same immutable settings built the current client. If a
                // future dependency makes rebuilding fail, keep serving through
                // the current generation and retry only after another interval.
                self.generation_requests.store(0, Ordering::Relaxed);
                self.generation_started_ms
                    .store(monotonic_millis(), Ordering::Relaxed);
                tracing::warn!(
                    backend = %self.factory.backend.id,
                    error = %error,
                    "could not recycle upstream client pool; keeping current generation"
                );
            }
        }
        self.rotating.store(false, Ordering::Release);
    }
}

pub fn client(backend: &Backend, policy: AddressPolicy, websocket: bool) -> Result<ProxyClient> {
    client_with_keepalive(
        backend,
        policy,
        websocket,
        &crate::balance::Keepalive::default(),
    )
}

pub fn client_with_keepalive(
    backend: &Backend,
    policy: AddressPolicy,
    websocket: bool,
    keepalive: &crate::balance::Keepalive,
) -> Result<ProxyClient> {
    client_with_timeout(backend, policy, websocket, keepalive, CONNECT_TIMEOUT)
}

pub fn client_with_timeout(
    backend: &Backend,
    policy: AddressPolicy,
    websocket: bool,
    keepalive: &crate::balance::Keepalive,
    timeout: Duration,
) -> Result<ProxyClient> {
    let factory = ClientFactory {
        backend: backend.clone(),
        policy,
        websocket,
        keepalive: keepalive.clone(),
        timeout,
    };
    let current = build_raw_client(&factory)?;
    Ok(ProxyClient::new(current, factory))
}

fn build_raw_client(factory: &ClientFactory) -> Result<RawProxyClient> {
    let backend = &factory.backend;
    let policy = &factory.policy;
    let websocket = factory.websocket;
    let keepalive = &factory.keepalive;
    let timeout = factory.timeout;
    let is_https = url::Url::parse(&backend.url)?.scheme() == "https";
    let protocol = if websocket {
        UpstreamProtocol::Http1
    } else if backend.protocol == UpstreamProtocol::Http2 && is_https {
        UpstreamProtocol::Auto
    } else {
        backend.protocol
    };
    let provider = Arc::new(rustls::crypto::ring::default_provider());
    let trust_roots = if backend.tls_verify {
        roots(backend)?
    } else {
        public_roots()
    };
    let mut config = ClientConfig::builder_with_provider(provider.clone())
        .with_safe_default_protocol_versions()?
        .with_root_certificates(trust_roots)
        .with_no_client_auth();
    if !backend.tls_verify {
        let supported_schemes = provider
            .signature_verification_algorithms
            .supported_schemes();
        config.dangerous().set_certificate_verifier(Arc::new(
            BackendCertificateVerificationDisabled { supported_schemes },
        ));
    }
    let connect_host = if backend.tls_server_name.is_some() {
        Some(
            url::Url::parse(&backend.url)?
                .host_str()
                .context("backend host missing")?
                .trim_start_matches('[')
                .trim_end_matches(']')
                .to_owned(),
        )
    } else {
        None
    };
    let mut http = hyper_util::client::legacy::connect::HttpConnector::new_with_resolver(
        super::GuardedResolver {
            policy: policy.clone(),
            connect_host,
        },
    );
    http.enforce_http(false);
    http.set_connect_timeout(Some(timeout));
    http.set_nodelay(true);
    http.set_keepalive(
        keepalive
            .enabled
            .then_some(Duration::from_secs(keepalive.tcp_keepalive_seconds)),
    );
    let builder = HttpsConnectorBuilder::new()
        .with_tls_config(config)
        .https_or_http();
    let inner = match protocol {
        UpstreamProtocol::Http1 => builder.enable_http1().wrap_connector(http),
        UpstreamProtocol::Http2 => builder.enable_http2().wrap_connector(http),
        UpstreamProtocol::Auto => builder.enable_http1().enable_http2().wrap_connector(http),
    };
    let mut builder = Client::builder(TokioExecutor::new());
    builder
        .http1_title_case_headers(true)
        .http1_max_buf_size(keepalive.http1_max_buffer_bytes)
        .pool_timer(TokioTimer::new())
        .pool_idle_timeout(Duration::from_secs(keepalive.idle_timeout_seconds))
        .pool_max_idle_per_host(if keepalive.enabled {
            keepalive.max_idle_per_host
        } else {
            0
        })
        .retry_canceled_requests(false)
        .http2_only(protocol == UpstreamProtocol::Http2)
        .timer(TokioTimer::new())
        .http2_keep_alive_interval(
            keepalive
                .enabled
                .then_some(Duration::from_secs(keepalive.tcp_keepalive_seconds)),
        )
        .http2_keep_alive_timeout(Duration::from_secs(10))
        .http2_keep_alive_while_idle(false);
    Ok(builder.build(TimedConnector {
        inner,
        policy: policy.clone(),
        timeout,
    }))
}

#[cfg(test)]
include!("client_tests.rs");