openserve 2.0.3

//! Server Module
//! 
//! Main server implementation with AI-enhanced features.

use anyhow::Result;
use axum::{
    routing::{get, post, put},
    Router,
    middleware,

};
use std::sync::Arc;
use tokio::net::TcpListener;
use tokio::signal;
use tower::ServiceBuilder;
use tower_http::{
    compression::CompressionLayer,
    cors::CorsLayer,
    services::ServeDir,
    trace::TraceLayer,
};
use tracing::{info, warn};

use crate::{
    config::Config,
    handlers,
    middleware as app_middleware,
    services::{FileService, SearchService},
    ai::AiService,
};

/// The main server struct that orchestrates all application components.
///
/// This struct holds references to all the core services and configuration
/// needed to run the OpenServe application, including file operations,
/// search functionality, and optional AI features.
pub struct Server {
    config: Arc<Config>,
    file_service: Arc<FileService>,
    search_service: Arc<SearchService>,
    ai_service: Option<Arc<AiService>>,
}

impl Server {
    /// Creates a new server instance with all necessary services initialized.
    ///
    /// This method initializes the file service, search service, and optionally
    /// the AI service based on the configuration. It performs all necessary
    /// setup and validation before returning the server instance.
    ///
    /// # Arguments
    ///
    /// * `config` - The application configuration wrapped in an `Arc`.
    ///
    /// # Returns
    ///
    /// A `Result` containing the initialized `Server` instance or an error
    /// if initialization fails.
    pub async fn new(config: Arc<Config>) -> Result<Self> {
        info!("Initializing server components");

        let file_service = Arc::new(FileService::new(config.clone()).await?);
        let search_service = Arc::new(SearchService::new(config.clone()).await?);
        
        let ai_service = if config.ai.enabled {
            info!("Initializing AI service");
            Some(Arc::new(AiService::new(Arc::new(config.ai.clone())).await?))
        } else {
            info!("AI features disabled");
            None
        };

        Ok(Self {
            config,
            file_service,
            search_service,
            ai_service,
        })
    }

    /// Builds the Axum router with all application routes and middleware.
    ///
    /// This method constructs the complete routing table for the application,
    /// including file operations, search endpoints, WebDAV support, and
    /// optionally AI endpoints. It also applies all necessary middleware
    /// layers for logging, compression, CORS, and authentication.
    ///
    /// # Returns
    ///
    /// A configured `Router` ready to handle HTTP requests.
    fn build_router(&self) -> Router {
        let mut router = Router::new()
            // File operations
            .route("/api/files/*path", get(handlers::files::list).post(handlers::files::upload))
            .route("/api/files/*path", put(handlers::files::update).delete(handlers::files::delete))
            .route("/api/files/download/*path", get(handlers::files::download))
            .route("/api/files/metadata/*path", get(handlers::files::metadata))
            
            // Search endpoints
            .route("/api/search", get(handlers::search::search))
            .route("/api/search/semantic", post(handlers::search::semantic_search))
            .route("/api/search/files", get(handlers::search::file_search))
            
            // WebDAV support
            .route("/dav/*path", axum::routing::any(handlers::webdav::handle_webdav))
            
            // Static file serving
            .nest_service("/", ServeDir::new(&self.config.server.root_dir)
                .precompressed_br()
                .precompressed_gzip())
            
            // Health check
            .route("/health", get(handlers::health::health_check))
            
            // Metrics
            .route("/metrics", get(handlers::metrics::metrics));

        // Add AI endpoints if enabled
        if self.ai_service.is_some() {
            router = router
                .route("/api/ai/analyze", post(handlers::ai::analyze_content))
                .route("/api/ai/summarize", post(handlers::ai::summarize))
                .route("/api/ai/organize", post(handlers::ai::suggest_organization))
                .route("/api/ai/chat", post(handlers::ai::chat));
        }

        // Apply middleware
        router
            .layer(
                ServiceBuilder::new()
                    .layer(TraceLayer::new_for_http())
                    .layer(CompressionLayer::new())
                    .layer(CorsLayer::permissive())
                    .layer(middleware::from_fn_with_state(
                        self.build_app_state(),
                        app_middleware::auth::auth_middleware
                    ))
                    .layer(middleware::from_fn(app_middleware::logging::logging_middleware))
            )
            .with_state(self.build_app_state())
    }

    /// Builds the application state that will be shared across all handlers.
    ///
    /// This method creates an `AppState` instance containing references to
    /// all the services and configuration that handlers need to access.
    ///
    /// # Returns
    ///
    /// An `AppState` instance with all necessary service references.
    fn build_app_state(&self) -> AppState {
        AppState {
            config: self.config.clone(),
            file_service: self.file_service.clone(),
            search_service: self.search_service.clone(),
            ai_service: self.ai_service.clone(),
        }
    }

    /// Starts the HTTP server and runs it until a shutdown signal is received.
    ///
    /// This method binds to the configured host and port, starts serving
    /// HTTP requests, and handles graceful shutdown when termination
    /// signals are received.
    ///
    /// # Returns
    ///
    /// A `Result` indicating whether the server ran successfully or
    /// encountered an error during startup or operation.
    pub async fn run(&self) -> Result<()> {
        let app = self.build_router();
        let addr = format!("{}:{}", self.config.server.host, self.config.server.port);
        
        info!("Starting server on {}", addr);
        
        let listener = TcpListener::bind(&addr).await?;
        
        axum::serve(listener, app.into_make_service())
            .with_graceful_shutdown(shutdown_signal())
            .await?;

        info!("Server stopped");
        Ok(())
    }
}

/// Application state shared across all HTTP handlers.
///
/// This struct contains references to all the services and configuration
/// that handlers need to process requests. It is cloned and passed to
/// each handler invocation.
#[derive(Clone)]
pub struct AppState {
    /// Application configuration
    pub config: Arc<Config>,
    /// File operations service
    pub file_service: Arc<FileService>,
    /// Search functionality service
    pub search_service: Arc<SearchService>,
    /// Optional AI service (if enabled)
    pub ai_service: Option<Arc<AiService>>,
}

/// Handles graceful shutdown signals (Ctrl+C and SIGTERM).
///
/// This function listens for shutdown signals and returns when one is
/// received, allowing the server to perform a graceful shutdown.
async fn shutdown_signal() {
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install signal handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }

    warn!("Shutdown signal received, starting graceful shutdown");
}