wsforge_core/

router.rs

//! Routing and server management for WebSocket connections with middleware support.
//!
//! This module provides the core routing infrastructure for WsForge, allowing you to
//! define handlers for different message patterns, manage application state, add middleware
//! layers, and configure server behavior. The router handles both WebSocket connections
//! and static file serving on the same port.
//!
//! # Overview
//!
//! The [`Router`] is the main entry point for building a WebSocket server. It provides
//! a builder-style API for:
//! - Registering message handlers for specific routes
//! - Adding global and per-route middleware layers
//! - Managing shared application state
//! - Serving static files (HTML, CSS, JS)
//! - Configuring connection lifecycle callbacks
//! - Managing WebSocket connections
//!
//! # Architecture
//!
//! ```
//! ┌─────────────────┐
//! │  TCP Listener   │
//! └────────┬────────┘
//!          │
//!          ├──→ HTTP Request → Static File Handler
//!          │
//!          └──→ WebSocket Upgrade → Connection Manager
//!                                    │
//!                                    ├──→ on_connect callback
//!                                    │
//!                                    ├──→ Global Middleware Chain
//!                                    │    ├──→ Middleware 1
//!                                    │    ├──→ Middleware 2
//!                                    │    └──→ Route-Specific Middleware
//!                                    │         └──→ Handler
//!                                    │
//!                                    └──→ on_disconnect callback
//! ```
//!
//! # Middleware System
//!
//! WsForge supports a powerful middleware system with two types:
//! - **Global middleware**: Applied to all routes
//! - **Per-route middleware**: Applied only to specific routes
//!
//! Middleware execute in order: global first, then per-route, then the handler.
//!
//! # Examples
//!
//! ## Simple Echo Server
//!
//! ```
//! use wsforge::prelude::*;
//!
//! async fn echo(msg: Message) -> Result<Message> {
//!     Ok(msg)
//! }
//!
//! # async fn example() -> Result<()> {
//! let router = Router::new()
//!     .default_handler(handler(echo));
//!
//! router.listen("127.0.0.1:8080").await?;
//! # Ok(())
//! # }
//! ```
//!
//! ## With Global Middleware
//!
//! ```
//! use wsforge::prelude::*;
//!
//! async fn echo(msg: Message) -> Result<Message> {
//!     Ok(msg)
//! }
//!
//! # async fn example() -> Result<()> {
//! let router = Router::new()
//!     .layer(LoggerMiddleware::new())  // Logs all messages
//!     .default_handler(handler(echo));
//!
//! router.listen("127.0.0.1:8080").await?;
//! # Ok(())
//! # }
//! ```
//!
//! ## Multiple Routes with Per-Route Middleware
//!
//! ```
//! use wsforge::prelude::*;
//!
//! async fn public_handler(msg: Message) -> Result<String> {
//!     Ok("Public response".to_string())
//! }
//!
//! async fn admin_handler(msg: Message) -> Result<String> {
//!     Ok("Admin response".to_string())
//! }
//!
//! # async fn example() -> Result<()> {
//! # use std::sync::Arc;
//! let router = Router::new()
//!     // Public route - no auth needed
//!     .route("/public", handler(public_handler))
//!
//!     // Admin route - with auth middleware
//!     .route_with_layers(
//!         "/admin",
//!         vec![auth_middleware()],
//!         handler(admin_handler)
//!     );
//!
//! router.listen("127.0.0.1:8080").await?;
//! # Ok(())
//! # }
//! # fn auth_middleware() -> Arc<dyn Middleware> { unimplemented!() }
//! ```
//!
//! ## With State and Callbacks
//!
//! ```
//! use wsforge::prelude::*;
//! use std::sync::Arc;
//!
//! async fn broadcast(msg: Message, State(manager): State<Arc<ConnectionManager>>) -> Result<()> {
//!     manager.broadcast(msg);
//!     Ok(())
//! }
//!
//! # async fn example() -> Result<()> {
//! let router = Router::new()
//!     .layer(LoggerMiddleware::new())
//!     .default_handler(handler(broadcast))
//!     .on_connect(|manager, conn_id| {
//!         println!("✅ User {} connected (Total: {})", conn_id, manager.count());
//!     })
//!     .on_disconnect(|manager, conn_id| {
//!         println!("❌ User {} disconnected", conn_id);
//!     });
//!
//! router.listen("127.0.0.1:8080").await?;
//! # Ok(())
//! # }
//! ```

use crate::connection::{ConnectionId, ConnectionManager, handle_websocket};
use crate::error::{Error, Result};
use crate::extractor::Extensions;
use crate::handler::Handler;
use crate::message::Message;
use crate::middleware::{Middleware, MiddlewareChain};
use crate::state::AppState;
use dashmap::DashMap;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::net::{TcpListener, TcpStream};
use tokio_tungstenite::accept_async;
use tracing::{error, info};

/// Represents a single route with its path and middleware chain.
///
/// Routes map message patterns (paths) to handler functions with optional
/// middleware layers. This is typically used internally by the [`Router`].
///
/// # Examples
///
/// ```
/// use wsforge::prelude::*;
///
/// async fn my_handler() -> Result<String> {
///     Ok("response".to_string())
/// }
///
/// # fn example() {
/// let mut chain = MiddlewareChain::new();
/// chain.handler(handler(my_handler));
///
/// let route = Route {
///     path: "/api/message".to_string(),
///     chain: Arc::new(chain),
/// };
/// # }
/// ```
pub struct Route {
    /// The route path (e.g., "/chat", "/api/users")
    pub path: String,
    /// The middleware chain and handler for this route
    pub chain: Arc<MiddlewareChain>,
}

/// The main router for WebSocket servers with middleware support.
///
/// `Router` is the central component that manages routing, middleware, state, connections,
/// and server lifecycle. It uses a builder pattern for configuration and supports both
/// WebSocket and HTTP static file serving on the same port.
///
/// # Middleware System
///
/// The router supports two types of middleware:
/// - **Global middleware** (added with [`layer()`](Self::layer)) - Applied to all routes
/// - **Per-route middleware** (added with [`route_with_layers()`](Self::route_with_layers)) - Applied to specific routes
///
/// Execution order: Global middleware → Per-route middleware → Handler
///
/// # Thread Safety
///
/// Router is thread-safe and can be cloned cheaply (uses `Arc` internally).
/// All connections share the same router instance.
///
/// # Examples
///
/// ## Basic Setup
///
/// ```
/// use wsforge::prelude::*;
///
/// async fn handler(msg: Message) -> Result<String> {
///     Ok("received".to_string())
/// }
///
/// # async fn example() -> Result<()> {
/// let router = Router::new()
///     .default_handler(handler(handler));
///
/// router.listen("0.0.0.0:8080").await?;
/// # Ok(())
/// # }
/// ```
///
/// ## With Middleware Layers
///
/// ```
/// use wsforge::prelude::*;
///
/// # async fn handler(msg: Message) -> Result<String> { Ok("".to_string()) }
/// # async fn example() -> Result<()> {
/// let router = Router::new()
///     .layer(LoggerMiddleware::new())           // Global: logs all
///     .layer(auth_middleware())                  // Global: auth all
///     .route("/public", handler(handler))        // No extra middleware
///     .route_with_layers(
///         "/admin",
///         vec![admin_only_middleware()],         // Per-route: admin only
///         handler(handler)
///     );
///
/// router.listen("127.0.0.1:8080").await?;
/// # Ok(())
/// # }
/// # fn auth_middleware() -> Arc<dyn Middleware> { unimplemented!() }
/// # fn admin_only_middleware() -> Arc<dyn Middleware> { unimplemented!() }
/// ```
pub struct Router {
    routes: Arc<DashMap<String, Arc<MiddlewareChain>>>,
    global_middlewares: Vec<Arc<dyn Middleware>>,
    state: AppState,
    connection_manager: Arc<ConnectionManager>,
    on_connect: Option<Arc<dyn Fn(&Arc<ConnectionManager>, ConnectionId) + Send + Sync>>,
    on_disconnect: Option<Arc<dyn Fn(&Arc<ConnectionManager>, ConnectionId) + Send + Sync>>,
    default_chain: Option<Arc<MiddlewareChain>>,
    static_handler: Option<crate::static_files::StaticFileHandler>,
}

impl Router {
    /// Creates a new empty router.
    ///
    /// The router starts with no routes, no middleware, no state, and no handlers.
    /// Use the builder methods to configure it.
    ///
    /// # Examples
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// let router = Router::new();
    /// ```
    pub fn new() -> Self {
        Self {
            routes: Arc::new(DashMap::new()),
            global_middlewares: Vec::new(),
            state: AppState::new(),
            connection_manager: Arc::new(ConnectionManager::new()),
            on_connect: None,
            on_disconnect: None,
            default_chain: None,
            static_handler: None,
        }
    }

    /// Add a global middleware layer that applies to all routes.
    ///
    /// Global middleware are executed before per-route middleware and handlers.
    /// They are executed in the order they are added.
    ///
    /// # Arguments
    ///
    /// * `middleware` - The middleware to add
    ///
    /// # Examples
    ///
    /// ## Single Middleware
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// # fn example() {
    /// let router = Router::new()
    ///     .layer(LoggerMiddleware::new());
    /// # }
    /// ```
    ///
    /// ## Multiple Middleware
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// # fn example() {
    /// let router = Router::new()
    ///     .layer(LoggerMiddleware::new())        // First: logging
    ///     .layer(auth_middleware())               // Second: authentication
    ///     .layer(rate_limit_middleware());        // Third: rate limiting
    /// # }
    /// # fn auth_middleware() -> Arc<dyn Middleware> { unimplemented!() }
    /// # fn rate_limit_middleware() -> Arc<dyn Middleware> { unimplemented!() }
    /// ```
    pub fn layer(mut self, middleware: Arc<dyn Middleware>) -> Self {
        self.global_middlewares.push(middleware);
        self
    }

    /// Registers a handler for a specific route without additional middleware.
    ///
    /// Global middleware will still apply to this route. For route-specific middleware,
    /// use [`route_with_layers()`](Self::route_with_layers).
    ///
    /// Routes are matched against the beginning of incoming messages.
    /// For example, a message like `/chat hello` would match route `/chat`.
    ///
    /// # Arguments
    ///
    /// * `path` - The route path (e.g., "/chat", "/api/users")
    /// * `handler` - The handler function wrapped with `handler()`
    ///
    /// # Examples
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// async fn chat_handler(msg: Message) -> Result<String> {
    ///     Ok("chat response".to_string())
    /// }
    ///
    /// async fn api_handler(msg: Message) -> Result<String> {
    ///     Ok("api response".to_string())
    /// }
    ///
    /// # fn example() {
    /// let router = Router::new()
    ///     .route("/chat", handler(chat_handler))
    ///     .route("/api", handler(api_handler));
    /// # }
    /// ```
    pub fn route(self, path: impl Into<String>, handler: Arc<dyn Handler>) -> Self {
        let mut chain = MiddlewareChain::new();

        // Add global middlewares first
        for middleware in &self.global_middlewares {
            chain = chain.layer(middleware.clone());
        }

        // Add handler
        chain = chain.handler(handler);

        self.routes.insert(path.into(), Arc::new(chain));
        self
    }

    // Add a route with per-route middleware layers.
    ///
    /// Per-route middleware are executed after global middleware but before the handler.
    /// This is useful for route-specific concerns like authorization or validation.
    ///
    /// # Arguments
    ///
    /// * `path` - The route path (e.g., "/admin", "/api/users")
    /// * `layers` - Vector of middleware to apply to this route
    /// * `handler` - The handler function wrapped with `handler()`
    ///
    /// # Examples
    ///
    /// ## Admin Route with Auth
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// async fn admin_handler(msg: Message) -> Result<String> {
    ///     Ok("Admin panel".to_string())
    /// }
    ///
    /// # fn example() {
    /// # use std::sync::Arc;
    /// let router = Router::new()
    ///     .route_with_layers(
    ///         "/admin",
    ///         vec![
    ///             auth_middleware(),
    ///             admin_check_middleware(),
    ///         ],
    ///         handler(admin_handler)
    ///     );
    /// # }
    /// # fn auth_middleware() -> Arc<dyn Middleware> { unimplemented!() }
    /// # fn admin_check_middleware() -> Arc<dyn Middleware> { unimplemented!() }
    /// ```
    ///
    /// ## API Route with Validation
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// async fn api_handler(msg: Message) -> Result<String> {
    ///     Ok("API response".to_string())
    /// }
    ///
    /// # fn example() {
    /// # use std::sync::Arc;
    /// let router = Router::new()
    ///     .route_with_layers(
    ///         "/api/users",
    ///         vec![validate_json_middleware()],
    ///         handler(api_handler)
    ///     );
    /// # }
    /// # fn validate_json_middleware() -> Arc<dyn Middleware> { unimplemented!() }
    /// ```
    pub fn route_with_layers(
        self,
        path: impl Into<String>,
        layers: Vec<Arc<dyn Middleware>>,
        handler: Arc<dyn Handler>,
    ) -> Self {
        let mut chain = MiddlewareChain::new();

        // Add global middlewares first
        for middleware in &self.global_middlewares {
            chain = chain.layer(middleware.clone());
        }

        // Add route-specific middlewares
        for middleware in layers {
            chain = chain.layer(middleware);
        }

        // Add handler
        chain = chain.handler(handler);

        self.routes.insert(path.into(), Arc::new(chain));
        self
    }

    /// Adds shared state to the router.
    ///
    /// State is shared across all connections and can be extracted in handlers
    /// using the [`State`] extractor. Any type that is `Send + Sync + 'static`
    /// can be used as state.
    ///
    /// # Type Safety
    ///
    /// Multiple different types can be added as state. Each type is stored
    /// separately and retrieved by type.
    ///
    /// # Arguments
    ///
    /// * `data` - The state data to share (wrapped in `Arc`)
    ///
    /// # Examples
    ///
    /// ## Single State
    ///
    /// ```
    /// use wsforge::prelude::*;
    /// use std::sync::Arc;
    ///
    /// struct Database {
    ///     // database fields
    /// }
    ///
    /// async fn handler(State(db): State<Arc<Database>>) -> Result<String> {
    ///     Ok("query result".to_string())
    /// }
    ///
    /// # fn example() {
    /// let db = Arc::new(Database {});
    ///
    /// let router = Router::new()
    ///     .with_state(db)
    ///     .default_handler(handler(handler));
    /// # }
    /// ```
    ///
    /// ## Multiple States
    ///
    /// ```
    /// use wsforge::prelude::*;
    /// use std::sync::Arc;
    ///
    /// struct Config {
    ///     port: u16,
    /// }
    ///
    /// struct Database {
    ///     // database fields
    /// }
    ///
    /// # fn example() {
    /// let router = Router::new()
    ///     .with_state(Arc::new(Config { port: 8080 }))
    ///     .with_state(Arc::new(Database {}));
    /// # }
    /// ```
    pub fn with_state<T: Send + Sync + 'static>(self, data: Arc<T>) -> Self {
        self.state.insert(data);
        self
    }

    /// Sets a callback to be called when a new connection is established.
    ///
    /// The callback receives a reference to the connection manager and the
    /// connection ID. This is useful for logging, sending welcome messages,
    /// or updating user lists.
    ///
    /// # Arguments
    ///
    /// * `f` - Callback function with signature `Fn(&Arc<ConnectionManager>, ConnectionId)`
    ///
    /// # Examples
    ///
    /// ## Simple Logging
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// # fn example() {
    /// let router = Router::new()
    ///     .on_connect(|manager, conn_id| {
    ///         println!("New connection: {} (Total: {})", conn_id, manager.count());
    ///     });
    /// # }
    /// ```
    ///
    /// ## Send Welcome Message
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// # fn example() {
    /// let router = Router::new()
    ///     .on_connect(|manager, conn_id| {
    ///         if let Some(conn) = manager.get(&conn_id) {
    ///             let _ = conn.send_text("Welcome to the server!");
    ///         }
    ///     });
    /// # }
    /// ```
    ///
    /// ## Broadcast Join Notification
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// # fn example() {
    /// let router = Router::new()
    ///     .on_connect(|manager, conn_id| {
    ///         let msg = format!("User {} joined", conn_id);
    ///         manager.broadcast(Message::text(msg));
    ///     });
    /// # }
    /// ```
    pub fn on_connect<F>(mut self, f: F) -> Self
    where
        F: Fn(&Arc<ConnectionManager>, ConnectionId) + Send + Sync + 'static,
    {
        self.on_connect = Some(Arc::new(f));
        self
    }

    /// Sets a callback to be called when a connection is closed.
    ///
    /// The callback receives a reference to the connection manager and the
    /// connection ID. Note that the connection is already removed from the
    /// manager when this is called.
    ///
    /// # Arguments
    ///
    /// * `f` - Callback function with signature `Fn(&Arc<ConnectionManager>, ConnectionId)`
    ///
    /// # Examples
    ///
    /// ## Logging Disconnections
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// # fn example() {
    /// let router = Router::new()
    ///     .on_disconnect(|manager, conn_id| {
    ///         println!("Connection {} closed (Remaining: {})", conn_id, manager.count());
    ///     });
    /// # }
    /// ```
    ///
    /// ## Broadcast Leave Notification
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// # fn example() {
    /// let router = Router::new()
    ///     .on_disconnect(|manager, conn_id| {
    ///         let msg = format!("User {} left", conn_id);
    ///         manager.broadcast(Message::text(msg));
    ///     });
    /// # }
    /// ```
    pub fn on_disconnect<F>(mut self, f: F) -> Self
    where
        F: Fn(&Arc<ConnectionManager>, ConnectionId) + Send + Sync + 'static,
    {
        self.on_disconnect = Some(Arc::new(f));
        self
    }

    /// Sets the default handler for messages that don't match any route.
    ///
    /// This handler is called when no route matches the incoming message.
    /// Global middleware will still be applied to the default handler.
    ///
    /// # Arguments
    ///
    /// * `handler` - The default handler wrapped with `handler()`
    ///
    /// # Examples
    ///
    /// ## Echo Server
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// async fn echo(msg: Message) -> Result<Message> {
    ///     Ok(msg)
    /// }
    ///
    /// # fn example() {
    /// let router = Router::new()
    ///     .default_handler(handler(echo));
    /// # }
    /// ```
    ///
    /// ## Error Handler
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// async fn not_found() -> Result<String> {
    ///     Ok("Unknown command".to_string())
    /// }
    ///
    /// # fn example() {
    /// let router = Router::new()
    ///     .route("/known", handler(not_found))
    ///     .default_handler(handler(not_found));
    /// # }
    /// ```
    pub fn default_handler(mut self, handler: Arc<dyn Handler>) -> Self {
        let mut chain = MiddlewareChain::new();

        // Add global middlewares to default handler too
        for middleware in &self.global_middlewares {
            chain = chain.layer(middleware.clone());
        }

        chain = chain.handler(handler);
        self.default_chain = Some(Arc::new(chain));
        self
    }

    /// Enables static file serving from a directory.
    ///
    /// When enabled, the router will serve static files (HTML, CSS, JavaScript, images)
    /// from the specified directory for HTTP requests, while still handling WebSocket
    /// connections on the same port.
    ///
    /// # Path Resolution
    ///
    /// - Requests to `/` serve `index.html` from the directory
    /// - Other requests map directly to files (e.g., `/style.css` → `directory/style.css`)
    /// - MIME types are automatically detected
    ///
    /// # Security
    ///
    /// Path traversal attempts (e.g., `../../etc/passwd`) are automatically blocked.
    ///
    /// # Arguments
    ///
    /// * `path` - Path to the directory containing static files
    ///
    /// # Examples
    ///
    /// ## Serve Static Files
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// async fn ws_handler(msg: Message) -> Result<Message> {
    ///     Ok(msg)
    /// }
    ///
    /// # fn example() {
    /// let router = Router::new()
    ///     .serve_static("public")  // Serve files from ./public
    ///     .default_handler(handler(ws_handler));
    ///
    /// // Now you can access:
    /// // http://localhost:8080/          -> public/index.html
    /// // http://localhost:8080/app.js    -> public/app.js
    /// // ws://localhost:8080             -> WebSocket handler
    /// # }
    /// ```
    pub fn serve_static(mut self, path: impl Into<PathBuf>) -> Self {
        self.static_handler = Some(crate::static_files::StaticFileHandler::new(path.into()));
        self
    }

    /// Returns a reference to the connection manager.
    ///
    /// The connection manager is automatically created with the router.
    /// Use this to get access to it for storing in state or elsewhere.
    ///
    /// # Examples
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// # fn example() {
    /// let router = Router::new();
    /// let manager = router.connection_manager();
    ///
    /// // Now you can use the manager
    /// println!("Active connections: {}", manager.count());
    /// # }
    /// ```
    pub fn connection_manager(&self) -> Arc<ConnectionManager> {
        self.connection_manager.clone()
    }

    /// Starts the WebSocket server and listens for connections.
    ///
    /// This method consumes the router and starts the server loop. It will
    /// run indefinitely until the process is terminated or an error occurs.
    ///
    /// The connection manager is automatically inserted into the router's state
    /// before the server starts, making it available to all handlers via the
    /// `State<Arc<ConnectionManager>>` extractor.
    ///
    /// # Arguments
    ///
    /// * `addr` - The address to bind to (e.g., "127.0.0.1:8080", "0.0.0.0:3000")
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The address format is invalid
    /// - The port is already in use
    /// - Permission is denied (e.g., ports < 1024 on Unix)
    ///
    /// # Examples
    ///
    /// ## Basic Usage
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// # async fn example() -> Result<()> {
    /// let router = Router::new();
    /// router.listen("127.0.0.1:8080").await?;
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// ## All Interfaces
    ///
    /// ```
    /// use wsforge::prelude::*;
    ///
    /// # async fn example() -> Result<()> {
    /// let router = Router::new();
    /// router.listen("0.0.0.0:8080").await?;  // Accept connections from anywhere
    /// # Ok(())
    /// # }
    /// ```
    pub async fn listen(self, addr: impl AsRef<str>) -> Result<()> {
        let addr: SocketAddr = addr
            .as_ref()
            .parse()
            .map_err(|e| Error::custom(format!("Invalid address: {}", e)))?;

        self.state.insert(self.connection_manager.clone());

        let listener = TcpListener::bind(addr).await?;
        info!("WebSocket server listening on {}", addr);

        let router = Arc::new(self);

        loop {
            let (stream, peer_addr) = listener.accept().await?;
            let router = router.clone();

            tokio::spawn(async move {
                if let Err(e) = router.handle_connection(stream, peer_addr).await {
                    error!("Connection error: {}", e);
                }
            });
        }
    }

    async fn handle_connection(&self, stream: TcpStream, peer_addr: SocketAddr) -> Result<()> {
        let mut buffer = [0u8; 1024];

        let n = tokio::time::timeout(std::time::Duration::from_secs(5), stream.peek(&mut buffer))
            .await
            .map_err(|_| Error::custom("Connection timeout"))?
            .map_err(|e| Error::custom(format!("Failed to read: {}", e)))?;

        let header = String::from_utf8_lossy(&buffer[..n]);

        if header.contains("Upgrade: websocket") || header.contains("upgrade: websocket") {
            self.handle_websocket_connection(stream, peer_addr).await
        } else if let Some(ref static_handler) = self.static_handler {
            self.handle_http_request(stream, static_handler, &header)
                .await
        } else {
            Err(Error::custom("No handler for HTTP requests"))
        }
    }

    async fn handle_http_request(
        &self,
        mut stream: TcpStream,
        static_handler: &crate::static_files::StaticFileHandler,
        header: &str,
    ) -> Result<()> {
        use crate::static_files::http_response;
        use tokio::io::AsyncWriteExt;

        let path = header
            .lines()
            .next()
            .and_then(|line| {
                let parts: Vec<&str> = line.split_whitespace().collect();
                if parts.len() >= 2 && (parts[0] == "GET" || parts[0] == "HEAD") {
                    Some(parts[1])
                } else {
                    None
                }
            })
            .unwrap_or("/");

        let response = match static_handler.serve(path).await {
            Ok((content, mime_type)) => {
                info!("Served: {} ({} bytes)", path, content.len());
                http_response(200, &mime_type, content)
            }
            Err(e) => {
                tracing::warn!("File not found: {} - {}", path, e);
                let html = b"<html><body><h1>404 Not Found</h1></body></html>".to_vec();
                http_response(404, "text/html", html)
            }
        };

        stream.write_all(&response).await?;
        stream.flush().await?;
        Ok(())
    }

    async fn handle_websocket_connection(
        &self,
        stream: TcpStream,
        peer_addr: SocketAddr,
    ) -> Result<()> {
        let ws_stream = accept_async(stream).await?;
        let conn_id = Self::generate_connection_id();

        let router = self.clone();
        let manager = self.connection_manager.clone();

        let on_message = Arc::new(move |conn_id: ConnectionId, message: Message| {
            let router = router.clone();
            tokio::spawn(async move {
                if let Err(e) = router.handle_message(conn_id, message).await {
                    error!("Message handling error: {}", e);
                }
            });
        });

        let manager_ref = manager.clone();
        let on_connect = self
            .on_connect
            .clone()
            .map(move |cb| {
                let manager = manager_ref.clone();
                Arc::new(move |conn_id: ConnectionId| {
                    cb(&manager, conn_id);
                }) as Arc<dyn Fn(ConnectionId) + Send + Sync>
            })
            .unwrap_or_else(|| {
                Arc::new(|conn_id: ConnectionId| {
                    info!("Client connected: {}", conn_id);
                })
            });

        let manager_ref = manager.clone();
        let on_disconnect = self
            .on_disconnect
            .clone()
            .map(move |cb| {
                let manager = manager_ref.clone();
                Arc::new(move |conn_id: ConnectionId| {
                    cb(&manager, conn_id);
                }) as Arc<dyn Fn(ConnectionId) + Send + Sync>
            })
            .unwrap_or_else(|| {
                Arc::new(|conn_id: ConnectionId| {
                    info!("Client disconnected: {}", conn_id);
                })
            });

        handle_websocket(
            ws_stream,
            conn_id,
            peer_addr,
            manager,
            on_message,
            on_connect,
            on_disconnect,
        )
        .await;

        Ok(())
    }

    async fn handle_message(&self, conn_id: ConnectionId, message: Message) -> Result<()> {
        let conn = self
            .connection_manager
            .get(&conn_id)
            .ok_or_else(|| Error::ConnectionNotFound(conn_id.clone()))?;

        let extensions = Extensions::new();

        let chain = if let Some(text) = message.as_text() {
            if text.starts_with('/') {
                if let Some((route, _)) = text.split_once(' ') {
                    self.routes.get(route).map(|c| c.value().clone())
                } else {
                    self.routes.get(text).map(|c| c.value().clone())
                }
            } else {
                None
            }
        } else {
            None
        };

        let chain = chain.or_else(|| self.default_chain.clone());

        if let Some(chain) = chain {
            match chain
                .execute(message, conn.clone(), self.state.clone(), extensions)
                .await
            {
                Ok(Some(response)) => {
                    if let Err(e) = conn.send(response) {
                        error!("Failed to send response to {}: {}", conn_id, e);
                    }
                }
                Ok(None) => {
                    tracing::debug!("Handler processed message without response");
                }
                Err(e) => {
                    error!("Handler error for {}: {}", conn_id, e);
                }
            }
        } else {
            tracing::warn!("No handler found for message from {}", conn_id);
        }

        Ok(())
    }

    fn generate_connection_id() -> ConnectionId {
        use std::sync::atomic::{AtomicU64, Ordering};
        static COUNTER: AtomicU64 = AtomicU64::new(0);
        format!("conn_{}", COUNTER.fetch_add(1, Ordering::SeqCst))
    }
}

impl Clone for Router {
    fn clone(&self) -> Self {
        Self {
            routes: self.routes.clone(),
            global_middlewares: self.global_middlewares.clone(),
            state: self.state.clone(),
            connection_manager: self.connection_manager.clone(),
            on_connect: self.on_connect.clone(),
            on_disconnect: self.on_disconnect.clone(),
            default_chain: self.default_chain.clone(),
            static_handler: self.static_handler.clone(),
        }
    }
}

impl Default for Router {
    fn default() -> Self {
        Self::new()
    }
}