mod session;
mod ws_handler;
mod api;
mod auth;

use std::sync::Arc;
use axum::{
    Router,
    routing::{get, post},
    middleware,
};
use tokio::net::TcpListener;
use tracing::info;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

use session::SessionStore;
use auth::require_api_key;

#[derive(Clone)]
pub struct AppState {
    pub sessions: Arc<SessionStore>,
    pub api_key: String,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    tracing_subscriber::registry()
        .with(tracing_subscriber::EnvFilter::try_from_default_env()
            .unwrap_or_else(|_| "helios_server=debug,tower_http=info".into()))
        .with(tracing_subscriber::fmt::layer())
        .init();

    const GIT_COMMIT: &str = env!("GIT_COMMIT");
    info!("helios-server ({GIT_COMMIT})");

    let api_key = std::env::var("HELIOS_API_KEY")
        .unwrap_or_else(|_| "dev-secret".to_string());

    let bind_addr = std::env::var("HELIOS_BIND")
        .unwrap_or_else(|_| "0.0.0.0:3000".to_string());

    let state = AppState {
        sessions: Arc::new(SessionStore::new()),
        api_key,
    };

    let protected = Router::new()
        .route("/devices", get(api::list_devices))
        .route("/devices/:label/screenshot", post(api::request_screenshot))
        .route("/devices/:label/exec", post(api::request_exec))
        .route("/devices/:label/prompt", post(api::prompt_user))
        .route("/devices/:label/windows", get(api::list_windows))
        .route("/devices/:label/windows/minimize-all", post(api::minimize_all))
        .route("/devices/:label/logs", get(api::logs))
        .route("/devices/:label/windows/:window_id/screenshot", post(api::window_screenshot))
        .route("/devices/:label/windows/:window_id/focus", post(api::focus_window))
        .route("/devices/:label/windows/:window_id/maximize", post(api::maximize_and_focus))
        .route("/devices/:label/version", get(api::client_version))
        .route("/devices/:label/upload", post(api::upload_file))
        .route("/devices/:label/download", get(api::download_file))
        .route("/devices/:label/run", post(api::run_program))
        .route("/devices/:label/clipboard", get(api::clipboard_get))
        .route("/devices/:label/clipboard", post(api::clipboard_set))
        .layer(middleware::from_fn_with_state(state.clone(), require_api_key));

    let app = Router::new()
        .route("/ws", get(ws_handler::ws_upgrade))
        .route("/version", get(api::server_version))
        .merge(protected)
        .with_state(state);

    info!("helios-server listening on {bind_addr}");
    let listener = TcpListener::bind(&bind_addr).await?;
    axum::serve(listener, app).await?;
    Ok(())
}