botui/src/ui_server/mod.rs

//! UI Server module for BotUI
//!
//! Serves the web UI (suite, minimal) and handles API proxying.

use axum::{
    body::Body,
    extract::{
        ws::{Message as AxumMessage, WebSocket, WebSocketUpgrade},
        OriginalUri, Query, State,
    },
    http::{Request, StatusCode},
    response::{Html, IntoResponse, Response},
    routing::{any, get},
    Router,
};
use futures_util::{SinkExt, StreamExt};
use log::{debug, error, info};
use serde::Deserialize;
use std::{fs, path::PathBuf};
use tokio_tungstenite::{
    connect_async_tls_with_config, tungstenite::protocol::Message as TungsteniteMessage,
};

use crate::shared::AppState;

/// Serve the index page (suite UI)
pub async fn index() -> impl IntoResponse {
    serve_suite().await
}

/// Handler for minimal UI
pub async fn serve_minimal() -> impl IntoResponse {
    match fs::read_to_string("ui/minimal/index.html") {
        Ok(html) => (StatusCode::OK, [("content-type", "text/html")], Html(html)),
        Err(e) => {
            error!("Failed to load minimal UI: {}", e);
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                [("content-type", "text/plain")],
                Html("Failed to load minimal interface".to_string()),
            )
        }
    }
}

/// Handler for suite UI
pub async fn serve_suite() -> impl IntoResponse {
    match fs::read_to_string("ui/suite/index.html") {
        Ok(html) => (StatusCode::OK, [("content-type", "text/html")], Html(html)),
        Err(e) => {
            error!("Failed to load suite UI: {}", e);
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                [("content-type", "text/plain")],
                Html("Failed to load suite interface".to_string()),
            )
        }
    }
}

/// Health check endpoint - checks BotServer connectivity
async fn health(State(state): State<AppState>) -> (StatusCode, axum::Json<serde_json::Value>) {
    match state.health_check().await {
        true => (
            StatusCode::OK,
            axum::Json(serde_json::json!({
                "status": "healthy",
                "service": "botui",
                "mode": "web"
            })),
        ),
        false => (
            StatusCode::SERVICE_UNAVAILABLE,
            axum::Json(serde_json::json!({
                "status": "unhealthy",
                "service": "botui",
                "error": "botserver unreachable"
            })),
        ),
    }
}

/// API health check endpoint
async fn api_health() -> (StatusCode, axum::Json<serde_json::Value>) {
    (
        StatusCode::OK,
        axum::Json(serde_json::json!({
            "status": "ok",
            "version": env!("CARGO_PKG_VERSION")
        })),
    )
}

/// Extract app context from Referer header or path
fn extract_app_context(headers: &axum::http::HeaderMap, path: &str) -> Option<String> {
    // Try to extract from Referer header first
    if let Some(referer) = headers.get("referer") {
        if let Ok(referer_str) = referer.to_str() {
            // Match /apps/{app_name}/ pattern
            if let Some(start) = referer_str.find("/apps/") {
                let after_apps = &referer_str[start + 6..];
                if let Some(end) = after_apps.find('/') {
                    return Some(after_apps[..end].to_string());
                } else if !after_apps.is_empty() {
                    return Some(after_apps.to_string());
                }
            }
        }
    }

    // Try to extract from path (for /apps/{app}/api/* routes)
    if path.starts_with("/apps/") {
        let after_apps = &path[6..];
        if let Some(end) = after_apps.find('/') {
            return Some(after_apps[..end].to_string());
        }
    }

    None
}

/// Proxy API requests to botserver
async fn proxy_api(
    State(state): State<AppState>,
    original_uri: OriginalUri,
    req: Request<Body>,
) -> Response<Body> {
    let path = original_uri.path();
    let query = original_uri
        .query()
        .map(|q| format!("?{}", q))
        .unwrap_or_default();
    let method = req.method().clone();
    let headers = req.headers().clone();

    // Extract app context from request
    let app_context = extract_app_context(&headers, path);

    let target_url = format!("{}{}{}", state.client.base_url(), path, query);
    debug!(
        "Proxying {} {} to {} (app: {:?})",
        method, path, target_url, app_context
    );

    // Build the proxied request with self-signed cert support
    let client = reqwest::Client::builder()
        .danger_accept_invalid_certs(true)
        .build()
        .unwrap_or_else(|_| reqwest::Client::new());
    let mut proxy_req = client.request(method.clone(), &target_url);

    // Copy headers (excluding host)
    for (name, value) in headers.iter() {
        if name != "host" {
            if let Ok(v) = value.to_str() {
                proxy_req = proxy_req.header(name.as_str(), v);
            }
        }
    }

    // Inject X-App-Context header if we detected an app
    if let Some(app) = app_context {
        proxy_req = proxy_req.header("X-App-Context", app);
    }

    // Copy body for non-GET requests
    let body_bytes = match axum::body::to_bytes(req.into_body(), usize::MAX).await {
        Ok(bytes) => bytes,
        Err(e) => {
            error!("Failed to read request body: {}", e);
            return Response::builder()
                .status(StatusCode::INTERNAL_SERVER_ERROR)
                .body(Body::from("Failed to read request body"))
                .unwrap();
        }
    };

    if !body_bytes.is_empty() {
        proxy_req = proxy_req.body(body_bytes.to_vec());
    }

    // Execute the request
    match proxy_req.send().await {
        Ok(resp) => {
            let status = resp.status();
            let headers = resp.headers().clone();

            match resp.bytes().await {
                Ok(body) => {
                    let mut response = Response::builder().status(status);

                    // Copy response headers
                    for (name, value) in headers.iter() {
                        response = response.header(name, value);
                    }

                    response.body(Body::from(body)).unwrap_or_else(|_| {
                        Response::builder()
                            .status(StatusCode::INTERNAL_SERVER_ERROR)
                            .body(Body::from("Failed to build response"))
                            .unwrap()
                    })
                }
                Err(e) => {
                    error!("Failed to read response body: {}", e);
                    Response::builder()
                        .status(StatusCode::BAD_GATEWAY)
                        .body(Body::from(format!("Failed to read response: {}", e)))
                        .unwrap()
                }
            }
        }
        Err(e) => {
            error!("Proxy request failed: {}", e);
            Response::builder()
                .status(StatusCode::BAD_GATEWAY)
                .body(Body::from(format!("Proxy error: {}", e)))
                .unwrap()
        }
    }
}

/// Create API proxy router
fn create_api_router() -> Router<AppState> {
    Router::new()
        .route("/health", get(api_health))
        .route("/chat", any(proxy_api))
        .route("/sessions", any(proxy_api))
        .route("/sessions/{id}", any(proxy_api))
        .route("/sessions/{id}/history", any(proxy_api))
        .route("/sessions/{id}/start", any(proxy_api))
        .route("/drive/files", any(proxy_api))
        .route("/drive/files/{path}", any(proxy_api))
        .route("/drive/upload", any(proxy_api))
        .route("/drive/download/{path}", any(proxy_api))
        .route("/tasks", any(proxy_api))
        .route("/tasks/{id}", any(proxy_api))
        .fallback(any(proxy_api))
}

/// WebSocket query parameters
#[derive(Debug, Deserialize)]
struct WsQuery {
    session_id: String,
    user_id: String,
}

/// WebSocket proxy handler
async fn ws_proxy(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    Query(params): Query<WsQuery>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_ws_proxy(socket, state, params))
}

/// Handle WebSocket proxy connection
async fn handle_ws_proxy(client_socket: WebSocket, state: AppState, params: WsQuery) {
    let backend_url = format!(
        "{}/ws?session_id={}&user_id={}",
        state
            .client
            .base_url()
            .replace("https://", "wss://")
            .replace("http://", "ws://"),
        params.session_id,
        params.user_id
    );

    info!("Proxying WebSocket to: {}", backend_url);

    // Create TLS connector that accepts self-signed certs
    let tls_connector = native_tls::TlsConnector::builder()
        .danger_accept_invalid_certs(true)
        .danger_accept_invalid_hostnames(true)
        .build()
        .expect("Failed to build TLS connector");

    let connector = tokio_tungstenite::Connector::NativeTls(tls_connector);

    // Connect to backend WebSocket
    let backend_result =
        connect_async_tls_with_config(&backend_url, None, false, Some(connector)).await;

    let backend_socket = match backend_result {
        Ok((socket, _)) => socket,
        Err(e) => {
            error!("Failed to connect to backend WebSocket: {}", e);
            return;
        }
    };

    info!("Connected to backend WebSocket");

    // Split both sockets
    let (mut client_tx, mut client_rx) = client_socket.split();
    let (mut backend_tx, mut backend_rx) = backend_socket.split();

    // Forward messages from client to backend
    let client_to_backend = async {
        while let Some(msg) = client_rx.next().await {
            match msg {
                Ok(AxumMessage::Text(text)) => {
                    if backend_tx
                        .send(TungsteniteMessage::Text(text))
                        .await
                        .is_err()
                    {
                        break;
                    }
                }
                Ok(AxumMessage::Binary(data)) => {
                    if backend_tx
                        .send(TungsteniteMessage::Binary(data))
                        .await
                        .is_err()
                    {
                        break;
                    }
                }
                Ok(AxumMessage::Ping(data)) => {
                    if backend_tx
                        .send(TungsteniteMessage::Ping(data))
                        .await
                        .is_err()
                    {
                        break;
                    }
                }
                Ok(AxumMessage::Pong(data)) => {
                    if backend_tx
                        .send(TungsteniteMessage::Pong(data))
                        .await
                        .is_err()
                    {
                        break;
                    }
                }
                Ok(AxumMessage::Close(_)) | Err(_) => break,
            }
        }
    };

    // Forward messages from backend to client
    let backend_to_client = async {
        while let Some(msg) = backend_rx.next().await {
            match msg {
                Ok(TungsteniteMessage::Text(text)) => {
                    if client_tx.send(AxumMessage::Text(text)).await.is_err() {
                        break;
                    }
                }
                Ok(TungsteniteMessage::Binary(data)) => {
                    if client_tx.send(AxumMessage::Binary(data)).await.is_err() {
                        break;
                    }
                }
                Ok(TungsteniteMessage::Ping(data)) => {
                    if client_tx.send(AxumMessage::Ping(data)).await.is_err() {
                        break;
                    }
                }
                Ok(TungsteniteMessage::Pong(data)) => {
                    if client_tx.send(AxumMessage::Pong(data)).await.is_err() {
                        break;
                    }
                }
                Ok(TungsteniteMessage::Close(_)) | Err(_) => break,
                _ => {}
            }
        }
    };

    // Run both forwarding tasks concurrently
    tokio::select! {
        _ = client_to_backend => info!("Client connection closed"),
        _ = backend_to_client => info!("Backend connection closed"),
    }
}

/// Create WebSocket proxy router
fn create_ws_router() -> Router<AppState> {
    Router::new().fallback(any(ws_proxy))
}

/// Create apps proxy router - proxies /apps/* to botserver
fn create_apps_router() -> Router<AppState> {
    Router::new()
        // Proxy all /apps/* requests to botserver
        // botserver serves static files from site_path
        // API calls from apps also go through here with X-App-Context header
        .fallback(any(proxy_api))
}

/// Create UI HTMX proxy router (for HTML fragment endpoints)
fn create_ui_router() -> Router<AppState> {
    Router::new()
        // Email UI endpoints
        .route("/email/accounts", any(proxy_api))
        .route("/email/list", any(proxy_api))
        .route("/email/folders", any(proxy_api))
        .route("/email/compose", any(proxy_api))
        .route("/email/labels", any(proxy_api))
        .route("/email/templates", any(proxy_api))
        .route("/email/signatures", any(proxy_api))
        .route("/email/rules", any(proxy_api))
        .route("/email/search", any(proxy_api))
        .route("/email/auto-responder", any(proxy_api))
        .route("/email/{id}", any(proxy_api))
        .route("/email/{id}/delete", any(proxy_api))
        // Calendar UI endpoints
        .route("/calendar/list", any(proxy_api))
        .route("/calendar/upcoming", any(proxy_api))
        .route("/calendar/event/new", any(proxy_api))
        .route("/calendar/new", any(proxy_api))
        // Fallback for any other /ui/* routes
        .fallback(any(proxy_api))
}

/// Configure and return the main router
pub fn configure_router() -> Router {
    let suite_path = PathBuf::from("./ui/suite");
    let _minimal_path = PathBuf::from("./ui/minimal");
    let state = AppState::new();

    Router::new()
        // Health check endpoints
        .route("/health", get(health))
        // API proxy routes
        .nest("/api", create_api_router())
        // UI HTMX proxy routes (for /ui/* endpoints that return HTML fragments)
        .nest("/ui", create_ui_router())
        // WebSocket proxy routes
        .nest("/ws", create_ws_router())
        // Apps proxy routes - proxy /apps/* to botserver (which serves from site_path)
        .nest("/apps", create_apps_router())
        // UI routes
        .route("/", get(index))
        .route("/minimal", get(serve_minimal))
        .route("/suite", get(serve_suite))
        // Suite static assets (when accessing /suite/*)
        .nest_service(
            "/suite/js",
            tower_http::services::ServeDir::new(suite_path.join("js")),
        )
        .nest_service(
            "/suite/css",
            tower_http::services::ServeDir::new(suite_path.join("css")),
        )
        .nest_service(
            "/suite/public",
            tower_http::services::ServeDir::new(suite_path.join("public")),
        )
        .nest_service(
            "/suite/drive",
            tower_http::services::ServeDir::new(suite_path.join("drive")),
        )
        .nest_service(
            "/suite/chat",
            tower_http::services::ServeDir::new(suite_path.join("chat")),
        )
        .nest_service(
            "/suite/mail",
            tower_http::services::ServeDir::new(suite_path.join("mail")),
        )
        .nest_service(
            "/suite/tasks",
            tower_http::services::ServeDir::new(suite_path.join("tasks")),
        )
        .nest_service(
            "/suite/calendar",
            tower_http::services::ServeDir::new(suite_path.join("calendar")),
        )
        .nest_service(
            "/suite/meet",
            tower_http::services::ServeDir::new(suite_path.join("meet")),
        )
        .nest_service(
            "/suite/paper",
            tower_http::services::ServeDir::new(suite_path.join("paper")),
        )
        .nest_service(
            "/suite/research",
            tower_http::services::ServeDir::new(suite_path.join("research")),
        )
        .nest_service(
            "/suite/analytics",
            tower_http::services::ServeDir::new(suite_path.join("analytics")),
        )
        .nest_service(
            "/suite/monitoring",
            tower_http::services::ServeDir::new(suite_path.join("monitoring")),
        )
        .nest_service(
            "/suite/admin",
            tower_http::services::ServeDir::new(suite_path.join("admin")),
        )
        .nest_service(
            "/suite/auth",
            tower_http::services::ServeDir::new(suite_path.join("auth")),
        )
        .nest_service(
            "/suite/settings",
            tower_http::services::ServeDir::new(suite_path.join("settings")),
        )
        .nest_service(
            "/suite/sources",
            tower_http::services::ServeDir::new(suite_path.join("sources")),
        )
        .nest_service(
            "/suite/attendant",
            tower_http::services::ServeDir::new(suite_path.join("attendant")),
        )
        .nest_service(
            "/suite/tools",
            tower_http::services::ServeDir::new(suite_path.join("tools")),
        )
        .nest_service(
            "/suite/assets",
            tower_http::services::ServeDir::new(suite_path.join("assets")),
        )
        .nest_service(
            "/suite/partials",
            tower_http::services::ServeDir::new(suite_path.join("partials")),
        )
        // Legacy paths for backward compatibility (serve suite assets)
        .nest_service(
            "/js",
            tower_http::services::ServeDir::new(suite_path.join("js")),
        )
        .nest_service(
            "/css",
            tower_http::services::ServeDir::new(suite_path.join("css")),
        )
        .nest_service(
            "/public",
            tower_http::services::ServeDir::new(suite_path.join("public")),
        )
        .nest_service(
            "/drive",
            tower_http::services::ServeDir::new(suite_path.join("drive")),
        )
        .nest_service(
            "/chat",
            tower_http::services::ServeDir::new(suite_path.join("chat")),
        )
        .nest_service(
            "/mail",
            tower_http::services::ServeDir::new(suite_path.join("mail")),
        )
        .nest_service(
            "/tasks",
            tower_http::services::ServeDir::new(suite_path.join("tasks")),
        )
        // Additional app routes
        .nest_service(
            "/paper",
            tower_http::services::ServeDir::new(suite_path.join("paper")),
        )
        .nest_service(
            "/calendar",
            tower_http::services::ServeDir::new(suite_path.join("calendar")),
        )
        .nest_service(
            "/research",
            tower_http::services::ServeDir::new(suite_path.join("research")),
        )
        .nest_service(
            "/meet",
            tower_http::services::ServeDir::new(suite_path.join("meet")),
        )
        .nest_service(
            "/analytics",
            tower_http::services::ServeDir::new(suite_path.join("analytics")),
        )
        .nest_service(
            "/monitoring",
            tower_http::services::ServeDir::new(suite_path.join("monitoring")),
        )
        .nest_service(
            "/admin",
            tower_http::services::ServeDir::new(suite_path.join("admin")),
        )
        .nest_service(
            "/auth",
            tower_http::services::ServeDir::new(suite_path.join("auth")),
        )
        .nest_service(
            "/settings",
            tower_http::services::ServeDir::new(suite_path.join("settings")),
        )
        .nest_service(
            "/sources",
            tower_http::services::ServeDir::new(suite_path.join("sources")),
        )
        .nest_service(
            "/tools",
            tower_http::services::ServeDir::new(suite_path.join("tools")),
        )
        .nest_service(
            "/assets",
            tower_http::services::ServeDir::new(suite_path.join("assets")),
        )
        .nest_service(
            "/partials",
            tower_http::services::ServeDir::new(suite_path.join("partials")),
        )
        .nest_service(
            "/attendant",
            tower_http::services::ServeDir::new(suite_path.join("attendant")),
        )
        // Fallback for other static files (serve suite by default)
        .fallback_service(
            tower_http::services::ServeDir::new(suite_path.clone()).fallback(
                tower_http::services::ServeDir::new(suite_path)
                    .append_index_html_on_directories(true),
            ),
        )
        .with_state(state)
}