use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

use axum::extract::State;
use axum::http::StatusCode;
use axum::response::IntoResponse;
use axum::routing::{get, post};
use axum::{Json, Router};

use crate::dto::{EmbedRequest, EmbedResponse, ModelSize};
use crate::worker::WorkRequest;

// ---------------------------------------------------------------------------
// Shared application state
// ---------------------------------------------------------------------------

pub struct ModelWorker {
    pub tx: std::sync::mpsc::SyncSender<WorkRequest>,
    pub n_embd: usize,
}

pub struct AppState {
    /// One worker channel per loaded model size.
    pub workers: HashMap<ModelSize, ModelWorker>,
    /// Set to true once all workers are ready.
    pub ready: AtomicBool,
    /// The default model size when the request omits `size`.
    pub default_size: ModelSize,
}

// ---------------------------------------------------------------------------
// Router
// ---------------------------------------------------------------------------

pub fn router(state: Arc<AppState>) -> Router {
    Router::new()
        .route("/healthz", get(healthz))
        .route("/embed", post(embed))
        .with_state(state)
}

// ---------------------------------------------------------------------------
// GET /healthz
// ---------------------------------------------------------------------------

async fn healthz(State(state): State<Arc<AppState>>) -> impl IntoResponse {
    if state.ready.load(Ordering::Relaxed) {
        (StatusCode::OK, "ok")
    } else {
        (StatusCode::SERVICE_UNAVAILABLE, "warming")
    }
}

// ---------------------------------------------------------------------------
// POST /embed
// ---------------------------------------------------------------------------

async fn embed(
    State(state): State<Arc<AppState>>,
    Json(req): Json<EmbedRequest>,
) -> Result<Json<EmbedResponse>, (StatusCode, String)> {
    // Validate texts non-empty
    if req.texts.is_empty() {
        return Err((
            StatusCode::BAD_REQUEST,
            "texts must be non-empty".to_string(),
        ));
    }

    // Resolve which model to use
    let size = req.size.unwrap_or(state.default_size);

    let worker = state.workers.get(&size).ok_or_else(|| {
        let available: Vec<String> = state.workers.keys().map(|s| s.to_string()).collect();
        (
            StatusCode::BAD_REQUEST,
            format!(
                "model size {size} not loaded; available: [{}]",
                available.join(", ")
            ),
        )
    })?;

    // Create oneshot for the response
    let (reply_tx, reply_rx) = tokio::sync::oneshot::channel();

    let work = WorkRequest {
        texts: req.texts,
        input_kind: req.input_type,
        reply: reply_tx,
    };

    // Send to worker
    worker.tx.send(work).map_err(|_| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            "worker channel closed".to_string(),
        )
    })?;

    // Await response
    let vectors = reply_rx
        .await
        .map_err(|_| {
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                "worker dropped reply".to_string(),
            )
        })?
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("{e:#}")))?;

    let dim = worker.n_embd;
    Ok(Json(EmbedResponse {
        vectors,
        dim,
        model: size.label().to_string(),
    }))
}