pub mod audio;
pub mod batches;
pub mod chat;
pub mod embeddings;
pub mod files;
pub mod health;
pub mod images;
pub mod models;
pub mod moderations;
pub mod ocr;
pub mod realtime;
pub mod rerank;
pub mod responses;
pub mod search;

use axum::Router;
use axum::extract::DefaultBodyLimit;
use axum::http::HeaderValue;
use axum::http::header::AUTHORIZATION;
use axum::middleware;
use axum::routing::{get, post};
use tower_http::catch_panic::CatchPanicLayer;
use tower_http::compression::CompressionLayer;
use tower_http::cors::{AllowOrigin, Any, CorsLayer};
use tower_http::sensitive_headers::SetSensitiveHeadersLayer;
use tower_http::trace::TraceLayer;

use tower::Service;

use liter_llm::tower::types::{LlmRequest, LlmResponse};

use crate::auth;
use crate::auth::KeyContext;
use crate::error::ProxyError;
use crate::state::AppState;

/// Check model access for the authenticated key, attach the tenant identifier,
/// and dispatch the request through the Tower service stack.
///
/// `tenant_id` is propagated via [`LlmRequest::with_tenant_id`] so that every
/// Tower layer downstream — [`BudgetLedger::Tenant`], [`TenantScopedStrategy`],
/// and [`UsageEvent`] — receives the correct tenant dimension without each
/// handler needing to wire it independently.
pub(crate) async fn dispatch(
    state: &AppState,
    key_ctx: &KeyContext,
    model: &str,
    request: LlmRequest,
) -> Result<LlmResponse, ProxyError> {
    if !key_ctx.can_access_model(model) {
        return Err(ProxyError::forbidden(format!(
            "key '{}' is not allowed to access model '{model}'",
            key_ctx.key_id
        )));
    }
    let request = request.with_tenant_id(key_ctx.tenant_id.clone());
    let mut svc = state.service_pool.get_service(model)?;
    Ok(svc.call(request).await?)
}

/// Build the full axum router with all routes, middleware, and shared state.
pub fn build_router(state: AppState) -> Router {
    let v1_routes = Router::new()
        // Chat & Completions
        .route("/v1/chat/completions", post(chat::chat_completions))
        // Embeddings
        .route("/v1/embeddings", post(embeddings::create_embedding))
        // Models
        .route("/v1/models", get(models::list_models))
        // Images
        .route("/v1/images/generations", post(images::create_image))
        // Audio
        .route("/v1/audio/speech", post(audio::create_speech))
        .route(
            "/v1/audio/transcriptions",
            post(audio::create_transcription),
        )
        // Moderations
        .route("/v1/moderations", post(moderations::create_moderation))
        // Extended endpoints
        .route("/v1/rerank", post(rerank::rerank))
        .route("/v1/search", post(search::search))
        .route("/v1/ocr", post(ocr::ocr))
        // Files
        .route("/v1/files", post(files::create_file).get(files::list_files))
        .route(
            "/v1/files/{file_id}",
            get(files::retrieve_file).delete(files::delete_file),
        )
        .route("/v1/files/{file_id}/content", get(files::file_content))
        // Batches
        .route(
            "/v1/batches",
            post(batches::create_batch).get(batches::list_batches),
        )
        .route("/v1/batches/{batch_id}", get(batches::retrieve_batch))
        .route(
            "/v1/batches/{batch_id}/cancel",
            post(batches::cancel_batch),
        )
        // Responses
        .route("/v1/responses", post(responses::create_response))
        .route(
            "/v1/responses/{response_id}",
            get(responses::retrieve_response),
        )
        .route(
            "/v1/responses/{response_id}/cancel",
            post(responses::cancel_response),
        )
        // Realtime WebSocket proxy
        .route("/v1/realtime", get(realtime::realtime_websocket))
        // Auth middleware on all /v1 routes
        .layer(middleware::from_fn_with_state(
            state.clone(),
            auth::validate_api_key,
        ));

    let health_routes = Router::new()
        // Legacy health endpoints (retained for backward compatibility).
        .route("/health", get(health::health))
        .route("/health/liveness", get(health::liveness))
        .route("/health/readiness", get(health::readiness))
        // v1.6 enriched endpoints.
        .route("/healthz", get(health::healthz))
        .route("/readyz", get(health::readyz))
        .route("/openapi.json", get(crate::openapi::openapi_schema));

    // Snapshot the current config for router-build decisions (CORS origins,
    // body limit). These are applied once at startup; a hot-reload of CORS
    // origins or body_limit_bytes requires a server restart. Request handlers
    // that need per-request config should call `state.config.load()`.
    let cfg_snapshot = state.config.load();

    // Build an optional CORS layer.  Empty cors_origins means CORS is disabled
    // entirely — no CorsLayer is added to the router.  A wildcard origin ("*")
    // is allowed but must NOT expose the Authorization header, which would
    // permit credentialed cross-origin requests from any origin.
    let cors_layer: Option<CorsLayer> = if cfg_snapshot.server.cors_origins.is_empty() {
        None
    } else if cfg_snapshot.server.cors_origins.iter().any(|o| o == "*") {
        Some(
            CorsLayer::new()
                .allow_origin(Any)
                .allow_methods(Any)
                // Deliberately exclude Authorization — wildcard origins must not
                // receive credentialed cross-origin access.
                .allow_headers([
                    axum::http::header::CONTENT_TYPE,
                    axum::http::header::ACCEPT,
                ]),
        )
    } else {
        let origins: Vec<HeaderValue> = cfg_snapshot
            .server
            .cors_origins
            .iter()
            .filter_map(|o| o.parse().ok())
            .collect();
        Some(
            CorsLayer::new()
                .allow_origin(AllowOrigin::list(origins))
                .allow_methods(Any)
                .allow_headers(Any),
        )
    };

    let mut router = Router::new()
        .merge(v1_routes)
        .merge(health_routes)
        .layer(SetSensitiveHeadersLayer::new([AUTHORIZATION]))
        .layer(DefaultBodyLimit::max(cfg_snapshot.server.body_limit_bytes))
        .layer(CompressionLayer::new())
        .layer(CatchPanicLayer::new())
        .layer(TraceLayer::new_for_http())
        .with_state(state);

    if let Some(layer) = cors_layer {
        router = router.layer(layer);
    }

    router
}
