//! Passes frames through unchanged, including trailers. No body collection.
use crate::{
    logging::{now_ms, LogEntry, Logs},
    metrics::Metrics,
};
use axum::body::Body;
use bytes::Bytes;
use http_body::{Body as HttpBody, Frame, SizeHint};
use std::{
    future::Future,
    pin::Pin,
    sync::{
        atomic::{AtomicBool, AtomicU16, AtomicU64, AtomicU8, Ordering},
        Arc, Mutex,
    },
    task::{Context, Poll},
    time::{Duration, Instant},
};
use tokio::{
    sync::OwnedSemaphorePermit,
    time::{sleep_until, Sleep},
};

/// Compact source tags avoid a mutex on every completed exchange.
#[repr(u8)]
#[derive(Clone, Copy)]
pub enum ExchangeSource {
    Unresolved = 0,
    LocalResponse = 1,
    LocalRedirect = 2,
    WebServer = 3,
    ReverseProxy = 4,
}
impl ExchangeSource {
    fn name(value: u8) -> &'static str {
        match value {
            1 => "local_response",
            2 => "local_redirect",
            3 => "web_server",
            4 => "reverse_proxy",
            _ => "unresolved",
        }
    }
}

// 0 = unknown, 1 = Some(false), 2 = Some(true). Preserve log schema.
fn read_effective(value: u8) -> Option<bool> {
    match value {
        1 => Some(false),
        2 => Some(true),
        _ => None,
    }
}

pub struct Exchange {
    pub vhost_metrics: Option<Arc<crate::metrics::VhostMetrics>>,
    pub vhost: Option<String>,
    pub route: Mutex<Option<String>>,
    pub source: AtomicU8,
    pub cache_effective_enabled: AtomicU8,
    pub transform_effective_enabled: AtomicU8,
    pub backend: Mutex<Option<Arc<crate::balance::Lease>>>,
    pub id: u64,
    pub request_id: String,
    pub method: String,
    pub host: String,
    pub path: String,
    pub http_version: &'static str,
    pub transport: &'static str,
    pub websocket_bootstrap: Option<&'static str>,
    pub websocket_phase: Mutex<Option<&'static str>>,
    pub websocket_error: Mutex<Option<&'static str>>,
    pub h3_reset_code: Mutex<Option<&'static str>>,
    pub upstream_status: AtomicU16,
    pub upstream_error: Mutex<Option<String>>,
    pub client_ip: Mutex<Option<String>>,
    pub user_agent: Option<String>,
    pub country: Mutex<Option<String>>,
    pub started: Instant,
    pub status: AtomicU16,
    pub cache_status: AtomicU16,
    pub cache_reason: Mutex<&'static str>,
    pub failed: AtomicBool,
    pub aborted: AtomicBool,
    pub request_bytes: AtomicU64,
    pub response_bytes: AtomicU64,
    pub metrics: Arc<Metrics>,
    pub logs: Arc<Logs>,
    pub traffic: Arc<crate::traffic::Analytics>,
    pub _permit: OwnedSemaphorePermit,
}
impl Exchange {
    pub fn set_source(&self, source: ExchangeSource) {
        self.source.store(source as u8, Ordering::Relaxed);
    }
    pub fn set_cache_enabled(&self, enabled: bool) {
        self.cache_effective_enabled
            .store(if enabled { 2 } else { 1 }, Ordering::Relaxed);
    }
    pub fn set_transform_enabled(&self, enabled: bool) {
        self.transform_effective_enabled
            .store(if enabled { 2 } else { 1 }, Ordering::Relaxed);
    }
    pub fn websocket_phase(&self, phase: &'static str) {
        if self.websocket_bootstrap.is_some() {
            *self
                .websocket_phase
                .lock()
                .unwrap_or_else(|e| e.into_inner()) = Some(phase);
        }
    }
    pub fn websocket_error(&self, phase: &'static str, error: &'static str) {
        if self.websocket_bootstrap.is_some() {
            *self
                .websocket_phase
                .lock()
                .unwrap_or_else(|e| e.into_inner()) = Some(phase);
            *self
                .websocket_error
                .lock()
                .unwrap_or_else(|e| e.into_inner()) = Some(error);
        }
    }
    pub fn h3_reset(&self, code: &'static str) {
        *self.h3_reset_code.lock().unwrap_or_else(|e| e.into_inner()) = Some(code);
    }
    pub fn backend_bytes(&self, bytes: u64, incoming: bool) {
        if let Some(backend) = self
            .backend
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .as_ref()
        {
            backend.bytes(bytes, incoming);
        }
    }
    fn backend_error(&self, reason: &'static str) {
        if let Some(backend) = self
            .backend
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .as_ref()
        {
            backend.error(reason, self.started.elapsed());
        }
    }
}
impl Drop for Exchange {
    fn drop(&mut self) {
        let status = self.status.load(Ordering::Relaxed);
        let failed = self.failed.load(Ordering::Relaxed);
        let aborted = self.aborted.load(Ordering::Relaxed);
        let duration = self.started.elapsed();
        if let Some(m) = &self.vhost_metrics {
            m.complete(
                status,
                failed || aborted,
                duration,
                self.request_bytes.load(Ordering::Relaxed),
                self.response_bytes.load(Ordering::Relaxed),
            );
        }
        self.metrics.active.fetch_sub(1, Ordering::Relaxed);
        self.metrics.completed.fetch_add(1, Ordering::Relaxed);
        self.metrics
            .duration_us
            .fetch_add(duration.as_micros() as u64, Ordering::Relaxed);
        if failed || aborted || status >= 500 {
            self.metrics.errors.fetch_add(1, Ordering::Relaxed);
        }
        match status / 100 {
            1 => &self.metrics.status_1xx,
            2 => &self.metrics.status_2xx,
            3 => &self.metrics.status_3xx,
            4 => &self.metrics.status_4xx,
            _ => &self.metrics.status_5xx,
        }
        .fetch_add(1, Ordering::Relaxed);
        // Final Arc drop has exclusive access: move strings without locking.
        let backend = self
            .backend
            .get_mut()
            .unwrap_or_else(|e| e.into_inner())
            .as_ref()
            .map(|l| l.member.backend.id.clone());
        let websocket_phase = *self.websocket_phase.get_mut().unwrap_or_else(|e| e.into_inner());
        let route = self.route.get_mut().unwrap_or_else(|e| e.into_inner()).take();
        let source = ExchangeSource::name(self.source.load(Ordering::Relaxed));
        let cache_effective_enabled = read_effective(self.cache_effective_enabled.load(Ordering::Relaxed));
        let transform_effective_enabled = read_effective(self.transform_effective_enabled.load(Ordering::Relaxed));
        let client_ip = self
            .client_ip
            .get_mut()
            .unwrap_or_else(|e| e.into_inner())
            .take();
        let country = self
            .country
            .get_mut()
            .unwrap_or_else(|e| e.into_inner())
            .take();
        let request_bytes = self.request_bytes.load(Ordering::Relaxed);
        let response_bytes = self.response_bytes.load(Ordering::Relaxed);
        let duration_ms = duration.as_millis() as u64;
        let cache = match self.cache_status.load(Ordering::Relaxed) {
            1 => "HIT",
            2 => "MISS",
            3 => "BYPASS",
            _ => "-",
        };
        let outcome = if failed {
            "body_error"
        } else if aborted {
            "cancelled"
        } else if status >= 500 {
            "proxy_or_upstream_error"
        } else if status == 101 || websocket_phase == Some("tunnel_closed") {
            "websocket_closed"
        } else {
            "complete"
        };
        if let Some(vhost) = self.vhost.as_ref() {
            self.traffic.push_lazy(|| crate::traffic::Sample {
                vhost: vhost.clone(),
                route: route.clone(),
                method: self.method.clone(),
                path: self.path.clone(),
                client_ip: client_ip.clone(),
                status,
                duration_ms,
                request_bytes,
                response_bytes,
                cache,
                source,
                user_agent: self.user_agent.clone(),
                country: country.clone(),
            });
        }
        self.logs.push(LogEntry {
            event: "request",
            protocol_error: None,
            reason: None,
            error_detail: None,
            error: None,
            local_addr: None,
            tls_sni: None,
            sni: None,
            authority: None,
            matched_by: None,
            cache_reason: *self.cache_reason.get_mut().unwrap_or_else(|e| e.into_inner()),
            vhost: self.vhost.take(),
            route,
            cache_effective_enabled,
            transform_effective_enabled,
            backend,
            id: self.id,
            request_id: std::mem::take(&mut self.request_id),
            time_ms: now_ms(),
            method: std::mem::take(&mut self.method),
            host: std::mem::take(&mut self.host),
            path: std::mem::take(&mut self.path),
            http_version: self.http_version,
            transport: self.transport,
            websocket_bootstrap: self.websocket_bootstrap,
            websocket_phase,
            websocket_error: *self
                .websocket_error
                .get_mut()
                .unwrap_or_else(|e| e.into_inner()),
            h3_reset_code: *self.h3_reset_code.get_mut().unwrap_or_else(|e| e.into_inner()),
            upstream_status: match self.upstream_status.load(Ordering::Relaxed) {
                0 => None,
                value => Some(value),
            },
            upstream_error: self
                .upstream_error
                .get_mut()
                .unwrap_or_else(|e| e.into_inner())
                .take(),
            client_ip,
            user_agent: self.user_agent.take(),
            country,
            status,
            duration_ms,
            cache,
            request_bytes,
            response_bytes,
            outcome,
        });
    }
}

#[derive(Clone, Copy)]
pub enum Direction {
    Request,
    Response,
}
pub struct MeteredBody {
    inner: Body,
    exchange: Option<Arc<Exchange>>,
    /// For a response, backend selection is final before this wrapper is made.
    /// Keep the selected lease alive and avoid locking Exchange.backend per frame.
    response_backend: Option<Arc<crate::balance::Lease>>,
    request_backend: Option<Arc<crate::balance::Lease>>,
    direction: Direction,
    timer: Option<Pin<Box<Sleep>>>,
    deadline: tokio::time::Instant,
    idle: Duration,
    finished: bool,
}
impl MeteredBody {
    pub fn new(inner: Body, exchange: Arc<Exchange>, direction: Direction, idle: Duration) -> Self {
        let finished = inner.is_end_stream();
        let response_backend = matches!(direction, Direction::Response)
            .then(|| exchange.backend.lock().unwrap_or_else(|e| e.into_inner()).clone())
            .flatten();
        Self {
            inner,
            exchange: Some(exchange),
            response_backend,
            request_backend: None,
            direction,
            // Ready/empty bodies never need a timer allocation. Keep the
            // original deadline so a delayed first poll still times out.
            timer: None,
            deadline: tokio::time::Instant::now() + idle,
            idle,
            finished,
        }
    }
    fn finish(&mut self) {
        self.finished = true;
        self.timer.take();
        self.response_backend.take();
        self.request_backend.take();
        self.exchange.take();
    }
}
impl HttpBody for MeteredBody {
    type Data = Bytes;
    type Error = axum::Error;
    fn poll_frame(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<Frame<Bytes>, Self::Error>>> {
        let this = self.get_mut();
        if this.finished {
            this.exchange.take();
            return Poll::Ready(None);
        }
        match Pin::new(&mut this.inner).poll_frame(cx) {
            Poll::Ready(Some(Ok(frame))) => {
                if let (Some(data), Some(exchange)) = (frame.data_ref(), this.exchange.as_ref()) {
                    let count = data.len() as u64;
                    match this.direction {
                        Direction::Request => {
                            // Backend is selected before a non-empty upstream
                            // upload is polled. No retry of a streamed write.
                            if this.request_backend.is_none() {
                                this.request_backend = exchange.backend
                                    .lock().unwrap_or_else(|e| e.into_inner()).clone();
                            }
                            if let Some(lease) = &this.request_backend {
                                lease.bytes(count, true);
                            }
                        }
                        Direction::Response => {
                            if let Some(lease) = &this.response_backend {
                                lease.bytes(count, false);
                            }
                        }
                    }
                    match this.direction {
                        Direction::Request => {
                            exchange.request_bytes.fetch_add(count, Ordering::Relaxed);
                            exchange
                                .metrics
                                .bytes_in
                                .fetch_add(count, Ordering::Relaxed);
                        }
                        Direction::Response => {
                            exchange.response_bytes.fetch_add(count, Ordering::Relaxed);
                            exchange
                                .metrics
                                .bytes_out
                                .fetch_add(count, Ordering::Relaxed);
                        }
                    }
                }
                if this.inner.is_end_stream() {
                    this.finish();
                } else {
                    this.deadline = tokio::time::Instant::now() + this.idle;
                    if let Some(timer) = this.timer.as_mut() {
                        timer.as_mut().reset(this.deadline);
                    }
                }
                Poll::Ready(Some(Ok(frame)))
            }
            Poll::Ready(Some(Err(error))) => {
                if let Some(exchange) = &this.exchange {
                    exchange.failed.store(true, Ordering::Relaxed);
                    exchange.backend_error(match this.direction {
                        Direction::Request => "send_error",
                        Direction::Response => "read_error",
                    });
                }
                this.finish();
                Poll::Ready(Some(Err(error)))
            }
            Poll::Ready(None) => {
                this.finish();
                Poll::Ready(None)
            }
            Poll::Pending => {
                let timer = this
                    .timer
                    .get_or_insert_with(|| Box::pin(sleep_until(this.deadline)));
                if timer.as_mut().poll(cx).is_ready() {
                    if let Some(exchange) = &this.exchange {
                        exchange.failed.store(true, Ordering::Relaxed);
                        exchange.backend_error(match this.direction {
                            Direction::Request => "send_timeout",
                            Direction::Response => "read_timeout",
                        });
                    }
                    this.finish();
                    Poll::Ready(Some(Err(axum::Error::new(std::io::Error::new(
                        std::io::ErrorKind::TimedOut,
                        "body idle timeout",
                    )))))
                } else {
                    Poll::Pending
                }
            }
        }
    }
    fn is_end_stream(&self) -> bool {
        self.finished || self.inner.is_end_stream()
    }
    fn size_hint(&self) -> SizeHint {
        self.inner.size_hint()
    }
}
impl Drop for MeteredBody {
    fn drop(&mut self) {
        if !self.finished
            && !self.inner.is_end_stream()
            && matches!(self.direction, Direction::Response)
        {
            if let Some(exchange) = &self.exchange {
                exchange.aborted.store(true, Ordering::Relaxed);
            }
        }
    }
}

#[cfg(test)]
include!("body/tests.rs");

/// Retains the first inspected frame without collecting the rest of a request.
/// A trailer/empty DATA frame still prevents storing a request-body variant.
pub struct ReplayBody {
    head: Option<http_body::Frame<bytes::Bytes>>,
    tail: axum::body::Body,
}
impl ReplayBody {
    pub fn new(head: http_body::Frame<bytes::Bytes>, tail: axum::body::Body) -> Self {
        Self {
            head: Some(head),
            tail,
        }
    }
}
impl http_body::Body for ReplayBody {
    type Data = bytes::Bytes;
    type Error = axum::Error;
    fn poll_frame(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
        if let Some(head) = self.head.take() {
            return std::task::Poll::Ready(Some(Ok(head)));
        }
        std::pin::Pin::new(&mut self.tail).poll_frame(cx)
    }
    fn is_end_stream(&self) -> bool {
        self.head.is_none() && self.tail.is_end_stream()
    }
    fn size_hint(&self) -> http_body::SizeHint {
        let mut hint = http_body::SizeHint::new();
        if self.head.is_none() {
            return self.tail.size_hint();
        }
        if let Some(data) = self.head.as_ref().and_then(|f| f.data_ref()) {
            hint.set_lower(data.len() as u64);
        }
        hint
    }
}