use crate::control::Hydrant;
use axum::Router;
use axum::routing::get;
use axum::{
    Json,
    extract::{Path, Query, State},
    http::{HeaderValue, StatusCode, header},
    response::{IntoResponse, Response},
};
use serde::{Deserialize, Serialize};
use tower_http::services::ServeFile;

use super::ws::{
    Message, WebSocket, WebSocketUpgrade, WsAction, control_frame_only_limits, run_socket,
};

pub fn router() -> Router<Hydrant> {
    Router::new()
        .route("/", get(handle_stream))
        .route("/head", get(handle_head))
        .route("/manifest", get(handle_manifest))
        .route("/chunk/{sha256}", get(handle_chunk))
}

/// the newest stored event id. identity and account events take ids too but
/// are never replayed, so this is the furthest a cursor replay can reach.
#[derive(Serialize)]
struct StreamHead<'a> {
    id: Option<u64>,
    stream: &'a str,
}

pub async fn handle_head(State(hydrant): State<Hydrant>) -> Response {
    match hydrant.stream_head() {
        Ok(id) => Json(StreamHead {
            id,
            stream: hydrant.stream_id(),
        })
        .into_response(),
        Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response(),
    }
}

pub async fn handle_manifest(State(hydrant): State<Hydrant>) -> Response {
    match hydrant.stream_chunks() {
        Some(chunks) => Json(chunks.manifest()).into_response(),
        None => (StatusCode::NOT_FOUND, "stream chunks are off").into_response(),
    }
}

pub async fn handle_chunk(
    State(hydrant): State<Hydrant>,
    Path(sha256): Path<String>,
    request: axum::extract::Request,
) -> Response {
    let Some(path) = hydrant.stream_chunks().and_then(|c| c.servable(&sha256)) else {
        return StatusCode::NOT_FOUND.into_response();
    };
    let served = ServeFile::new(path)
        .with_buf_chunk_size(256 * 1024)
        .try_call(request)
        .await;
    let mut response = match served {
        Ok(response) => response.map(axum::body::Body::new),
        Err(e) => return (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response(),
    };
    if !response.status().is_success() {
        return response;
    }
    let headers = response.headers_mut();
    headers.insert(
        header::CONTENT_TYPE,
        HeaderValue::from_static("application/zstd"),
    );
    headers.insert(
        header::CACHE_CONTROL,
        HeaderValue::from_static("public, max-age=31536000, immutable"),
    );
    if let Ok(etag) = HeaderValue::try_from(format!("\"{sha256}\"")) {
        headers.insert(header::ETAG, etag);
    }
    response
}

#[derive(Deserialize)]
pub struct StreamQuery {
    pub cursor: Option<u64>,
    /// leave it out and the cursor is not checked against this stream
    pub stream: Option<String>,
}

pub async fn handle_stream(
    State(hydrant): State<Hydrant>,
    Query(query): Query<StreamQuery>,
    ws: WebSocketUpgrade,
) -> impl IntoResponse {
    ws.limits(control_frame_only_limits())
        .on_upgrade(move |socket| handle_socket(socket, hydrant, query))
}

async fn handle_socket(socket: WebSocket, hydrant: Hydrant, query: StreamQuery) {
    let send_timeout = hydrant.stream_send_timeout();
    let draining = hydrant.state.drain.subscribe();
    if let Some(stream) = query.stream.filter(|s| s != hydrant.stream_id()) {
        let json = serde_json::json!({
            "type": "error",
            "error": "StreamMismatch",
            "message": format!("cursor is from stream {stream}, this is {}", hydrant.stream_id()),
        });
        let refusal = Message::text(json.to_string());
        run_socket(
            socket,
            futures::stream::iter([refusal]),
            |refusal| WsAction::Close(Some(refusal)),
            send_timeout,
            |_| None,
            draining,
        )
        .await;
        return;
    }
    let frames = hydrant.subscribe_frames(query.cursor);
    run_socket(
        socket,
        frames,
        |item| match item {
            Ok(frame) => WsAction::Send(Message::text(frame)),
            Err(err) => {
                let json = serde_json::json!({
                    "type": "error",
                    "error": err.code(),
                    "message": err.to_string(),
                });
                WsAction::Close(Some(Message::text(json.to_string())))
            }
        },
        send_timeout,
        |timeout_dur| {
            let json = serde_json::json!({
                "type": "error",
                "error": "ConsumerTooSlow",
                "message": format!(
                    "stream socket send blocked for at least {} seconds",
                    timeout_dur.as_secs()
                ),
            });
            Some(Message::text(json.to_string()))
        },
        draining,
    )
    .await;
}

#[cfg(test)]
mod tests {
    use super::*;
    use jacquard_common::types::did::Did;
    use miette::IntoDiagnostic;

    async fn head(hydrant: &Hydrant) -> miette::Result<serde_json::Value> {
        let response = handle_head(State(hydrant.clone())).await;
        assert_eq!(response.status(), StatusCode::OK);
        let body = axum::body::to_bytes(response.into_body(), 1024)
            .await
            .into_diagnostic()?;
        serde_json::from_slice(&body).into_diagnostic()
    }

    #[tokio::test]
    async fn head_tracks_durable_events_not_the_id_counter() -> miette::Result<()> {
        let tmp = tempfile::tempdir().into_diagnostic()?;
        let hydrant = Hydrant::new(crate::config::Config {
            database_path: tmp.path().to_path_buf(),
            ..Default::default()
        })
        .await?;
        let stream = hydrant.stream_id().to_owned();
        assert_eq!(
            head(&hydrant).await?,
            serde_json::json!({ "id": null, "stream": stream })
        );

        let ids = &hydrant.state.db.stream.ids;
        hydrant.seed_events_for_bench(3);
        let last_seeded = ids.head().expect("seeded events took ids");
        assert_eq!(
            head(&hydrant).await?,
            serde_json::json!({ "id": last_seeded, "stream": stream })
        );

        let did = Did::new_static("did:plc:aaaaaaaaaaaaaaaaaaaaaaaa").into_diagnostic()?;
        crate::ops::announce_identity(&hydrant.state.db, &did, None).await;
        assert_eq!(
            ids.head(),
            Some(last_seeded + 1),
            "the identity event must have taken an id"
        );
        assert_eq!(
            head(&hydrant).await?,
            serde_json::json!({ "id": last_seeded, "stream": stream })
        );
        Ok(())
    }

    #[tokio::test]
    async fn the_manifest_lists_the_chunks_and_is_404_while_they_are_off() -> miette::Result<()> {
        let manifest = |hydrant: &Hydrant| handle_manifest(State(hydrant.clone()));
        let off = tempfile::tempdir().into_diagnostic()?;
        let hydrant = Hydrant::new(crate::config::Config {
            database_path: off.path().to_path_buf(),
            ..Default::default()
        })
        .await?;
        assert_eq!(manifest(&hydrant).await.status(), StatusCode::NOT_FOUND);

        let on = tempfile::tempdir().into_diagnostic()?;
        let hydrant = Hydrant::new(crate::config::Config {
            database_path: on.path().to_path_buf(),
            stream_chunk_span: 4,
            ..Default::default()
        })
        .await?;
        hydrant.seed_events_for_bench(8);
        crate::control::stream::chunks::build(&hydrant.state)?;
        let listed = manifest(&hydrant).await;
        assert_eq!(listed.status(), StatusCode::OK);
        let body = axum::body::to_bytes(listed.into_body(), 1 << 20)
            .await
            .into_diagnostic()?;
        let chunks = hydrant.stream_chunks().expect("chunks are on").manifest();
        assert!(!chunks.chunks.is_empty());
        assert_eq!(
            serde_json::from_slice::<serde_json::Value>(&body).into_diagnostic()?,
            serde_json::to_value(chunks).into_diagnostic()?
        );
        Ok(())
    }

    #[tokio::test]
    async fn a_chunk_is_served_whole_or_by_range() -> miette::Result<()> {
        let tmp = tempfile::tempdir().into_diagnostic()?;
        let hydrant = Hydrant::new(crate::config::Config {
            database_path: tmp.path().to_path_buf(),
            stream_chunk_span: 4,
            ..Default::default()
        })
        .await?;
        hydrant.seed_events_for_bench(8);
        crate::control::stream::chunks::build(&hydrant.state)?;
        let chunks = hydrant.stream_chunks().expect("chunks are on");
        let sha256 = chunks.manifest().chunks[0].sha256.clone();
        let file = std::fs::read(chunks.path(&sha256)).into_diagnostic()?;
        let get = |sha256: &str, range: Option<&str>| {
            let mut request = axum::extract::Request::builder().uri("/");
            if let Some(range) = range {
                request = request.header(header::RANGE, range);
            }
            handle_chunk(
                State(hydrant.clone()),
                Path(sha256.to_owned()),
                request.body(axum::body::Body::empty()).unwrap(),
            )
        };
        let body = |response: Response| axum::body::to_bytes(response.into_body(), 1 << 20);

        let whole = get(&sha256, None).await;
        assert_eq!(whole.status(), StatusCode::OK);
        let headers = whole.headers();
        assert_eq!(headers[header::CONTENT_TYPE], "application/zstd");
        assert_eq!(
            headers[header::CACHE_CONTROL],
            "public, max-age=31536000, immutable"
        );
        assert_eq!(headers[header::ETAG], format!("\"{sha256}\"").as_str());
        assert_eq!(body(whole).await.into_diagnostic()?, file);

        let part = get(&sha256, Some("bytes=4-9")).await;
        assert_eq!(part.status(), StatusCode::PARTIAL_CONTENT);
        assert_eq!(body(part).await.into_diagnostic()?, file[4..10]);

        assert_eq!(get("00", None).await.status(), StatusCode::NOT_FOUND);
        Ok(())
    }
}