//! WebSocket upgrade
//!
//! Provides WebSocket handshake and upgrade functionality following RFC 6455.
//!
//! # Examples
//!
//! ```rust,ignore
//! use ruxno::upgrade::websocket::WebSocketUpgrade;
//!
//! app.get("/ws", |ctx: Context| async move {
//!     let ws = WebSocketUpgrade::new(ctx);
//!     ws.upgrade(|socket| async move {
//!         // Handle WebSocket connection
//!         while let Some(Ok(msg)) = socket.recv().await {
//!             socket.send(msg).await.ok();
//!         }
//!     }).await
//! });
//! ```

use crate::core::CoreError;
use crate::domain::{Context, Response};
use crate::upgrade::websocket::socket::{ConnectionId, WebSocket};
use hyper_tungstenite::{is_upgrade_request, upgrade};
use sha1::{Digest, Sha1};
use tokio::sync::broadcast;

/// WebSocket GUID constant from RFC 6455
///
/// This magic string is concatenated with the client's Sec-WebSocket-Key
/// and hashed to produce the Sec-WebSocket-Accept value.
const WEBSOCKET_GUID: &str = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

/// WebSocket upgrade
///
/// Handles the WebSocket handshake and upgrade process according to RFC 6455.
/// Validates the upgrade request, computes the accept key, and returns a
/// 101 Switching Protocols response.
///
/// # Design
///
/// The upgrade process follows RFC 6455:
/// 1. Validate the upgrade request (Upgrade + Sec-WebSocket-Key headers)
/// 2. Compute Sec-WebSocket-Accept from the client's key
/// 3. Return 101 Switching Protocols with proper headers
/// 4. Spawn a task to handle the WebSocket connection
///
/// # Examples
///
/// ```rust,ignore
/// use ruxno::upgrade::websocket::WebSocketUpgrade;
///
/// app.get("/ws", |ctx: Context| async move {
///     let ws = WebSocketUpgrade::new(ctx);
///     ws.upgrade(|mut socket| async move {
///         while let Some(Ok(msg)) = socket.recv().await {
///             socket.send(msg).await.ok();
///         }
///     }).await
/// });
/// ```
pub struct WebSocketUpgrade<E = ()> {
    ctx: Context<E>,
    broadcast_capacity: Option<usize>,
}

impl<E> From<Context<E>> for WebSocketUpgrade<E> {
    fn from(value: Context<E>) -> Self {
        Self::new(value)
    }
}

impl<E> WebSocketUpgrade<E> {
    /// Create new WebSocket upgrade from context
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let ws = WebSocketUpgrade::new(ctx);
    /// ```
    pub fn new(ctx: Context<E>) -> Self {
        Self {
            ctx,
            broadcast_capacity: None,
        }
    }

    /// Enable broadcasting with the specified channel capacity
    pub fn with_broadcast(mut self, capacity: usize) -> Self {
        self.broadcast_capacity = Some(capacity);
        self
    }

    /// Get the WebSocket accept key
    ///
    /// Computes the `Sec-WebSocket-Accept` value from the client's
    /// `Sec-WebSocket-Key` header according to RFC 6455.
    ///
    /// # Algorithm (RFC 6455 Section 1.3)
    ///
    /// 1. Validate the client's key (must be 16 bytes base64-encoded)
    /// 2. Concatenate the client's key with the WebSocket GUID
    /// 3. Compute SHA-1 hash of the concatenated string
    /// 4. Base64 encode the hash
    ///
    /// # Validation (RFC 6455 Section 4.2.1)
    ///
    /// The `Sec-WebSocket-Key` header must:
    /// - Be present in the request
    /// - Be a valid base64-encoded string
    /// - Decode to exactly 16 bytes
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let ws = WebSocketUpgrade::new(ctx);
    /// let accept_key = ws.accept_key()?;
    /// ```
    ///
    /// # Returns
    ///
    /// Returns `Ok(String)` with the base64-encoded accept key, or
    /// `Err(CoreError)` if the key is missing or invalid.
    pub fn accept_key(&self) -> Result<String, CoreError> {
        // Get the Sec-WebSocket-Key header from the request
        let key = self
            .ctx
            .req
            .headers()
            .get("sec-websocket-key")
            .ok_or_else(|| {
                CoreError::bad_request("Missing Sec-WebSocket-Key header for WebSocket upgrade")
            })?;

        // Validate key length (must be 24 characters for 16 bytes base64-encoded)
        // RFC 6455 Section 4.2.1: The value must be a base64-encoded value that,
        // when decoded, is 16 bytes in length
        if key.len() != 24 {
            return Err(CoreError::bad_request(format!(
                "Invalid Sec-WebSocket-Key length: expected 24 characters, got {}",
                key.len()
            )));
        }

        // Validate base64 encoding and decode to verify it's exactly 16 bytes
        let decoded = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, key)
            .map_err(|e| {
                CoreError::bad_request(format!(
                    "Invalid Sec-WebSocket-Key: not valid base64 encoding: {}",
                    e
                ))
            })?;

        // Verify decoded length is exactly 16 bytes (RFC 6455 requirement)
        if decoded.len() != 16 {
            return Err(CoreError::bad_request(format!(
                "Invalid Sec-WebSocket-Key: decoded to {} bytes, expected 16 bytes",
                decoded.len()
            )));
        }

        // Concatenate with the WebSocket GUID and hash
        let mut hasher = Sha1::new();
        hasher.update(key.as_bytes());
        hasher.update(WEBSOCKET_GUID.as_bytes());
        let hash = hasher.finalize();

        // Base64 encode the result
        Ok(base64::Engine::encode(
            &base64::engine::general_purpose::STANDARD,
            hash,
        ))
    }

    /// Perform the WebSocket upgrade
    ///
    /// Validates the upgrade request, performs the handshake, and spawns
    /// a task to handle the WebSocket connection.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// use ruxno::upgrade::websocket::{WebSocketUpgrade, Message};
    ///
    /// app.get("/ws", |ctx: Context| async move {
    ///     let ws = WebSocketUpgrade::new(ctx);
    ///     ws.upgrade(|mut socket| async move {
    ///         while let Some(Ok(msg)) = socket.recv().await {
    ///             match msg {
    ///                 Message::Text(text) => {
    ///                     println!("Received: {}", text);
    ///                     socket.send(Message::Text(text)).await.ok();
    ///                 }
    ///                 Message::Binary(data) => {
    ///                     socket.send(Message::Binary(data)).await.ok();
    ///                 }
    ///                 _ => {}
    ///             }
    ///         }
    ///     }).await
    /// });
    /// ```
    ///
    /// # Returns
    ///
    /// Returns a `Response` with status 101 (Switching Protocols) and the
    /// required WebSocket headers, or an error if the upgrade fails.
    pub async fn upgrade<F, Fut>(self, handler: F) -> Result<Response, CoreError>
    where
        F: FnOnce(WebSocket) -> Fut + Send + 'static,
        Fut: std::future::Future<Output = ()> + Send + 'static,
    {
        // Get the raw hyper request for upgrade
        let mut raw_guard = self.ctx.req.raw_request().lock().await;

        if raw_guard.is_none() {
            return Err(CoreError::internal(
                "WebSocket upgrade requires original hyper request",
            ));
        }

        let mut hyper_req = raw_guard.take().unwrap();

        // Validate upgrade request
        if !is_upgrade_request(&hyper_req) {
            return Err(CoreError::bad_request(
                "Not a valid WebSocket upgrade request",
            ));
        }

        // Perform the upgrade using hyper-tungstenite.
        // CRITICAL: The returned `response` contains hyper upgrade extensions that
        // must be sent to the client as-is. Any reconstruction (collecting body,
        // rebuilding response) will lose those extensions and break the handshake.
        let (response, websocket) = upgrade(&mut hyper_req, None)
            .map_err(|e| CoreError::internal(format!("WebSocket upgrade failed: {}", e)))?;

        // Create broadcast channel if requested
        let broadcast_tx = self.broadcast_capacity.map(|cap| broadcast::channel(cap).0);

        // Spawn task to handle WebSocket connection after the upgrade completes.
        // The websocket future resolves once the 101 response has been sent.
        let conn_id = ConnectionId::new();
        let broadcast_tx_clone = broadcast_tx.clone();

        tokio::spawn(async move {
            match websocket.await {
                Ok(ws_stream) => {
                    let socket = WebSocket::new(ws_stream, conn_id, broadcast_tx_clone);
                    handler(socket).await;
                }
                Err(e) => {
                    eprintln!("WebSocket upgrade error: {}", e);
                }
            }
        });

        // Wrap the raw hyper response in the Upgrade variant so to_hyper_response
        // returns it unchanged, preserving all extensions needed for the upgrade.
        Ok(Response::upgrade(response))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::http::Headers;
    use crate::{domain::Request, Query};
    use bytes::Bytes;
    use http::{Method, Uri, Version};
    use std::sync::Arc;

    fn create_websocket_context() -> Context<()> {
        let mut headers = Headers::new();
        headers.set("upgrade", "websocket").unwrap();
        headers
            .set("sec-websocket-key", "dGhlIHNhbXBsZSBub25jZQ==")
            .unwrap();

        let request = Request::new(
            Method::GET,
            Uri::from_static("http://localhost/ws"),
            Version::HTTP_11,
            Query::new(),
            headers,
            Bytes::new(),
        );

        Context::new(request, ())
    }

    #[test]
    fn test_accept_key_rfc6455() {
        // Test vector from RFC 6455 Section 1.3
        // Client sends: Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==
        // Server should respond with: Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=

        let ctx = create_websocket_context();
        let ws = WebSocketUpgrade::new(ctx);

        let accept_key = ws.accept_key().unwrap();
        assert_eq!(accept_key, "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=");
    }

    #[test]
    fn test_accept_key_missing_header() {
        let mut headers = Headers::new();
        headers.set("upgrade", "websocket").unwrap();
        // Missing sec-websocket-key

        let request = Request::new(
            Method::GET,
            Uri::from_static("http://localhost/ws"),
            Version::HTTP_11,
            Query::new(),
            headers,
            Bytes::new(),
        );

        let ctx = Context::new(request, Arc::new(()));
        let ws = WebSocketUpgrade::new(ctx);

        let result = ws.accept_key();
        assert!(result.is_err());
    }

    /// Build a domain `Request` that has a real hyper request (with upgrade
    /// extensions) stored in its `raw_request` slot.
    ///
    /// Uses `tokio::io::duplex` to create an in-memory I/O pair and runs a
    /// hyper HTTP/1.1 connection on the server half with `.with_upgrades()`.
    /// The server-side handler captures the hyper request and stores it in the
    /// domain request before returning a placeholder response. The client half
    /// sends a raw WebSocket upgrade HTTP/1.1 request so hyper wires up the
    /// upgrade extensions on the server side.
    async fn create_websocket_context_with_hyper_request() -> Context<()> {
        use hyper::body::Incoming;
        use hyper::server::conn::http1;
        use hyper::service::service_fn;
        use hyper_util::rt::TokioIo;
        use tokio::sync::oneshot;

        let (tx, rx) = oneshot::channel::<hyper::Request<Incoming>>();
        let tx = Arc::new(tokio::sync::Mutex::new(Some(tx)));

        // In-memory duplex stream — no TCP socket needed.
        let (client_io, server_io) = tokio::io::duplex(4096);

        // Server half: serve one request, capture the hyper Request.
        tokio::spawn(async move {
            let tx = Arc::clone(&tx);
            let svc = service_fn(move |req: hyper::Request<Incoming>| {
                let tx = Arc::clone(&tx);
                async move {
                    if let Some(sender) = tx.lock().await.take() {
                        let _ = sender.send(req);
                    }
                    // Minimal 101 placeholder — hyper-tungstenite overwrites this.
                    Ok::<_, std::convert::Infallible>(
                        hyper::Response::builder()
                            .status(101)
                            .body(http_body_util::Empty::<bytes::Bytes>::new())
                            .unwrap(),
                    )
                }
            });
            let _ = http1::Builder::new()
                .serve_connection(TokioIo::new(server_io), svc)
                .with_upgrades()
                .await;
        });

        // Client half: send a raw WS upgrade request so hyper wires up extensions.
        tokio::spawn(async move {
            use tokio::io::{AsyncReadExt, AsyncWriteExt};
            let mut client = client_io;
            let req = concat!(
                "GET /ws HTTP/1.1\r\n",
                "Host: localhost\r\n",
                "Upgrade: websocket\r\n",
                "Connection: Upgrade\r\n",
                "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n",
                "Sec-WebSocket-Version: 13\r\n",
                "\r\n"
            );
            let _ = client.write_all(req.as_bytes()).await;
            // Drain the response so hyper can finish writing it.
            let mut buf = vec![0u8; 4096];
            let _ = client.read(&mut buf).await;
        });

        // Wait for the server handler to capture the hyper request.
        let hyper_req = rx.await.expect("server handler should send hyper request");

        // Build the domain Request with matching WS headers.
        let mut headers = Headers::new();
        headers.set("upgrade", "websocket").unwrap();
        headers
            .set("sec-websocket-key", "dGhlIHNhbXBsZSBub25jZQ==")
            .unwrap();

        let domain_req = Request::new(
            Method::GET,
            Uri::from_static("http://localhost/ws"),
            Version::HTTP_11,
            crate::domain::Query::new(),
            headers,
            Bytes::new(),
        );

        // Inject the real hyper request so upgrade() can call hyper-tungstenite.
        *domain_req.raw_request().lock().await = Some(hyper_req);

        Context::new(domain_req, ())
    }

    #[tokio::test]
    async fn test_upgrade_returns_101() {
        let ctx = create_websocket_context_with_hyper_request().await;
        let ws = WebSocketUpgrade::new(ctx);

        let response = ws.upgrade(|_socket| async move {}).await.unwrap();

        assert_eq!(response.status(), http::StatusCode::SWITCHING_PROTOCOLS);
    }

    #[tokio::test]
    async fn test_upgrade_sets_headers() {
        let ctx = create_websocket_context_with_hyper_request().await;
        let ws = WebSocketUpgrade::new(ctx);

        let response = ws.upgrade(|_socket| async move {}).await.unwrap();

        // The Upgrade response variant wraps the raw hyper response — the 101
        // status is the observable signal that the handshake succeeded.
        // Full header assertions (upgrade, connection, sec-websocket-accept)
        // belong in an end-to-end test where the response is sent over the wire.
        assert_eq!(response.status(), http::StatusCode::SWITCHING_PROTOCOLS);
    }

    #[tokio::test]
    async fn test_upgrade_missing_key() {
        let mut headers = Headers::new();
        headers.set("upgrade", "websocket").unwrap();
        // Missing sec-websocket-key

        let request = Request::new(
            Method::GET,
            Uri::from_static("http://localhost/ws"),
            Version::HTTP_11,
            Query::new(),
            headers,
            Bytes::new(),
        );

        let ctx = Context::new(request, Arc::new(()));
        let ws = WebSocketUpgrade::new(ctx);

        let result = ws
            .upgrade(|_socket| async move {
                // Handler
            })
            .await;

        assert!(result.is_err());
    }

    #[test]
    fn test_accept_key_different_keys() {
        // Test with different client keys to ensure uniqueness
        let keys = vec![
            "x3JJHMbDL1EzLkh9GBhXDw==", // Valid 16-byte key
            "dGhlIHNhbXBsZSBub25jZQ==", // RFC 6455 example
            "AAAAAAAAAAAAAAAAAAAAAA==", // All zeros (16 bytes)
        ];

        let mut accept_keys = Vec::new();

        for key in keys {
            let mut headers = Headers::new();
            headers.set("upgrade", "websocket").unwrap();
            headers.set("sec-websocket-key", key).unwrap();

            let request = Request::new(
                Method::GET,
                Uri::from_static("http://localhost/ws"),
                Version::HTTP_11,
                Query::new(),
                headers,
                Bytes::new(),
            );

            let ctx = Context::new(request, Arc::new(()));
            let ws = WebSocketUpgrade::new(ctx);

            accept_keys.push(ws.accept_key().unwrap());
        }

        // All accept keys should be unique
        assert_eq!(accept_keys.len(), 3);
        assert_ne!(accept_keys[0], accept_keys[1]);
        assert_ne!(accept_keys[1], accept_keys[2]);
        assert_ne!(accept_keys[0], accept_keys[2]);
    }

    #[test]
    fn test_websocket_guid_constant() {
        // Verify the GUID constant matches RFC 6455
        assert_eq!(WEBSOCKET_GUID, "258EAFA5-E914-47DA-95CA-C5AB0DC85B11");
    }

    // RFC 6455 Compliance Tests

    #[test]
    fn test_accept_key_invalid_length_too_short() {
        let mut headers = Headers::new();
        headers.set("upgrade", "websocket").unwrap();
        headers.set("sec-websocket-key", "short").unwrap(); // Too short

        let request = Request::new(
            Method::GET,
            Uri::from_static("http://localhost/ws"),
            Version::HTTP_11,
            Query::new(),
            headers,
            Bytes::new(),
        );

        let ctx = Context::new(request, Arc::new(()));
        let ws = WebSocketUpgrade::new(ctx);

        let result = ws.accept_key();
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("Invalid Sec-WebSocket-Key length"));
    }

    #[test]
    fn test_accept_key_invalid_length_too_long() {
        let mut headers = Headers::new();
        headers.set("upgrade", "websocket").unwrap();
        headers
            .set("sec-websocket-key", "dGhlIHNhbXBsZSBub25jZQ==EXTRA")
            .unwrap(); // Too long

        let request = Request::new(
            Method::GET,
            Uri::from_static("http://localhost/ws"),
            Version::HTTP_11,
            Query::new(),
            headers,
            Bytes::new(),
        );

        let ctx = Context::new(request, Arc::new(()));
        let ws = WebSocketUpgrade::new(ctx);

        let result = ws.accept_key();
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("Invalid Sec-WebSocket-Key length"));
    }

    #[test]
    fn test_accept_key_invalid_base64() {
        let mut headers = Headers::new();
        headers.set("upgrade", "websocket").unwrap();
        headers
            .set("sec-websocket-key", "not!valid@base64#chars")
            .unwrap(); // Invalid base64

        let request = Request::new(
            Method::GET,
            Uri::from_static("http://localhost/ws"),
            Version::HTTP_11,
            Query::new(),
            headers,
            Bytes::new(),
        );

        let ctx = Context::new(request, Arc::new(()));
        let ws = WebSocketUpgrade::new(ctx);

        let result = ws.accept_key();
        assert!(result.is_err());
        let error_msg = result.unwrap_err().to_string();
        assert!(
            error_msg.contains("not valid base64") || error_msg.contains("Invalid"),
            "Expected base64 error, got: {}",
            error_msg
        );
    }

    #[test]
    fn test_accept_key_valid_base64_wrong_decoded_length() {
        let mut headers = Headers::new();
        headers.set("upgrade", "websocket").unwrap();
        // Valid base64, 24 chars, but decodes to 17 bytes instead of 16
        // "dGVzdGluZzEyMzQ1Njc4OTAx" = "testing1234567890" (17 bytes)
        headers
            .set("sec-websocket-key", "dGVzdGluZzEyMzQ1Njc4OTAx")
            .unwrap();

        let request = Request::new(
            Method::GET,
            Uri::from_static("http://localhost/ws"),
            Version::HTTP_11,
            Query::new(),
            headers,
            Bytes::new(),
        );

        let ctx = Context::new(request, Arc::new(()));
        let ws = WebSocketUpgrade::new(ctx);

        let result = ws.accept_key();
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("decoded to"));
    }

    #[test]
    fn test_accept_key_rfc6455_compliance() {
        // Test multiple valid RFC 6455 compliant keys
        let valid_keys = vec![
            "dGhlIHNhbXBsZSBub25jZQ==", // RFC 6455 example (16 bytes)
            "x3JJHMbDL1EzLkh9GBhXDw==", // Valid 16-byte key
            "AAAAAAAAAAAAAAAAAAAAAA==", // All zeros (16 bytes)
        ];

        for key in valid_keys {
            let mut headers = Headers::new();
            headers.set("upgrade", "websocket").unwrap();
            headers.set("sec-websocket-key", key).unwrap();

            let request = Request::new(
                Method::GET,
                Uri::from_static("http://localhost/ws"),
                Version::HTTP_11,
                Query::new(),
                headers,
                Bytes::new(),
            );

            let ctx = Context::new(request, Arc::new(()));
            let ws = WebSocketUpgrade::new(ctx);

            let result = ws.accept_key();
            assert!(
                result.is_ok(),
                "Valid key '{}' should be accepted, but got error: {:?}",
                key,
                result.unwrap_err()
            );
        }
    }

    #[test]
    fn test_accept_key_empty_string() {
        let mut headers = Headers::new();
        headers.set("upgrade", "websocket").unwrap();
        headers.set("sec-websocket-key", "").unwrap(); // Empty

        let request = Request::new(
            Method::GET,
            Uri::from_static("http://localhost/ws"),
            Version::HTTP_11,
            Query::new(),
            headers,
            Bytes::new(),
        );

        let ctx = Context::new(request, Arc::new(()));
        let ws = WebSocketUpgrade::new(ctx);

        let result = ws.accept_key();
        assert!(result.is_err());
    }

    #[test]
    fn test_accept_key_whitespace() {
        let mut headers = Headers::new();
        headers.set("upgrade", "websocket").unwrap();
        headers
            .set("sec-websocket-key", "                        ")
            .unwrap(); // 24 spaces

        let request = Request::new(
            Method::GET,
            Uri::from_static("http://localhost/ws"),
            Version::HTTP_11,
            Query::new(),
            headers,
            Bytes::new(),
        );

        let ctx = Context::new(request, Arc::new(()));
        let ws = WebSocketUpgrade::new(ctx);

        let result = ws.accept_key();
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("not valid base64 encoding"));
    }
}
