mod cli;
mod config;
mod db;
mod egress_tcp;
mod http;
mod json_util;
mod logger;
mod microservice_proxy;
mod overview;
mod performance;
mod provider_registry;
mod scheduler;
mod ssh_data_channel;
mod ssh_bridge;
mod state;
mod task_dispatcher;
mod types;

use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Context;
use axum::extract::ws::Message;
use axum::extract::{Request, State, WebSocketUpgrade};
use axum::http::HeaderMap;
use axum::response::Response;
use axum::routing::get;
use axum::Router;
use tokio::net::TcpListener;

use crate::config::read_config;
use crate::db::init_database;
use crate::http::{error_json, json_response};
use crate::logger::Logger;
use crate::performance::record_request_performance;
use crate::provider_registry::{mark_stale_providers_offline, provider_ws};
use crate::scheduler::{recover_scheduled_runs, run_due_scheduled_tasks};
use crate::ssh_bridge::ssh_ws;
use crate::ssh_data_channel::spawn_ssh_data_listener;
use crate::state::AppState;
use crate::task_dispatcher::mark_stale_tasks_failed;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    if cli::handle_cli().await? {
        return Ok(());
    }

    let config = read_config()?;
    let logger = Logger::new("backend-core", &config.log_file)?;
    let state = Arc::new(AppState::new(config, logger).await?);

    match init_database(&state).await {
        Ok(()) => {
            state.set_db_ready(true);
            state.log("info", "database_ready", None);
            if let Err(error) = recover_scheduled_runs(&state).await {
                state.log(
                    "error",
                    "scheduled_recovery_failed",
                    Some(error_json(&error)),
                );
            }
        }
        Err(error) => {
            state.log("error", "database_init_failed", Some(error_json(&error)));
            return Err(error).context("database initialization failed");
        }
    }

    spawn_sweeps(state.clone());

    let api_addr = SocketAddr::from(([0, 0, 0, 0], state.config.port));
    let provider_addr = SocketAddr::from(([0, 0, 0, 0], state.config.provider_port));
    let provider_data_addr = SocketAddr::from(([0, 0, 0, 0], state.config.provider_data_port));
    let api = Router::new()
        .route("/ws/ssh", get(ssh_ws_handler))
        .fallback(api_handler)
        .with_state(state.clone());
    let provider = Router::new()
        .route("/", get(provider_health_handler))
        .route("/health", get(provider_health_handler))
        .route("/ws/provider", get(provider_ws_handler))
        .fallback(provider_not_found_handler)
        .with_state(state.clone());

    let api_listener = TcpListener::bind(api_addr).await?;
    let provider_listener = TcpListener::bind(provider_addr).await?;
    let provider_data_listener = TcpListener::bind(provider_data_addr).await?;
    state.log(
        "info",
        "server_listening",
        Some(serde_json::json!({
            "apiUrl": format!("http://0.0.0.0:{}", state.config.port),
            "providerIngressUrl": format!("ws://0.0.0.0:{}/ws/provider", state.config.provider_port),
            "providerDataTcpUrl": format!("tcp://0.0.0.0:{}", state.config.provider_data_port),
            "logFile": state.config.log_file,
        })),
    );

    let api_server = axum::serve(api_listener, api);
    let provider_server = axum::serve(provider_listener, provider);
    let provider_data_server = spawn_ssh_data_listener(state.clone(), provider_data_listener);
    tokio::try_join!(api_server, provider_server, provider_data_server)?;
    Ok(())
}

fn spawn_sweeps(state: Arc<AppState>) {
    tokio::spawn({
        let state = state.clone();
        async move {
            let mut interval = tokio::time::interval(Duration::from_secs(10));
            loop {
                interval.tick().await;
                if let Err(error) = mark_stale_providers_offline(&state).await {
                    state.log("error", "heartbeat_sweep_failed", Some(error_json(&error)));
                }
            }
        }
    });
    tokio::spawn({
        let state = state.clone();
        async move {
            let period = Duration::from_millis(state.config.task_pending_timeout_ms.min(60_000));
            let mut interval = tokio::time::interval(period);
            loop {
                interval.tick().await;
                if let Err(error) = mark_stale_tasks_failed(&state).await {
                    state.log(
                        "error",
                        "task_timeout_sweep_failed",
                        Some(error_json(&error)),
                    );
                }
            }
        }
    });
    tokio::spawn(async move {
        let mut interval = tokio::time::interval(Duration::from_secs(30));
        loop {
            interval.tick().await;
            if let Err(error) = run_due_scheduled_tasks(&state).await {
                state.log(
                    "error",
                    "scheduled_task_sweep_failed",
                    Some(error_json(&error)),
                );
            }
        }
    });
}

async fn api_handler(State(state): State<Arc<AppState>>, request: Request) -> Response {
    let method = request.method().clone();
    let path = request.uri().path().to_string();
    let started = std::time::Instant::now();
    let response = crate::http::route_api(state.clone(), request).await;
    record_request_performance(
        &state,
        method.as_str(),
        &path,
        response.status().as_u16(),
        started.elapsed(),
    );
    response
}

async fn ssh_ws_handler(
    State(state): State<Arc<AppState>>,
    ws: WebSocketUpgrade,
    headers: HeaderMap,
    request: Request,
) -> Response {
    let uri = request.uri().clone();
    let started = std::time::Instant::now();
    let response = ssh_ws(state.clone(), ws, headers, uri).await;
    record_request_performance(
        &state,
        "GET",
        "/ws/ssh",
        response.status().as_u16(),
        started.elapsed(),
    );
    response
}

async fn provider_health_handler(State(state): State<Arc<AppState>>) -> Response {
    json_response(
        serde_json::json!({
            "ok": true,
            "service": "unidesk-provider-ingress",
            "activeSocketCount": state.active_provider_count().await,
            "providerDataPort": state.config.provider_data_port,
        }),
        200,
    )
}

async fn provider_ws_handler(
    State(state): State<Arc<AppState>>,
    ws: WebSocketUpgrade,
    headers: HeaderMap,
    request: Request,
) -> Response {
    let uri = request.uri().clone();
    let started = std::time::Instant::now();
    let response = provider_ws(state.clone(), ws, headers, uri).await;
    record_request_performance(
        &state,
        "GET",
        "/ws/provider",
        response.status().as_u16(),
        started.elapsed(),
    );
    response
}

async fn provider_not_found_handler() -> Response {
    json_response(
        serde_json::json!({ "ok": false, "error": "not found" }),
        404,
    )
}

pub(crate) fn ws_text(value: serde_json::Value) -> Message {
    Message::Text(value.to_string())
}
