use axum::{body::Body, http, response::Response};
use bytes::Bytes;
use futures::stream::StreamExt;
use serde_json::Value;
use std::{
    sync::{Arc, Mutex},
    time::Duration,
};

mod error_summary;
#[cfg(test)]
mod tests;

use self::error_summary::{summarize_upstream_body_bytes, summarize_upstream_json_value};
use super::{
    error::ProxyError,
    metrics::estimate_tokens_from_bytes,
    providers::{
        streaming::create_anthropic_sse_stream,
        streaming_responses::create_anthropic_sse_stream_from_responses,
    },
};

pub struct PreparedResponse {
    pub response: Response,
    pub stream_completion: Option<StreamCompletion>,
    pub estimated_output_tokens: u64,
    pub upstream_error_summary: Option<String>,
    pub body_bytes: Option<Bytes>,
}

impl PreparedResponse {
    fn buffered(
        response: Response,
        estimated_output_tokens: u64,
        upstream_error_summary: Option<String>,
        body_bytes: Bytes,
    ) -> Self {
        Self {
            response,
            stream_completion: None,
            estimated_output_tokens,
            upstream_error_summary,
            body_bytes: Some(body_bytes),
        }
    }

    fn streaming(response: Response, stream_completion: StreamCompletion) -> Self {
        Self {
            response,
            stream_completion: Some(stream_completion),
            estimated_output_tokens: 0,
            upstream_error_summary: None,
            body_bytes: None,
        }
    }
}

#[derive(Clone, Default)]
pub struct StreamCompletion {
    inner: Arc<Mutex<Option<Result<(), String>>>>,
}

impl StreamCompletion {
    pub fn record_success(&self) {
        let mut outcome = self.inner.lock().expect("lock stream completion");
        if outcome.is_none() {
            *outcome = Some(Ok(()));
        }
    }

    pub fn record_error(&self, message: String) {
        let mut outcome = self.inner.lock().expect("lock stream completion");
        if outcome.is_none() {
            *outcome = Some(Err(message));
        }
    }

    pub fn outcome(&self) -> Option<Result<(), String>> {
        self.inner.lock().expect("lock stream completion").clone()
    }
}

pub fn is_sse_response(response: &reqwest::Response) -> bool {
    response
        .headers()
        .get("content-type")
        .and_then(|v| v.to_str().ok())
        .map(|ct| ct.contains("text/event-stream"))
        .unwrap_or(false)
}

pub async fn build_passthrough_response(
    response: reqwest::Response,
    first_byte_timeout: Option<Duration>,
    idle_timeout: Option<Duration>,
) -> Result<PreparedResponse, ProxyError> {
    let status = response.status();
    let headers = response.headers().clone();
    let mut builder = Response::builder().status(status);
    copy_headers(&mut builder, &headers, false);

    if is_sse_response(&response) {
        let stream_completion = StreamCompletion::default();
        let stream = with_stream_timeouts(
            response.bytes_stream(),
            first_byte_timeout,
            idle_timeout,
            Some(stream_completion.clone()),
        );
        return builder
            .body(Body::from_stream(stream))
            .map(|response| PreparedResponse::streaming(response, stream_completion))
            .map_err(|error| {
                ProxyError::RequestFailed(format!("build streaming response failed: {error}"))
            });
    }

    let body = read_buffered_body(response, first_byte_timeout).await?;
    let upstream_error_summary = if !status.is_success() {
        summarize_upstream_body_bytes(&body)
    } else {
        None
    };
    let estimated_output_tokens = estimate_tokens_from_bytes(&body);
    let response_bytes = body.clone();
    builder
        .body(Body::from(body))
        .map(|response| {
            PreparedResponse::buffered(
                response,
                estimated_output_tokens,
                upstream_error_summary,
                response_bytes,
            )
        })
        .map_err(|error| {
            ProxyError::RequestFailed(format!("build passthrough response failed: {error}"))
        })
}

pub async fn build_json_response<F>(
    response: reqwest::Response,
    first_byte_timeout: Option<Duration>,
    transform: F,
) -> Result<PreparedResponse, ProxyError>
where
    F: FnOnce(Value) -> Result<Value, ProxyError>,
{
    let status = response.status();
    let headers = response.headers().clone();
    let body = read_buffered_body(response, first_byte_timeout).await?;
    build_buffered_json_response_inner(status, &headers, body, transform)
}

pub fn build_buffered_passthrough_response(
    status: reqwest::StatusCode,
    headers: &reqwest::header::HeaderMap,
    body: Bytes,
) -> Result<PreparedResponse, ProxyError> {
    let upstream_error_summary = if !status.is_success() {
        summarize_upstream_body_bytes(&body)
    } else {
        None
    };
    let estimated_output_tokens = estimate_tokens_from_bytes(&body);
    let mut builder = Response::builder().status(status);
    copy_headers(&mut builder, headers, false);
    let response_bytes = body.clone();
    builder
        .body(Body::from(body))
        .map(|response| {
            PreparedResponse::buffered(
                response,
                estimated_output_tokens,
                upstream_error_summary,
                response_bytes,
            )
        })
        .map_err(|error| {
            ProxyError::RequestFailed(format!("build passthrough response failed: {error}"))
        })
}

pub fn build_buffered_json_response<F>(
    status: reqwest::StatusCode,
    headers: &reqwest::header::HeaderMap,
    body: Bytes,
    transform: F,
) -> Result<PreparedResponse, ProxyError>
where
    F: FnOnce(Value) -> Result<Value, ProxyError>,
{
    build_buffered_json_response_inner(status, headers, body, transform)
}

pub fn build_anthropic_stream_response(
    response: reqwest::Response,
    first_byte_timeout: Option<Duration>,
    idle_timeout: Option<Duration>,
    api_format: &str,
) -> Result<PreparedResponse, ProxyError> {
    let status = response.status();
    let headers = response.headers().clone();
    let mut builder = Response::builder().status(status);
    copy_headers(&mut builder, &headers, true);

    let stream_completion = StreamCompletion::default();
    let timed_stream = with_stream_timeouts(
        response.bytes_stream(),
        first_byte_timeout,
        idle_timeout,
        None,
    );
    let stream: std::pin::Pin<
        Box<dyn futures::Stream<Item = Result<Bytes, std::io::Error>> + Send>,
    > = if api_format == "openai_responses" {
        Box::pin(create_anthropic_sse_stream_from_responses(
            timed_stream,
            stream_completion.clone(),
        ))
    } else {
        Box::pin(create_anthropic_sse_stream(
            timed_stream,
            stream_completion.clone(),
        ))
    };
    builder
        .body(Body::from_stream(stream))
        .map(|response| PreparedResponse::streaming(response, stream_completion))
        .map_err(|error| {
            ProxyError::RequestFailed(format!("build anthropic stream response failed: {error}"))
        })
}

fn with_stream_timeouts(
    stream: impl futures::Stream<Item = Result<Bytes, reqwest::Error>> + Send + 'static,
    first_byte_timeout: Option<Duration>,
    idle_timeout: Option<Duration>,
    stream_completion: Option<StreamCompletion>,
) -> impl futures::Stream<Item = Result<Bytes, std::io::Error>> + Send {
    async_stream::stream! {
        tokio::pin!(stream);
        let mut is_first_chunk = true;

        while let Some(next) = next_chunk_with_timeout(
            &mut stream,
            if is_first_chunk { first_byte_timeout } else { idle_timeout },
            if is_first_chunk {
                StreamTimeoutPhase::FirstByte
            } else {
                StreamTimeoutPhase::Idle
            },
        ).await {
            match next {
                Ok(chunk) => {
                    is_first_chunk = false;
                    yield Ok(chunk);
                }
                Err(error) => {
                    if let Some(stream_completion) = &stream_completion {
                        stream_completion.record_error(error.to_string());
                    }
                    yield Err(error);
                    return;
                }
            }
        }

        if let Some(stream_completion) = &stream_completion {
            stream_completion.record_success();
        }
    }
}

async fn read_buffered_body(
    response: reqwest::Response,
    timeout_duration: Option<Duration>,
) -> Result<Bytes, ProxyError> {
    match timeout_duration {
        Some(timeout) => match tokio::time::timeout(timeout, response.bytes()).await {
            Ok(result) => result.map_err(|error| {
                ProxyError::RequestFailed(format!("read response body failed: {error}"))
            }),
            Err(_) => Err(ProxyError::Timeout(
                StreamTimeoutPhase::FirstByte.error_message(timeout),
            )),
        },
        None => response.bytes().await.map_err(|error| {
            ProxyError::RequestFailed(format!("read response body failed: {error}"))
        }),
    }
}

async fn next_chunk_with_timeout<S>(
    stream: &mut S,
    timeout_duration: Option<Duration>,
    phase: StreamTimeoutPhase,
) -> Option<Result<Bytes, std::io::Error>>
where
    S: futures::Stream<Item = Result<Bytes, reqwest::Error>> + Unpin,
{
    let next = match timeout_duration {
        Some(timeout) => match tokio::time::timeout(timeout, stream.next()).await {
            Ok(next) => next,
            Err(_) => {
                return Some(Err(std::io::Error::new(
                    std::io::ErrorKind::TimedOut,
                    phase.error_message(timeout),
                )));
            }
        },
        None => stream.next().await,
    };

    next.map(|result| result.map_err(std::io::Error::other))
}

#[derive(Clone, Copy)]
enum StreamTimeoutPhase {
    FirstByte,
    Idle,
}

impl StreamTimeoutPhase {
    fn error_message(self, timeout: Duration) -> String {
        let display_seconds = timeout.as_secs().max(u64::from(!timeout.is_zero()));
        match self {
            StreamTimeoutPhase::FirstByte => {
                format!("stream timeout after {}s", display_seconds)
            }
            StreamTimeoutPhase::Idle => {
                format!("stream idle timeout after {}s", display_seconds)
            }
        }
    }
}

fn copy_headers(
    builder: &mut http::response::Builder,
    headers: &reqwest::header::HeaderMap,
    force_sse_content_type: bool,
) {
    for (key, value) in headers {
        let lower = key.as_str().to_ascii_lowercase();
        if lower == "content-length" || lower == "transfer-encoding" {
            continue;
        }
        if force_sse_content_type && lower == "content-type" {
            continue;
        }
        *builder = std::mem::take(builder).header(key, value);
    }

    if force_sse_content_type {
        *builder = std::mem::take(builder).header("content-type", "text/event-stream");
    }
}

fn build_buffered_json_response_inner<F>(
    status: reqwest::StatusCode,
    headers: &reqwest::header::HeaderMap,
    body: Bytes,
    transform: F,
) -> Result<PreparedResponse, ProxyError>
where
    F: FnOnce(Value) -> Result<Value, ProxyError>,
{
    let upstream_body: Value = match serde_json::from_slice(&body) {
        Ok(body) => body,
        Err(_) if !status.is_success() => {
            return build_buffered_passthrough_response(status, headers, body);
        }
        Err(error) => {
            return Err(ProxyError::RequestFailed(format!(
                "parse upstream json failed: {error}"
            )));
        }
    };
    let upstream_error_summary = if !status.is_success() {
        summarize_upstream_json_value(&upstream_body)
    } else {
        None
    };
    let response_body = match transform(upstream_body) {
        Ok(body) => body,
        Err(error) if should_passthrough_transform_failure(status, &error) => {
            return build_buffered_passthrough_response(status, headers, body);
        }
        Err(error) => {
            if !status.is_success() {
                return Err(error);
            }
            return Err(ProxyError::RequestFailed(format!(
                "transform upstream json failed: {}",
                proxy_error_message(error)
            )));
        }
    };
    let response_body = match serde_json::to_vec(&response_body) {
        Ok(body) => body,
        Err(error) => {
            return Err(ProxyError::RequestFailed(format!(
                "serialize transformed json failed: {error}"
            )));
        }
    };
    let response_bytes = Bytes::from(response_body);
    let estimated_output_tokens = estimate_tokens_from_bytes(&response_bytes);

    let mut builder = Response::builder().status(status);
    copy_headers(&mut builder, headers, false);
    builder = builder.header("content-type", "application/json");

    builder
        .body(Body::from(response_bytes.clone()))
        .map(|response| {
            PreparedResponse::buffered(
                response,
                estimated_output_tokens,
                upstream_error_summary,
                response_bytes,
            )
        })
        .map_err(|error| {
            ProxyError::RequestFailed(format!("build transformed response failed: {error}"))
        })
}

fn should_passthrough_transform_failure(status: reqwest::StatusCode, error: &ProxyError) -> bool {
    !status.is_success() && matches!(error, ProxyError::TransformError(_))
}

fn proxy_error_message(error: ProxyError) -> String {
    match error {
        ProxyError::ConfigError(message)
        | ProxyError::AuthError(message)
        | ProxyError::RequestFailed(message)
        | ProxyError::TransformError(message)
        | ProxyError::ForwardFailed(message)
        | ProxyError::BindFailed(message)
        | ProxyError::StopFailed(message)
        | ProxyError::ProviderUnhealthy(message)
        | ProxyError::DatabaseError(message)
        | ProxyError::InvalidRequest(message)
        | ProxyError::Timeout(message)
        | ProxyError::Internal(message) => message,
        other => other.to_string(),
    }
}
