vex-web 0.6.5

use std::sync::Arc;

use anyhow::Result;
use axum::Router;
use axum::extract::ws::{Message, WebSocket};
use axum::extract::{State, WebSocketUpgrade};
use axum::http::HeaderMap;
use axum::response::{Html, IntoResponse, Response};
use axum::routing::get;
use futures_util::{SinkExt, StreamExt};
use rust_embed::Embed;
use tracing::info;
use vex_hub::{FrontendCommand, FrontendEvent, Hub};

#[derive(Embed)]
#[folder = "src/static/"]
struct Asset;

struct AppState {
    hub: Arc<Hub>,
}

pub async fn run(hub: Arc<Hub>, port: u16) -> Result<()> {
    let state = Arc::new(AppState { hub });

    let app = Router::new()
        .route("/", get(index_handler))
        .route("/ws", get(ws_handler))
        .route("/static/{*path}", get(static_handler))
        .fallback(get(index_handler))
        .with_state(state);

    let listener = tokio::net::TcpListener::bind(("0.0.0.0", port)).await?;
    info!("web UI listening on http://0.0.0.0:{}", port);

    axum::serve(listener, app).await?;
    Ok(())
}

async fn index_handler() -> Html<String> {
    match Asset::get("index.html") {
        Some(content) => Html(String::from_utf8_lossy(&content.data).to_string()),
        None => Html("<h1>vex web UI</h1><p>static files not found</p>".to_string()),
    }
}

async fn static_handler(axum::extract::Path(path): axum::extract::Path<String>) -> Response {
    match Asset::get(&path) {
        Some(content) => {
            let mime = mime_guess::from_path(&path).first_or_octet_stream();
            let mut headers = HeaderMap::new();
            headers.insert(
                axum::http::header::CONTENT_TYPE,
                mime.as_ref().parse().unwrap(),
            );
            if path.ends_with(".css") || path.ends_with(".js") {
                headers.insert(
                    axum::http::header::CACHE_CONTROL,
                    "public, max-age=3600".parse().unwrap(),
                );
            }
            (headers, content.data.to_vec()).into_response()
        }
        None => axum::http::StatusCode::NOT_FOUND.into_response(),
    }
}

async fn ws_handler(ws: WebSocketUpgrade, State(state): State<Arc<AppState>>) -> Response {
    ws.on_upgrade(move |socket| handle_ws(socket, state))
}

async fn handle_ws(socket: WebSocket, state: Arc<AppState>) {
    let (mut ws_tx, mut ws_rx) = socket.split();
    let mut event_rx = state.hub.event_rx();
    let command_tx = state.hub.command_tx();

    // Send initial state
    let initial_state = state.hub.state_rx().borrow().clone();
    let json = serde_json::json!({
        "type": "state_snapshot",
        "state": initial_state,
    })
    .to_string();
    let _ = ws_tx.send(Message::Text(json.into())).await;

    loop {
        tokio::select! {
            // Hub events -> WebSocket
            Ok(event) = event_rx.recv() => {
                let json = event_to_json(&event);
                if ws_tx.send(Message::Text(json.into())).await.is_err() {
                    break;
                }
            }
            // WebSocket messages -> Hub commands
            Some(Ok(msg)) = ws_rx.next() => {
                match msg {
                    Message::Text(text) => {
                        if let Ok(cmd) = serde_json::from_str::<FrontendCommand>(&text) {
                            let _ = command_tx.send(cmd).await;
                        }
                    }
                    Message::Close(_) => break,
                    _ => {}
                }
            }
            else => break,
        }
    }
}

fn event_to_json(event: &FrontendEvent) -> String {
    match event {
        FrontendEvent::StateSnapshot(state) => serde_json::json!({
            "type": "state_snapshot",
            "state": state,
        })
        .to_string(),
        FrontendEvent::ShellOutput { shell_id, data } => serde_json::json!({
            "type": "shell_output",
            "shell_id": shell_id.to_string(),
            "data": data,
        })
        .to_string(),
        FrontendEvent::AgentConversationLine { shell_id, line } => serde_json::json!({
            "type": "agent_conversation_line",
            "shell_id": shell_id.to_string(),
            "line": line,
        })
        .to_string(),
        FrontendEvent::AgentWatchEnd { shell_id } => serde_json::json!({
            "type": "agent_watch_end",
            "shell_id": shell_id.to_string(),
        })
        .to_string(),
        FrontendEvent::ShellBell { shell_id } => serde_json::json!({
            "type": "shell_bell",
            "shell_id": shell_id.to_string(),
        })
        .to_string(),
        FrontendEvent::IntrospectResult {
            path,
            suggested_name,
            git_remote,
            git_branch,
            children,
        } => serde_json::json!({
            "type": "introspect_result",
            "path": path,
            "suggested_name": suggested_name,
            "git_remote": git_remote,
            "git_branch": git_branch,
            "children": children,
        })
        .to_string(),
    }
}