use crate::backend_rust::RustBackend;
use crate::backend_swift::SwiftBackend;
use crate::config::Config;
use crate::dto::{EmbedRequest, EmbedResponse};
use crate::errors::OrchestratorError;
use crate::splitter;
use crate::tokenizer::TokenCounter;
use axum::{extract::State, http::StatusCode, routing::{get, post}, Json, Router};
use std::sync::Arc;

pub struct AppState {
    pub cfg: Arc<Config>,
    pub tokenizer: TokenCounter,
    pub swift: SwiftBackend,
    pub rust: RustBackend,
}

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

async fn embed_handler(
    State(state): State<Arc<AppState>>,
    Json(req): Json<EmbedRequest>,
) -> Result<Json<EmbedResponse>, OrchestratorError> {
    let resp = splitter::dispatch(
        req,
        &state.tokenizer,
        &state.swift,
        &state.rust,
        &state.cfg,
    )
    .await?;
    Ok(Json(resp))
}

async fn healthz(State(state): State<Arc<AppState>>) -> (StatusCode, &'static str) {
    if state.swift.healthz().await {
        (StatusCode::OK, "ok")
    } else {
        (StatusCode::SERVICE_UNAVAILABLE, "warming")
    }
}