// Core pool model and policies.
use crate::{config::CompiledBackend, vhost::PoolMember};
use anyhow::{ensure, Result};
use axum::body::Body;
use http::{HeaderMap, HeaderValue, Request};
use rand::{distr::Alphanumeric, RngExt};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::{
    net::IpAddr,
    sync::{
        atomic::{AtomicBool, AtomicU16, AtomicU64, Ordering},
        Arc, Mutex, MutexGuard, Weak,
    },
    time::{Duration, Instant},
};
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;

fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
    mutex.lock().unwrap_or_else(|e| e.into_inner())
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Algorithm {
    RoundRobin,
    #[default]
    LeastConnections,
    LeastTraffic,
    Adaptive,
    IpHash,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum HealthMode {
    Off,
    #[default]
    Tcp,
    Http,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StickyMode {
    #[default]
    Off,
    Cookie,
    IpHash,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct Keepalive {
    pub enabled: bool,
    pub idle_timeout_seconds: u64,
    pub max_idle_per_host: usize,
    pub tcp_keepalive_seconds: u64,
    /// Rebuild a backend Hyper client after this many requests. The counter is
    /// per backend/client generation, not per TCP connection. Zero disables
    /// request-count based recycling.
    pub pool_recycle_requests: u64,
    /// Maximum age of one backend Hyper client generation. Zero disables
    /// time-based recycling.
    pub pool_recycle_seconds: u64,
    /// Bound Hyper's adaptive HTTP/1 read buffer per connection. Hyper requires
    /// at least 8192 bytes. This affects upstream HTTP/1 only.
    pub http1_max_buffer_bytes: usize,
}
impl Default for Keepalive {
    fn default() -> Self {
        Self {
            enabled: true,
            idle_timeout_seconds: 30,
            max_idle_per_host: 4,
            tcp_keepalive_seconds: 30,
            pool_recycle_requests: 25_000,
            pool_recycle_seconds: 3_600,
            http1_max_buffer_bytes: 131_072,
        }
    }
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct RetryPolicy {
    pub max_retries: u8,
    pub retry_connect_errors: bool,
    pub retry_timeouts: bool,
    pub retry_statuses: Vec<u16>,
}
impl Default for RetryPolicy {
    fn default() -> Self {
        Self {
            max_retries: 1,
            retry_connect_errors: true,
            retry_timeouts: true,
            retry_statuses: vec![502, 503, 504],
        }
    }
}
impl RetryPolicy {
    pub fn retries_status(&self, status: u16) -> bool {
        self.max_retries > 0 && self.retry_statuses.contains(&status)
    }
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct HealthCheck {
    pub mode: HealthMode,
    pub path: String,
    pub host: Option<String>,
    pub interval_ms: u64,
    pub timeout_ms: u64,
    pub slow_response_ms: u64,
    pub fall: u16,
    pub rise: u16,
    pub expected_status_min: u16,
    pub expected_status_max: u16,
}
impl Default for HealthCheck {
    fn default() -> Self {
        Self {
            mode: HealthMode::Tcp,
            path: "/".into(),
            host: None,
            interval_ms: 5000,
            timeout_ms: 2000,
            slow_response_ms: 1500,
            fall: 3,
            rise: 2,
            expected_status_min: 200,
            expected_status_max: 399,
        }
    }
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct PassiveHealth {
    pub enabled: bool,
    pub count_connect_errors: bool,
    pub count_timeouts: bool,
    pub count_io_errors: bool,
    pub count_http_5xx: bool,
    pub count_slow_responses: bool,
    pub fall: Option<u16>,
    pub rise: Option<u16>,
}
impl Default for PassiveHealth {
    fn default() -> Self {
        Self {
            enabled: false,
            count_connect_errors: true,
            count_timeouts: true,
            count_io_errors: true,
            count_http_5xx: true,
            count_slow_responses: false,
            fall: None,
            rise: None,
        }
    }
}
impl PassiveHealth {
    fn fall(&self, active: &HealthCheck) -> u16 {
        self.fall.unwrap_or(active.fall)
    }
    fn rise(&self, active: &HealthCheck) -> u16 {
        self.rise.unwrap_or(active.rise)
    }
    fn counts(&self, reason: &str) -> bool {
        match reason {
            "connect_error" => self.count_connect_errors,
            "timeout" => self.count_timeouts,
            "io_error" => self.count_io_errors,
            "http_5xx" => self.count_http_5xx,
            "slow_response" => self.count_slow_responses,
            _ => false,
        }
    }
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct Sticky {
    pub mode: StickyMode,
    pub cookie_name: Option<String>,
    pub ttl_seconds: u64,
}
impl Default for Sticky {
    fn default() -> Self {
        Self {
            mode: StickyMode::Off,
            cookie_name: None,
            ttl_seconds: 3600,
        }
    }
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct Settings {
    pub algorithm: Algorithm,
    pub health_check: HealthCheck,
    pub passive_health: PassiveHealth,
    pub keepalive: Keepalive,
    pub retry: RetryPolicy,
    pub sticky: Sticky,
    pub slow_start_seconds: u64,
}
impl Default for Settings {
    fn default() -> Self {
        Self {
            algorithm: Algorithm::default(),
            health_check: HealthCheck::default(),
            passive_health: PassiveHealth::default(),
            keepalive: Keepalive::default(),
            retry: RetryPolicy::default(),
            sticky: Sticky::default(),
            slow_start_seconds: 10,
        }
    }
}
impl Settings {
    /// Resolve active health per member; a disabled monitor never probes.
    pub fn health_for(&self, member: &PoolMember) -> HealthCheck {
        let mut health = member
            .health_check
            .clone()
            .unwrap_or_else(|| self.health_check.clone());
        if health.host.is_none() && health.mode == HealthMode::Http {
            health.host = self.health_check.host.clone();
        }
        if !member.health_check_enabled {
            health.mode = HealthMode::Off;
        }
        health
    }
    /// Monitoring off also excludes passive ejection; failures remain logged.
    pub fn passive_for(&self, member: &PoolMember) -> PassiveHealth {
        let mut passive = self.passive_health.clone();
        if !member.health_check_enabled {
            passive.enabled = false;
        }
        passive
    }
    pub fn validate(&self) -> Result<()> {
        let h = &self.health_check;
        let k = &self.keepalive;
        ensure!(
            (250..=300_000).contains(&h.interval_ms),
            "health interval_ms must be 250..300000"
        );
        ensure!(
            (50..=60_000).contains(&h.timeout_ms) && h.timeout_ms <= h.interval_ms,
            "health timeout_ms must be 50..60000 and <= interval_ms"
        );
        ensure!(
            h.slow_response_ms <= h.timeout_ms,
            "slow_response_ms must be 0 (off) or <= timeout_ms"
        );
        ensure!(
            (1..=20).contains(&h.fall) && (1..=20).contains(&h.rise),
            "health fall/rise must be 1..20"
        );
        ensure!(
            h.expected_status_min >= 100
                && h.expected_status_max <= 599
                && h.expected_status_min <= h.expected_status_max,
            "invalid health HTTP status range"
        );
        ensure!(
            h.path.starts_with('/')
                && !h.path.starts_with("//")
                && h.path.len() <= 1024
                && !h.path.contains('#')
                && !h.path.bytes().any(|b| b <= 32 || b == 127),
            "invalid health HTTP path"
        );
        if let Some(host) = &h.host {
            crate::config::request_host(host)?;
            let _: HeaderValue = host.parse()?;
        }
        if let Some(fall) = self.passive_health.fall {
            ensure!(
                (1..=20).contains(&fall),
                "passive health fall must be 1..20"
            );
        }
        if let Some(rise) = self.passive_health.rise {
            ensure!(
                (1..=20).contains(&rise),
                "passive health rise must be 1..20"
            );
        }
        ensure!(
            (1..=3600).contains(&k.idle_timeout_seconds)
                && (1..=3600).contains(&k.tcp_keepalive_seconds)
                && (1..=64).contains(&k.max_idle_per_host)
                && (k.pool_recycle_requests == 0
                    || (1_000..=10_000_000).contains(&k.pool_recycle_requests))
                && (k.pool_recycle_seconds == 0
                    || (60..=604_800).contains(&k.pool_recycle_seconds))
                && (8_192..=1_048_576).contains(&k.http1_max_buffer_bytes),
            "invalid keepalive limits"
        );
        ensure!(
            self.retry.max_retries <= 3,
            "retry max_retries must be 0..3"
        );
        ensure!(
            self.retry.retry_statuses.len() <= 8
                && self
                    .retry
                    .retry_statuses
                    .iter()
                    .all(|status| (500..=599).contains(status)),
            "retry_statuses must contain at most 8 HTTP 5xx status codes"
        );
        ensure!(
            self.slow_start_seconds <= 300,
            "slow_start_seconds must be 0..300"
        );
        ensure!(
            (60..=604_800).contains(&self.sticky.ttl_seconds),
            "sticky ttl_seconds must be 60..604800"
        );
        if let Some(name) = &self.sticky.cookie_name {
            ensure!(!name.is_empty() && name.len() <= 96 && name.bytes().all(|b| b.is_ascii_alphanumeric() || matches!(b,b'_'|b'-')),
                "affinity cookie name must use 1..96 ASCII letters, numbers, underscores or hyphens");
            ensure!(
                !name.eq_ignore_ascii_case(crate::auth::COOKIE_NAME),
                "affinity cannot use the administrator session cookie name"
            );
        }
        Ok(())
    }
}
#[derive(Debug)]
struct Health {
    healthy: bool,
    active_healthy: bool,
    passive_healthy: bool,
    checked: bool,
    failures: u16,
    passive_failures: u16,
    successes: u16,
    passive_successes: u16,
    next_probe: Instant,
    probing: bool,
    healthy_since: Instant,
    result: &'static str,
    passive_result: &'static str,
    last_probe_ms: u64,
    latency_us: u64,
    ewma_us: u64,
    transitions: u64,
    unhealthy_source: &'static str,
    unhealthy_reason: &'static str,
    active_unhealthy_source: &'static str,
    active_unhealthy_reason: &'static str,
    passive_unhealthy_source: &'static str,
    passive_unhealthy_reason: &'static str,
    passive_first_failure: Option<FailureObservation>,
    active_first_failure: Option<FailureObservation>,
    passive_probe_successes: u16,
    passive_ejected_at: Option<Instant>,
    passive_cooldown: Duration,
    last_io_counted_at: Option<Instant>,
}

#[derive(Debug)]
struct Traffic {
    at: Instant,
    total: u64,
    rate: f64,
}
#[derive(Debug, Clone)]
struct FailureObservation {
    at_ms: u64,
    class: &'static str,
    detail: String,
}
impl Default for FailureObservation {
    fn default() -> Self {
        Self {
            at_ms: 0,
            class: "none",
            detail: String::new(),
        }
    }
}
#[derive(Debug)]
pub struct BackendState {
    health: Mutex<Health>,
    traffic: Mutex<Traffic>,
    last_failure: Mutex<FailureObservation>,
    pub active: AtomicU64,
    pub requests: AtomicU64,
    pub errors: AtomicU64,
    pub bytes_in: AtomicU64,
    pub bytes_out: AtomicU64,
    pub last_status: AtomicU16,
    pub active_health_failures_total: AtomicU64,
    pub passive_health_failures_total: AtomicU64,
    pub http_5xx_total: AtomicU64,
    pub timeouts_total: AtomicU64,
    pub connect_errors_total: AtomicU64,
    pub io_errors_total: AtomicU64,
    pub slow_responses_total: AtomicU64,
    pub retry_attempts_total: AtomicU64,
    pub retry_successes_total: AtomicU64,
    pub capacity_rejections_total: AtomicU64,
    pub queue_timeouts_total: AtomicU64,
    pub dns_errors_total: AtomicU64,
    pub tls_errors_total: AtomicU64,
    pub response_header_timeouts_total: AtomicU64,
    pub read_timeouts_total: AtomicU64,
    pub send_timeouts_total: AtomicU64,
    pub last_ejected_ms: AtomicU64,
    pub last_recovered_ms: AtomicU64,
}