use axum::body::Body;
use axum::extract::State;
use axum::http::StatusCode;
use axum::http::header::CONTENT_TYPE;
use axum::response::IntoResponse;
use axum::routing::{get, post};
use axum::{Json, Router};
use reqwest::Client;
use serde_json::{Value, json};
use sglangmux::{SgLangMux, SgLangMuxOptions};
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::{debug, error, info, warn};
use tracing_subscriber::EnvFilter;

#[derive(Clone)]
struct AppState {
    mux: SgLangMux,
    model_ports: HashMap<String, u16>,
    model_ids: Vec<String>,
    model_created_unix: u64,
    client: Client,
    upstream_timeout: Duration,
}

static REQUEST_COUNTER: AtomicU64 = AtomicU64::new(1);

fn init_tracing() {
    let filter = EnvFilter::try_from_default_env()
        .unwrap_or_else(|_| EnvFilter::new("sglangmux=debug,sglangmuxd=debug,warn"));

    let _ = tracing_subscriber::fmt()
        .with_env_filter(filter)
        .with_target(true)
        .with_thread_ids(true)
        .with_level(true)
        .try_init();
}

fn usage() -> ! {
    eprintln!(
        "Usage: sglangmuxd [--host HOST] [--listen-port PORT] [--upstream-timeout-secs SECONDS] [--model-ready-timeout-secs SECONDS] [--model-switch-timeout-secs SECONDS] [--log-dir DIR] <script1.sh> <script2.sh> ..."
    );
    std::process::exit(2);
}

fn parse_args() -> (String, u16, u64, u64, u64, PathBuf, Vec<PathBuf>) {
    let mut host = String::from("127.0.0.1");
    let mut listen_port: u16 = 30100;
    let mut upstream_timeout_secs: u64 = 120;
    let mut model_ready_timeout_secs: u64 = 120;
    let mut model_switch_timeout_secs: u64 = 60;
    let mut log_dir = PathBuf::from("sglangmux-logs");
    let mut scripts = Vec::new();

    let mut args = std::env::args().skip(1);
    while let Some(arg) = args.next() {
        match arg.as_str() {
            "--host" => {
                let Some(value) = args.next() else {
                    usage();
                };
                host = value;
            }
            "--listen-port" => {
                let Some(value) = args.next() else {
                    usage();
                };
                let Ok(port) = value.parse::<u16>() else {
                    usage();
                };
                listen_port = port;
            }
            "--upstream-timeout-secs" => {
                let Some(value) = args.next() else {
                    usage();
                };
                let Ok(seconds) = value.parse::<u64>() else {
                    usage();
                };
                if seconds == 0 {
                    usage();
                }
                upstream_timeout_secs = seconds;
            }
            "--model-ready-timeout-secs" => {
                let Some(value) = args.next() else {
                    usage();
                };
                let Ok(seconds) = value.parse::<u64>() else {
                    usage();
                };
                if seconds == 0 {
                    usage();
                }
                model_ready_timeout_secs = seconds;
            }
            "--model-switch-timeout-secs" => {
                let Some(value) = args.next() else {
                    usage();
                };
                let Ok(seconds) = value.parse::<u64>() else {
                    usage();
                };
                if seconds == 0 {
                    usage();
                }
                model_switch_timeout_secs = seconds;
            }
            "--log-dir" => {
                let Some(value) = args.next() else {
                    usage();
                };
                log_dir = PathBuf::from(value);
            }
            "-h" | "--help" => usage(),
            _ => scripts.push(PathBuf::from(arg)),
        }
    }

    if scripts.is_empty() {
        usage();
    }

    (
        host,
        listen_port,
        upstream_timeout_secs,
        model_ready_timeout_secs,
        model_switch_timeout_secs,
        log_dir,
        scripts,
    )
}

async fn health(State(state): State<AppState>) -> impl IntoResponse {
    let active_model = state.mux.active_model().await;
    debug!(active_model = ?active_model, "Health check served");
    Json(json!({
        "status": "ok",
        "active_model": active_model
    }))
}

async fn list_models(State(state): State<AppState>) -> impl IntoResponse {
    let data: Vec<Value> = state
        .model_ids
        .iter()
        .map(|model_id| {
            json!({
                "id": model_id,
                "object": "model",
                "created": state.model_created_unix,
                "owned_by": "sglangmux",
                "permission": [],
                "root": model_id,
                "parent": Value::Null
            })
        })
        .collect();

    debug!(model_count = data.len(), "Listing deployed models");
    Json(json!({
        "object": "list",
        "data": data
    }))
}

async fn proxy_generation_request(
    state: &AppState,
    body: Value,
    upstream_path: &str,
) -> axum::response::Response {
    let request_id = REQUEST_COUNTER.fetch_add(1, Ordering::Relaxed);

    let Some(model) = body.get("model").and_then(Value::as_str) else {
        warn!(
            request_id,
            upstream_path, "Rejecting request without required string field `model`"
        );
        return (
            StatusCode::BAD_REQUEST,
            Json(json!({"error": "missing string field: model"})),
        )
            .into_response();
    };

    let message_count = body
        .get("messages")
        .and_then(Value::as_array)
        .map(std::vec::Vec::len);
    let stream = body.get("stream").and_then(Value::as_bool).unwrap_or(false);
    info!(
        request_id,
        model = %model,
        upstream_path,
        message_count = ?message_count,
        stream,
        "Incoming generation request"
    );

    let mut attempt: u32 = 0;
    loop {
        attempt += 1;
        debug!(
            request_id,
            model = %model,
            attempt,
            "Ensuring requested model is ready"
        );
        if let Err(error) = state.mux.ensure_model_ready(model).await {
            error!(
                request_id,
                model = %model,
                attempt,
                error = %error,
                "Failed to ensure requested model is ready"
            );
            return (
                StatusCode::SERVICE_UNAVAILABLE,
                Json(json!({"error": format!("model not ready: {}", error)})),
            )
                .into_response();
        }

        let Some(_guard) = state.mux.acquire_in_flight(model) else {
            // Active model started draining between readiness check and guard acquire.
            // Re-enter readiness flow, same pattern used by llmux middleware.
            debug!(
                request_id,
                model = %model,
                attempt,
                "Model started draining between readiness and lock acquisition; retrying"
            );
            continue;
        };
        debug!(
            request_id,
            model = %model,
            attempt,
            "Acquired in-flight lock for request"
        );

        let Some(port) = state.model_ports.get(model).copied() else {
            error!(
                request_id,
                model = %model,
                "Model was not found in daemon port map"
            );
            return (
                StatusCode::NOT_FOUND,
                Json(json!({"error": format!("unknown model: {}", model)})),
            )
                .into_response();
        };

        let url = format!("http://127.0.0.1:{port}{upstream_path}");
        debug!(
            request_id,
            model = %model,
            attempt,
            port,
            url,
            "Forwarding request to upstream model server"
        );
        let upstream_start = Instant::now();
        let upstream = tokio::time::timeout(
            state.upstream_timeout,
            state.client.post(&url).json(&body).send(),
        )
        .await;

        return match upstream {
            Ok(Ok(response)) => {
                let status = response.status();
                let upstream_content_type = response.headers().get(CONTENT_TYPE).cloned();
                debug!(
                    request_id,
                    model = %model,
                    attempt,
                    status = %status,
                    latency_ms = upstream_start.elapsed().as_millis(),
                    "Upstream responded"
                );

                let body_bytes = match response.bytes().await {
                    Ok(bytes) => bytes,
                    Err(error) => {
                        error!(
                            request_id,
                            model = %model,
                            attempt,
                            error = %error,
                            "Failed reading upstream response body"
                        );
                        return (
                            StatusCode::BAD_GATEWAY,
                            Json(json!({"error": format!("failed reading upstream body: {}", error)})),
                        )
                            .into_response();
                    }
                };

                let content_type_str = upstream_content_type
                    .as_ref()
                    .and_then(|value| value.to_str().ok())
                    .unwrap_or("");
                let is_event_stream = content_type_str.starts_with("text/event-stream");

                if stream || is_event_stream {
                    info!(
                        request_id,
                        model = %model,
                        attempt,
                        status = %status,
                        content_type = content_type_str,
                        "Returning upstream streaming payload"
                    );
                    let mut out = axum::response::Response::new(Body::from(body_bytes));
                    *out.status_mut() = status;
                    if let Some(value) = upstream_content_type {
                        out.headers_mut().insert(CONTENT_TYPE, value);
                    }
                    return out;
                }

                match serde_json::from_slice::<Value>(&body_bytes) {
                    Ok(payload) => {
                        info!(
                            request_id,
                            model = %model,
                            attempt,
                            status = %status,
                            "Request completed successfully through mux"
                        );
                        (status, Json(payload)).into_response()
                    }
                    Err(error) => {
                        let preview = String::from_utf8_lossy(&body_bytes);
                        let preview = preview.chars().take(200).collect::<String>();
                        error!(
                            request_id,
                            model = %model,
                            attempt,
                            status = %status,
                            content_type = content_type_str,
                            error = %error,
                            preview = %preview,
                            "Upstream returned non-JSON payload"
                        );
                        (
                            StatusCode::BAD_GATEWAY,
                            Json(json!({"error": format!("invalid upstream response: {}", error)})),
                        )
                            .into_response()
                    }
                }
            }
            Ok(Err(error)) => {
                error!(
                    request_id,
                    model = %model,
                    attempt,
                    error = %error,
                    latency_ms = upstream_start.elapsed().as_millis(),
                    "Upstream request failed"
                );
                (
                    StatusCode::BAD_GATEWAY,
                    Json(json!({"error": format!("upstream request failed: {}", error)})),
                )
                    .into_response()
            }
            Err(_) => {
                error!(
                    request_id,
                    model = %model,
                    attempt,
                    timeout_s = state.upstream_timeout.as_secs(),
                    "Upstream request timed out"
                );
                (
                    StatusCode::GATEWAY_TIMEOUT,
                    Json(json!({"error": "upstream request timed out"})),
                )
                    .into_response()
            }
        };
    }
}

async fn chat_completions(
    State(state): State<AppState>,
    Json(body): Json<Value>,
) -> impl IntoResponse {
    debug!("Received /v1/chat/completions request");
    proxy_generation_request(&state, body, "/v1/chat/completions").await
}

async fn completions(State(state): State<AppState>, Json(body): Json<Value>) -> impl IntoResponse {
    debug!("Received /v1/completions request");
    proxy_generation_request(&state, body, "/v1/completions").await
}

async fn shutdown_signal(mux: SgLangMux) {
    if tokio::signal::ctrl_c().await.is_ok() {
        info!("Received Ctrl+C, shutting down all model processes");
        mux.shutdown_all().await;
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    init_tracing();

    let (
        host,
        listen_port,
        upstream_timeout_secs,
        model_ready_timeout_secs,
        model_switch_timeout_secs,
        log_dir,
        scripts,
    ) = parse_args();
    let upstream_timeout = Duration::from_secs(upstream_timeout_secs);
    info!(
        host = %host,
        listen_port,
        upstream_timeout_secs,
        model_ready_timeout_secs,
        model_switch_timeout_secs,
        log_dir = %log_dir.display(),
        script_count = scripts.len(),
        "Starting sglangmuxd"
    );

    let options = SgLangMuxOptions {
        ready_timeout: Duration::from_secs(model_ready_timeout_secs),
        request_timeout: Duration::from_secs(model_switch_timeout_secs),
        log_dir,
        ..SgLangMuxOptions::default()
    };

    let mux = SgLangMux::from_scripts(scripts, options)?;
    debug!(models = ?mux.model_specs(), "Parsed model specs");
    mux.bootstrap_sequential().await?;
    info!(active_model = ?mux.active_model().await, "Bootstrap complete");

    let model_specs = mux.model_specs();
    let model_ports: HashMap<String, u16> = model_specs
        .iter()
        .map(|spec| (spec.model_name.clone(), spec.port))
        .collect();
    let model_ids: Vec<String> = model_specs
        .into_iter()
        .map(|spec| spec.model_name)
        .collect();
    let model_created_unix = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs();

    info!(
        model_ids = ?model_ids,
        model_created_unix,
        "Built model registry"
    );
    info!(model_ports = ?model_ports, "Built model port map");

    let state = AppState {
        mux: mux.clone(),
        model_ports,
        model_ids,
        model_created_unix,
        client: Client::new(),
        upstream_timeout,
    };

    let app = Router::new()
        .route("/health", get(health))
        .route("/models", get(list_models))
        .route("/v1/models", get(list_models))
        .route("/v1/chat/completions", post(chat_completions))
        .route("/v1/completions", post(completions))
        .with_state(state);

    let listener = tokio::net::TcpListener::bind((host.as_str(), listen_port)).await?;
    info!("sglangmuxd listening on http://{}:{}", host, listen_port);
    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal(mux))
        .await?;

    Ok(())
}
