use std::fmt::Display;
use std::future::Future;
use std::time::Duration;

use axum::body::Body;
use axum::extract::FromRequestParts;
use axum::http::request::Parts;
use axum::http::{HeaderMap, HeaderName, HeaderValue, Method, StatusCode, header};
use axum::response::Response;
use bytes::Bytes;
use futures::{FutureExt, Sink, SinkExt, Stream, StreamExt};
use hyper_util::rt::TokioIo;
use tokio::time::{MissedTickBehavior, interval, timeout};
use tokio_websockets::CloseCode;
pub(crate) use tokio_websockets::{Limits, Message};
use tracing::{debug, warn};

use crate::util::drain::Draining;

pub(crate) type WebSocket = tokio_websockets::WebSocketStream<TokioIo<hyper::upgrade::Upgraded>>;

pub(crate) struct WebSocketUpgrade {
    limits: Limits,
    sec_websocket_key: HeaderValue,
    on_upgrade: hyper::upgrade::OnUpgrade,
}

impl<S: Send + Sync> FromRequestParts<S> for WebSocketUpgrade {
    type Rejection = StatusCode;

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        if parts.method != Method::GET {
            return Err(StatusCode::METHOD_NOT_ALLOWED);
        }
        if !header_contains(&parts.headers, header::CONNECTION, "upgrade")
            || !header_eq(&parts.headers, header::UPGRADE, "websocket")
            || !header_eq(&parts.headers, header::SEC_WEBSOCKET_VERSION, "13")
        {
            return Err(StatusCode::BAD_REQUEST);
        }
        let sec_websocket_key = parts
            .headers
            .get(header::SEC_WEBSOCKET_KEY)
            .ok_or(StatusCode::BAD_REQUEST)?
            .clone();
        let on_upgrade = parts
            .extensions
            .remove::<hyper::upgrade::OnUpgrade>()
            .ok_or(StatusCode::UPGRADE_REQUIRED)?;
        Ok(Self {
            limits: Limits::default(),
            sec_websocket_key,
            on_upgrade,
        })
    }
}

impl WebSocketUpgrade {
    pub(crate) fn limits(mut self, limits: Limits) -> Self {
        self.limits = limits;
        self
    }

    #[must_use]
    pub(crate) fn on_upgrade<C, Fut>(self, callback: C) -> Response
    where
        C: FnOnce(WebSocket) -> Fut + Send + 'static,
        Fut: Future<Output = ()> + Send + 'static,
    {
        let on_upgrade = self.on_upgrade;
        let limits = self.limits;
        tokio::spawn(async move {
            let Ok(upgraded) = on_upgrade.await else {
                return;
            };
            let socket = tokio_websockets::server::Builder::new()
                .limits(limits)
                .serve(TokioIo::new(upgraded));
            callback(socket).await;
        });

        Response::builder()
            .status(StatusCode::SWITCHING_PROTOCOLS)
            .header(header::CONNECTION, HeaderValue::from_static("upgrade"))
            .header(header::UPGRADE, HeaderValue::from_static("websocket"))
            .header(
                header::SEC_WEBSOCKET_ACCEPT,
                sec_websocket_accept(self.sec_websocket_key.as_bytes()),
            )
            .body(Body::empty())
            .unwrap()
    }
}

fn header_contains(headers: &HeaderMap, key: HeaderName, value: &'static str) -> bool {
    headers
        .get(&key)
        .and_then(|h| std::str::from_utf8(h.as_bytes()).ok())
        .is_some_and(|h| h.to_ascii_lowercase().contains(value))
}

fn header_eq(headers: &HeaderMap, key: HeaderName, value: &'static str) -> bool {
    headers
        .get(&key)
        .is_some_and(|h| h.as_bytes().eq_ignore_ascii_case(value.as_bytes()))
}

fn sec_websocket_accept(key: &[u8]) -> HeaderValue {
    let mut ctx = aws_lc_rs::digest::Context::new(&aws_lc_rs::digest::SHA1_FOR_LEGACY_USE_ONLY);
    ctx.update(key);
    ctx.update(b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11");
    HeaderValue::from_str(&data_encoding::BASE64.encode(ctx.finish().as_ref()))
        .expect("valid base64")
}

const PING_INTERVAL: Duration = Duration::from_secs(30);
const CLOSE_TIMEOUT: Duration = Duration::from_secs(1);
const CONTROL_FRAME_PAYLOAD_LIMIT: usize = 125;
/// a batch flushes at whichever of these it reaches first, so a reader that keeps up still
/// gets frames in steady writes rather than one big one
const BATCH_FRAMES: usize = 64;
const BATCH_BYTES: usize = 64 * 1024;

pub(super) fn control_frame_only_limits() -> Limits {
    Limits::default().max_payload_len(Some(CONTROL_FRAME_PAYLOAD_LIMIT))
}

/// 1012 is what makes a client reconnect now instead of backing off
pub(super) fn restarting() -> Message {
    Message::close(Some(CloseCode::SERVICE_RESTART), "hydrant is restarting")
}

pub(super) enum WsAction {
    Send(Message),
    Close(Option<Message>),
}

pub(super) async fn run_socket<S, F, G>(
    socket: WebSocket,
    events: S,
    to_action: F,
    send_timeout: Duration,
    slow_consumer: G,
    draining: Draining,
) where
    S: Stream + Unpin,
    F: FnMut(S::Item) -> WsAction,
    G: FnOnce(Duration) -> Option<Message>,
{
    let (sink, ws_recv) = socket.split();
    serve(
        sink,
        ws_recv,
        events,
        to_action,
        send_timeout,
        slow_consumer,
        draining,
    )
    .await;
}

async fn serve<Si, R, E, S, F, G>(
    mut sink: Si,
    mut ws_recv: R,
    mut events: S,
    mut to_action: F,
    send_timeout: Duration,
    slow_consumer: G,
    mut draining: Draining,
) where
    Si: Sink<Message, Error = E> + Unpin,
    R: Stream<Item = Result<Message, E>> + Unpin,
    E: Display,
    S: Stream + Unpin,
    F: FnMut(S::Item) -> WsAction,
    G: FnOnce(Duration) -> Option<Message>,
{
    let mut ping_timer = interval(PING_INTERVAL);
    ping_timer.set_missed_tick_behavior(MissedTickBehavior::Delay);
    ping_timer.tick().await;

    let mut slow_consumer = Some(slow_consumer);

    loop {
        tokio::select! {
            _ = draining.started() => {
                let _ = timeout(CLOSE_TIMEOUT, sink.send(restarting())).await;
                break;
            }
            inbound = ws_recv.next() => match inbound {
                Some(Ok(m)) if m.is_close() => break,
                Some(Ok(m)) if m.is_text() || m.is_binary() => {
                    debug!("client sent unsolicited data frame, closing");
                    break;
                }
                Some(Ok(m)) if m.is_ping() => {
                    if let Err(err) = sink.send(Message::pong(m.into_payload())).await {
                        warn!(err = %err, "ws pong send error");
                        break;
                    }
                }
                Some(Ok(_)) => {}
                Some(Err(e)) => {
                    warn!(err = %e, "ws recv error");
                    break;
                }
                None => break,
            },
            evt = events.next() => {
                let Some(item) = evt else { break };
                let first = match to_action(item) {
                    WsAction::Send(msg) => msg,
                    WsAction::Close(final_msg) => {
                        if let Some(m) = final_msg {
                            let _ = timeout(send_timeout, sink.send(m)).await;
                        }
                        break;
                    }
                };
                let batch = send_batch(&mut sink, first, &mut events, &mut to_action);
                match timeout(send_timeout, batch).await {
                    Ok(Ok(Batch::Open)) => {}
                    Ok(Ok(Batch::Ended(final_msg))) => {
                        if let Some(m) = final_msg {
                            let _ = timeout(send_timeout, sink.send(m)).await;
                        }
                        break;
                    }
                    Ok(Err(err)) => {
                        warn!(err = %err, "ws send error");
                        break;
                    }
                    Err(_) => {
                        if let Some(final_msg) = slow_consumer.take().and_then(|f| f(send_timeout))
                        {
                            let _ = timeout(CLOSE_TIMEOUT, sink.send(final_msg)).await;
                        }
                        break;
                    }
                }
            },
            _ = ping_timer.tick() => {
                if let Err(err) = sink.send(Message::ping(Bytes::new())).await {
                    warn!(err = %err, "ws ping send error");
                    break;
                }
            }
        }
    }
    let _ = timeout(CLOSE_TIMEOUT, sink.close()).await;
}

/// whether the events go on after a batch
enum Batch {
    Open,
    /// the events ended, with the message to close on if there is one
    Ended(Option<Message>),
}

/// sends `first` and whatever else is ready right now, up to a batch, then flushes once, so a
/// reader that keeps up doesn't cost a write per frame
async fn send_batch<Si, S, F>(
    sink: &mut Si,
    first: Message,
    events: &mut S,
    to_action: &mut F,
) -> Result<Batch, Si::Error>
where
    Si: Sink<Message> + Unpin,
    S: Stream + Unpin,
    F: FnMut(S::Item) -> WsAction,
{
    let mut frames = 1;
    let mut bytes = first.as_payload().len();
    sink.feed(first).await?;
    let after = loop {
        if frames >= BATCH_FRAMES || bytes >= BATCH_BYTES {
            break Batch::Open;
        }
        let Some(next) = events.next().now_or_never() else {
            break Batch::Open;
        };
        let Some(item) = next else {
            break Batch::Ended(None);
        };
        frames += 1;
        match to_action(item) {
            WsAction::Send(msg) => {
                bytes += msg.as_payload().len();
                sink.feed(msg).await?;
            }
            WsAction::Close(final_msg) => break Batch::Ended(final_msg),
        }
    };
    sink.flush().await?;
    Ok(after)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::pin::Pin;
    use std::sync::{Arc, Mutex};
    use std::task::{Context, Poll};

    use crate::util::drain::Drain;

    /// a client socket that keeps what it was sent and how many flushes wrote it, or that stops
    /// reading, so no flush ever finishes
    #[derive(Clone, Default)]
    struct Client {
        sent: Arc<Mutex<Vec<String>>>,
        flushes: Arc<Mutex<usize>>,
        stalled: bool,
    }

    impl Client {
        fn texts(&self) -> Vec<String> {
            self.sent.lock().unwrap().clone()
        }
    }

    impl Sink<Message> for Client {
        type Error = std::convert::Infallible;

        fn poll_ready(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            Poll::Ready(Ok(()))
        }

        fn start_send(self: Pin<&mut Self>, msg: Message) -> Result<(), Self::Error> {
            if let Some(text) = msg.as_text() {
                self.sent.lock().unwrap().push(text.to_owned());
            } else if let Some((code, reason)) = msg.as_close() {
                let code = u16::from(code);
                self.sent.lock().unwrap().push(format!("close {code} {reason}"));
            }
            Ok(())
        }

        fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            if self.stalled {
                return Poll::Pending;
            }
            *self.flushes.lock().unwrap() += 1;
            Poll::Ready(Ok(()))
        }

        fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            self.poll_flush(cx)
        }
    }

    async fn run(client: Client, events: Vec<WsAction>, send_timeout: Duration) {
        let drain = Drain::default();
        serve_until(client, futures::stream::iter(events), send_timeout, &drain).await;
    }

    async fn serve_until<S: Stream<Item = WsAction> + Unpin>(
        client: Client,
        events: S,
        send_timeout: Duration,
        drain: &Drain,
    ) {
        let quiet = futures::stream::pending::<Result<Message, std::convert::Infallible>>();
        serve(
            client,
            quiet,
            events,
            |action| action,
            send_timeout,
            |_| Some(Message::text("too slow")),
            drain.subscribe(),
        )
        .await;
    }

    fn text(n: usize) -> WsAction {
        WsAction::Send(Message::text(n.to_string()))
    }

    #[tokio::test]
    async fn frames_ready_together_go_out_in_order_with_a_flush_per_batch() {
        let client = Client::default();
        let mut events: Vec<_> = (0..200).map(text).collect();
        events.push(WsAction::Close(Some(Message::text("bye"))));
        run(client.clone(), events, Duration::from_secs(5)).await;

        let expected: Vec<_> = (0..200)
            .map(|n| n.to_string())
            .chain(["bye".into()])
            .collect();
        assert_eq!(client.texts(), expected);
        // four batches of up to 64 items, the closing message, then the close
        assert_eq!(*client.flushes.lock().unwrap(), 6);
    }

    #[tokio::test]
    async fn a_batch_flushes_at_its_byte_limit() {
        let client = Client::default();
        let big = "x".repeat(BATCH_BYTES / 2);
        let events = (0..4)
            .map(|_| WsAction::Send(Message::text(big.clone())))
            .collect();
        run(client.clone(), events, Duration::from_secs(5)).await;

        assert_eq!(client.texts().len(), 4);
        // two batches of two, then the close
        assert_eq!(*client.flushes.lock().unwrap(), 3);
    }

    #[tokio::test]
    async fn a_client_that_stops_reading_is_closed_after_the_send_timeout() {
        let client = Client {
            stalled: true,
            ..Client::default()
        };
        let send_timeout = Duration::from_millis(300);
        let started = std::time::Instant::now();
        run(client.clone(), (0..10).map(text).collect(), send_timeout).await;

        // then up to a close timeout each for the closing message and the close
        let took = started.elapsed();
        assert!(took >= send_timeout, "{took:?}");
        assert!(took < send_timeout + CLOSE_TIMEOUT * 3, "{took:?}");
        let sent = client.texts();
        assert_eq!(sent.last().map(String::as_str), Some("too slow"));
        assert_eq!(
            sent[..10],
            (0..10).map(|n| n.to_string()).collect::<Vec<_>>()
        );
    }

    #[tokio::test]
    async fn a_drain_closes_an_idle_stream_as_a_restart() {
        let client = Client::default();
        let drain = Drain::default();
        let events = futures::stream::iter((0..3).map(text)).chain(futures::stream::pending());
        let sent_all = async {
            while client.texts().len() < 3 {
                tokio::time::sleep(Duration::from_millis(5)).await;
            }
            drain.drain(Duration::from_secs(5)).await;
        };
        let served = serve_until(client.clone(), events, Duration::from_secs(5), &drain);
        tokio::time::timeout(Duration::from_secs(5), futures::future::join(served, sent_all))
            .await
            .expect("the drain must end the stream");

        assert_eq!(
            client.texts(),
            ["0", "1", "2", "close 1012 hydrant is restarting"]
        );
    }
}