//! Axum management and public gateway protocol adapters.
//!
//! The management router accepts JSON and multipart packages and exposes
//! lifecycle operations. The gateway router is intentionally smaller: it
//! accepts raw `application/octet-stream` input and returns raw bytes plus
//! `x-wasmeld-revision` and `x-wasmeld-latency-ms` response headers.
//!
//! Neither router installs authentication or authorization middleware. A
//! production deployment must place the appropriate identity, authorization,
//! TLS, rate-limit, and request-size policy in front of these routers. CORS is
//! a browser access policy and must not be treated as authentication.

use super::*;
use axum::{
    Json, Router,
    body::{Body, Bytes},
    extract::{DefaultBodyLimit, Multipart, Path as AxumPath, State, rejection::BytesRejection},
    http::{HeaderMap, HeaderValue, Method, StatusCode, header},
    response::{IntoResponse, Response},
    routing::{delete, get, post},
};
use serde_json::json;
use tower_http::{
    cors::CorsLayer,
    trace::{DefaultMakeSpan, DefaultOnResponse, TraceLayer},
};
use tracing::Level;
use wasmeld_package::wit_package::WitPackageError;

/// Builds the Axum management API around a shared [`Console`].
///
/// The returned router exposes Runtime control, component registration,
/// invocation, events, and WIT Registry endpoints.
///
/// `allowed_origins` only affects browser CORS responses. Passing an empty
/// vector does not make the API private; it merely omits allowed origins.
/// Multipart bodies receive one MiB of framing headroom above the larger
/// configured artifact limit.
pub fn app(console: Arc<Console>, allowed_origins: Vec<HeaderValue>) -> Router {
    let max_body_bytes = console
        .max_artifact_bytes()
        .max(console.max_wit_package_bytes())
        .saturating_add(1024 * 1024);
    let mut cors = CorsLayer::new()
        .allow_methods([Method::GET, Method::POST, Method::DELETE])
        .allow_headers([header::CONTENT_TYPE]);
    if !allowed_origins.is_empty() {
        cors = cors.allow_origin(allowed_origins);
    }

    Router::new()
        .route("/healthz", get(health))
        .route("/api/v1/runtime", get(runtime_status))
        .route("/api/v1/runtime/start", post(start_runtime))
        .route("/api/v1/runtime/stop", post(stop_runtime))
        .route("/api/v1/runtime/restart", post(restart_runtime))
        .route("/api/v1/services", get(list_services).post(register))
        .route(
            "/api/v1/services/{id}/{revision}",
            delete(unregister_service),
        )
        .route("/api/v1/deployments", get(list_deployments))
        .route("/api/v1/deployments/{id}", get(get_deployment))
        .route(
            "/api/v1/deployments/{id}/activate",
            post(activate_deployment),
        )
        .route(
            "/api/v1/wit/packages",
            get(list_wit_packages).post(publish_wit_package),
        )
        .route(
            "/api/v1/wit/packages/{namespace}/{name}/{version}",
            get(get_wit_package),
        )
        .route(
            "/api/v1/wit/packages/{namespace}/{name}/{version}/content",
            get(download_wit_package),
        )
        .route(
            "/api/v1/services/{id}/{revision}/start",
            post(start_service),
        )
        .route("/api/v1/services/{id}/{revision}/stop", post(stop_service))
        .route(
            "/api/v1/services/{id}/{revision}/restart",
            post(restart_service),
        )
        .route(
            "/api/v1/services/{id}/{revision}/invoke",
            post(invoke_service),
        )
        .route("/api/v1/events", get(list_events))
        .layer(DefaultBodyLimit::max(max_body_bytes))
        .layer(cors)
        .layer(
            TraceLayer::new_for_http()
                .make_span_with(DefaultMakeSpan::new().level(Level::INFO))
                .on_response(DefaultOnResponse::new().level(Level::INFO)),
        )
        .with_state(console)
        .merge(web::router())
}

/// Builds the public data-plane API around the same resident Runtime.
///
/// This router intentionally contains no registration, lifecycle, Registry,
/// or deployment management routes. Callers must send
/// `Content-Type: application/octet-stream`; the stable service ID is resolved
/// to its active immutable revision at the start of each request.
///
/// # Wire example
///
/// ```text
/// POST /v1/services/image-resize/invoke
/// Content-Type: application/octet-stream
///
/// <component-specific bytes>
/// ```
pub fn gateway_app(console: Arc<Console>) -> Router {
    let max_input_bytes = console.max_gateway_input_bytes();
    Router::new()
        .route("/healthz", get(health))
        .route("/v1/services/{id}/invoke", post(gateway_invoke))
        .layer(DefaultBodyLimit::max(max_input_bytes))
        .layer(
            TraceLayer::new_for_http()
                .make_span_with(DefaultMakeSpan::new().level(Level::INFO))
                .on_response(DefaultOnResponse::new().level(Level::INFO)),
        )
        .with_state(console)
}

async fn health() -> Json<serde_json::Value> {
    Json(json!({ "status": "ok" }))
}

async fn gateway_invoke(
    State(console): State<Arc<Console>>,
    AxumPath(service_id): AxumPath<String>,
    headers: HeaderMap,
    body: Result<Bytes, BytesRejection>,
) -> Result<Response, GatewayError> {
    validate_path_segment("id", &service_id)?;
    let content_type = headers
        .get(header::CONTENT_TYPE)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.split(';').next())
        .map(str::trim);
    if !content_type.is_some_and(|value| value.eq_ignore_ascii_case("application/octet-stream")) {
        return Err(GatewayError::UnsupportedMediaType);
    }
    let input = body.map_err(|_| GatewayError::PayloadTooLarge {
        limit: console.max_gateway_input_bytes(),
    })?;
    // Deployment resolution and Actor enqueue happen in the same blocking
    // closure. The returned revision header therefore identifies the exact
    // revision that received this request, even if deployment changes later.
    let invocation = run_blocking(console, move |console| {
        console.gateway_invoke(&service_id, input.to_vec())
    })
    .await?;

    let mut response_headers = HeaderMap::new();
    response_headers.insert(
        header::CONTENT_TYPE,
        HeaderValue::from_static("application/octet-stream"),
    );
    response_headers.insert(
        "x-wasmeld-revision",
        HeaderValue::from_str(&invocation.revision).map_err(|_| {
            ConsoleError::InvalidRequest("revision is not a valid header".to_owned())
        })?,
    );
    response_headers.insert(
        "x-wasmeld-latency-ms",
        HeaderValue::from_str(&invocation.result.latency_ms.to_string()).map_err(|_| {
            ConsoleError::InvalidRequest("latency is not a valid header".to_owned())
        })?,
    );
    Ok((response_headers, invocation.result.output).into_response())
}

async fn runtime_status(
    State(console): State<Arc<Console>>,
) -> Result<Json<RuntimeView>, ApiError> {
    Ok(Json(console.runtime_view()?))
}

async fn start_runtime(State(console): State<Arc<Console>>) -> Result<Json<RuntimeView>, ApiError> {
    Ok(Json(
        run_blocking_and_persist(console, |console| console.start_runtime()).await?,
    ))
}

async fn stop_runtime(State(console): State<Arc<Console>>) -> Result<Json<RuntimeView>, ApiError> {
    Ok(Json(
        run_blocking_and_persist(console, |console| console.stop_runtime()).await?,
    ))
}

async fn restart_runtime(
    State(console): State<Arc<Console>>,
) -> Result<Json<RuntimeView>, ApiError> {
    Ok(Json(
        run_blocking_and_persist(console, |console| console.restart_runtime()).await?,
    ))
}

async fn list_services(State(console): State<Arc<Console>>) -> Result<Json<ServiceList>, ApiError> {
    Ok(Json(ServiceList {
        services: console.services()?,
    }))
}

async fn list_deployments(
    State(console): State<Arc<Console>>,
) -> Result<Json<DeploymentList>, ApiError> {
    Ok(Json(DeploymentList {
        deployments: console.deployments()?,
    }))
}

async fn get_deployment(
    State(console): State<Arc<Console>>,
    AxumPath(service_id): AxumPath<String>,
) -> Result<Json<DeploymentView>, ApiError> {
    validate_path_segment("id", &service_id)?;
    Ok(Json(console.deployment(&service_id)?))
}

async fn activate_deployment(
    State(console): State<Arc<Console>>,
    AxumPath(service_id): AxumPath<String>,
    Json(request): Json<ActivateDeploymentRequest>,
) -> Result<Json<DeploymentView>, ApiError> {
    validate_path_segment("id", &service_id)?;
    validate_path_segment("revision", &request.revision)?;
    let init_config = decode_optional_base64(request.init_config_base64.as_deref())?;
    Ok(Json(
        run_blocking_and_persist(console, move |console| {
            console.activate_deployment(service_id, request.revision, init_config)
        })
        .await?,
    ))
}

async fn list_events(State(console): State<Arc<Console>>) -> Result<Json<EventList>, ApiError> {
    Ok(Json(EventList {
        events: console.events()?,
    }))
}

async fn list_wit_packages(
    State(console): State<Arc<Console>>,
) -> Result<Json<WitPackageList>, ApiError> {
    Ok(Json(WitPackageList {
        packages: console.wit_registry.list()?,
    }))
}

async fn get_wit_package(
    State(console): State<Arc<Console>>,
    AxumPath((namespace, name, version)): AxumPath<(String, String, String)>,
) -> Result<Json<WitPackageMetadata>, ApiError> {
    Ok(Json(
        console
            .wit_registry
            .metadata(&format!("{namespace}:{name}"), &version)?,
    ))
}

async fn download_wit_package(
    State(console): State<Arc<Console>>,
    AxumPath((namespace, name, version)): AxumPath<(String, String, String)>,
) -> Result<Response, ApiError> {
    let (metadata, bytes) = console
        .wit_registry
        .download(&format!("{namespace}:{name}"), &version)?;
    Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "application/wasm")
        .header(header::ETAG, format!("\"{}\"", metadata.sha256))
        .body(Body::from(bytes))
        .map_err(|error| {
            ConsoleError::InvalidRequest(format!("failed to build package response: {error}"))
                .into()
        })
}

async fn publish_wit_package(
    State(console): State<Arc<Console>>,
    mut multipart: Multipart,
) -> Result<(StatusCode, Json<WitPackageMetadata>), ApiError> {
    // Axum materializes this one field in memory. DefaultBodyLimit above is
    // only a transport guard; WitRegistry repeats the exact artifact limit and
    // validates the package identity before committing it to disk.
    let mut package = None;
    while let Some(field) = multipart
        .next_field()
        .await
        .map_err(|error| ConsoleError::InvalidRequest(format!("invalid multipart body: {error}")))?
    {
        match field.name() {
            Some("package") => {
                if package.is_some() {
                    return Err(ConsoleError::InvalidRequest(
                        "package field must occur exactly once".to_owned(),
                    )
                    .into());
                }
                package = Some(
                    field
                        .bytes()
                        .await
                        .map_err(|error| {
                            ConsoleError::InvalidRequest(format!(
                                "invalid WIT package field: {error}"
                            ))
                        })?
                        .to_vec(),
                );
            }
            Some(name) => {
                return Err(ConsoleError::InvalidRequest(format!(
                    "unsupported multipart field {name}"
                ))
                .into());
            }
            None => {}
        }
    }
    let package = package
        .ok_or_else(|| ConsoleError::InvalidRequest("package field is required".to_owned()))?;
    if package.len() > console.max_wit_package_bytes() {
        return Err(WitRegistryError::TooLarge {
            limit: console.max_wit_package_bytes(),
        }
        .into());
    }
    let metadata = run_blocking(console, move |console| {
        Ok(console.wit_registry.publish(&package)?)
    })
    .await?;
    Ok((StatusCode::CREATED, Json(metadata)))
}

async fn register(
    State(console): State<Arc<Console>>,
    mut multipart: Multipart,
) -> Result<(StatusCode, Json<ServiceView>), ApiError> {
    let mut package = None;

    while let Some(field) = multipart
        .next_field()
        .await
        .map_err(|error| ConsoleError::InvalidRequest(format!("invalid multipart body: {error}")))?
    {
        match field.name() {
            Some("package") => {
                if package.is_some() {
                    return Err(ConsoleError::InvalidRequest(
                        "package field must occur exactly once".to_owned(),
                    )
                    .into());
                }
                package = Some(
                    field
                        .bytes()
                        .await
                        .map_err(|error| {
                            ConsoleError::InvalidRequest(format!("invalid package field: {error}"))
                        })?
                        .to_vec(),
                );
            }
            Some(name) => {
                return Err(ConsoleError::InvalidRequest(format!(
                    "unsupported multipart field {name}"
                ))
                .into());
            }
            None => {}
        }
    }

    let package = package
        .ok_or_else(|| ConsoleError::InvalidRequest("package field is required".to_owned()))?;
    let service =
        run_blocking_and_persist(console, move |console| console.register_package(package)).await?;
    Ok((StatusCode::CREATED, Json(service)))
}

async fn start_service(
    State(console): State<Arc<Console>>,
    AxumPath((id, revision)): AxumPath<(String, String)>,
    request: Option<Json<StartRequest>>,
) -> Result<Json<ServiceView>, ApiError> {
    let key = api_key(id, revision)?;
    let init_config = decode_optional_base64(
        request
            .as_ref()
            .and_then(|Json(request)| request.init_config_base64.as_deref()),
    )?;
    Ok(Json(
        run_blocking_and_persist(console, move |console| console.start(&key, init_config)).await?,
    ))
}

async fn stop_service(
    State(console): State<Arc<Console>>,
    AxumPath((id, revision)): AxumPath<(String, String)>,
) -> Result<Json<ServiceView>, ApiError> {
    let key = api_key(id, revision)?;
    Ok(Json(
        run_blocking_and_persist(console, move |console| console.stop(&key)).await?,
    ))
}

async fn unregister_service(
    State(console): State<Arc<Console>>,
    AxumPath((id, revision)): AxumPath<(String, String)>,
) -> Result<Json<ServiceView>, ApiError> {
    let key = api_key(id, revision)?;
    Ok(Json(
        run_blocking_and_persist(console, move |console| console.unregister(&key)).await?,
    ))
}

async fn restart_service(
    State(console): State<Arc<Console>>,
    AxumPath((id, revision)): AxumPath<(String, String)>,
    request: Option<Json<StartRequest>>,
) -> Result<Json<ServiceView>, ApiError> {
    let key = api_key(id, revision)?;
    let init_config = decode_optional_base64(
        request
            .as_ref()
            .and_then(|Json(request)| request.init_config_base64.as_deref()),
    )?;
    Ok(Json(
        run_blocking_and_persist(console, move |console| console.restart(&key, init_config))
            .await?,
    ))
}

async fn invoke_service(
    State(console): State<Arc<Console>>,
    AxumPath((id, revision)): AxumPath<(String, String)>,
    Json(request): Json<InvokeRequest>,
) -> Result<Json<InvokeResponse>, ApiError> {
    let key = api_key(id, revision)?;
    let input = BASE64.decode(request.input_base64).map_err(|error| {
        ConsoleError::InvalidRequest(format!("input_base64 is invalid: {error}"))
    })?;
    let result = run_blocking(console, move |console| console.invoke(&key, input)).await?;
    Ok(Json(InvokeResponse {
        output_base64: BASE64.encode(&result.output),
        output_bytes: result.output.len(),
        latency_ms: result.latency_ms,
    }))
}

async fn run_blocking<T, F>(console: Arc<Console>, operation: F) -> Result<T, ConsoleError>
where
    T: Send + 'static,
    F: FnOnce(Arc<Console>) -> Result<T, ConsoleError> + Send + 'static,
{
    // Wasmtime compilation and synchronous Actor calls must not occupy a Tokio
    // async worker.
    tokio::task::spawn_blocking(move || operation(console))
        .await
        .map_err(ConsoleError::from)
        .and_then(|result| result)
}

async fn run_blocking_and_persist<T, F>(
    console: Arc<Console>,
    operation: F,
) -> Result<T, ConsoleError>
where
    T: Send + 'static,
    F: FnOnce(Arc<Console>) -> Result<T, ConsoleError> + Send + 'static,
{
    // Control operations mutate durable state. Persist even on failure because
    // an operation may have changed state before returning its error.
    let operation_console = Arc::clone(&console);
    let result = run_blocking(operation_console, operation).await;
    // The flush marker orders previously accepted metric deltas before the
    // full snapshot transaction. It means "attempted", not "durably written";
    // the telemetry worker logs write failures independently.
    console.flush_invocation_telemetry().await;
    let persisted = console.persist_snapshot().await;

    match result {
        Ok(value) => {
            persisted?;
            Ok(value)
        }
        Err(error) => {
            if let Err(persistence_error) = persisted {
                tracing::error!(
                    error = %persistence_error,
                    "failed to persist console state after operation error"
                );
            }
            Err(error)
        }
    }
}
enum GatewayError {
    Console(ConsoleError),
    UnsupportedMediaType,
    PayloadTooLarge { limit: usize },
}

impl From<ConsoleError> for GatewayError {
    fn from(error: ConsoleError) -> Self {
        Self::Console(error)
    }
}

impl IntoResponse for GatewayError {
    fn into_response(self) -> Response {
        let (status, code, message) = match self {
            Self::UnsupportedMediaType => (
                StatusCode::UNSUPPORTED_MEDIA_TYPE,
                "unsupported_media_type",
                "Content-Type must be application/octet-stream".to_owned(),
            ),
            Self::PayloadTooLarge { limit } => (
                StatusCode::PAYLOAD_TOO_LARGE,
                "payload_too_large",
                format!("request body exceeds the {limit}-byte limit"),
            ),
            Self::Console(error) => {
                let (status, code, message) = match &error {
                    ConsoleError::InvalidRequest(_) => (
                        StatusCode::BAD_REQUEST,
                        "invalid_request",
                        error.to_string(),
                    ),
                    ConsoleError::DeploymentNotFound(_)
                    | ConsoleError::ServiceNotFound(_)
                    | ConsoleError::Runtime(RuntimeError::ServiceNotRegistered(_)) => (
                        StatusCode::NOT_FOUND,
                        "service_not_found",
                        "service is not deployed".to_owned(),
                    ),
                    ConsoleError::RuntimeNotRunning
                    | ConsoleError::Runtime(RuntimeError::ActorUnavailable(_))
                    | ConsoleError::Runtime(RuntimeError::ActorStopped(_)) => (
                        StatusCode::SERVICE_UNAVAILABLE,
                        "service_unavailable",
                        "service is temporarily unavailable".to_owned(),
                    ),
                    ConsoleError::Runtime(RuntimeError::ActorOverloaded(_)) => (
                        StatusCode::TOO_MANY_REQUESTS,
                        "service_overloaded",
                        "service is overloaded".to_owned(),
                    ),
                    ConsoleError::Runtime(RuntimeError::InputTooLarge { .. }) => (
                        StatusCode::PAYLOAD_TOO_LARGE,
                        "payload_too_large",
                        error.to_string(),
                    ),
                    ConsoleError::Runtime(RuntimeError::DeadlineExceeded { .. }) => (
                        StatusCode::GATEWAY_TIMEOUT,
                        "deadline_exceeded",
                        "service execution deadline exceeded".to_owned(),
                    ),
                    ConsoleError::Runtime(
                        RuntimeError::OutputTooLarge { .. }
                        | RuntimeError::ActorFault { .. }
                        | RuntimeError::ComponentError { .. },
                    ) => (StatusCode::BAD_GATEWAY, "service_error", error.to_string()),
                    ConsoleError::ArtifactTooLarge { .. }
                    | ConsoleError::ActiveDeployment(_)
                    | ConsoleError::Package(_)
                    | ConsoleError::WitRegistry(_)
                    | ConsoleError::Storage { .. }
                    | ConsoleError::ManifestSerialize(_)
                    | ConsoleError::Database(_)
                    | ConsoleError::KvBackendStart(_)
                    | ConsoleError::InvalidDatabaseData(_)
                    | ConsoleError::Runtime(_)
                    | ConsoleError::Join(_)
                    | ConsoleError::LockPoisoned => (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        "internal_error",
                        "internal gateway error".to_owned(),
                    ),
                };
                (status, code, message)
            }
        };

        (
            status,
            Json(json!({
                "error": {
                    "code": code,
                    "message": message,
                }
            })),
        )
            .into_response()
    }
}

struct ApiError(ConsoleError);

impl From<ConsoleError> for ApiError {
    fn from(error: ConsoleError) -> Self {
        Self(error)
    }
}

impl From<WitRegistryError> for ApiError {
    fn from(error: WitRegistryError) -> Self {
        Self(error.into())
    }
}

impl IntoResponse for ApiError {
    fn into_response(self) -> Response {
        let (status, code) = match &self.0 {
            ConsoleError::InvalidRequest(_) => (StatusCode::BAD_REQUEST, "invalid_request"),
            ConsoleError::ServiceNotFound(_) => (StatusCode::NOT_FOUND, "service_not_found"),
            ConsoleError::DeploymentNotFound(_) => (StatusCode::NOT_FOUND, "deployment_not_found"),
            ConsoleError::ActiveDeployment(_) => (StatusCode::CONFLICT, "active_deployment"),
            ConsoleError::RuntimeNotRunning => (StatusCode::CONFLICT, "runtime_not_running"),
            ConsoleError::ArtifactTooLarge { .. }
            | ConsoleError::Package(PackageError::ComponentTooLarge { .. })
            | ConsoleError::WitRegistry(WitRegistryError::TooLarge { .. })
            | ConsoleError::Runtime(RuntimeError::InputTooLarge { .. })
            | ConsoleError::Runtime(RuntimeError::OutputTooLarge { .. }) => {
                (StatusCode::PAYLOAD_TOO_LARGE, "payload_too_large")
            }
            ConsoleError::Runtime(RuntimeError::ServiceNotRegistered(_)) => {
                (StatusCode::NOT_FOUND, "service_not_registered")
            }
            ConsoleError::Runtime(RuntimeError::ServiceAlreadyRegistered(_))
            | ConsoleError::Runtime(RuntimeError::ActorUnavailable(_))
            | ConsoleError::Runtime(RuntimeError::ActorOverloaded(_)) => {
                (StatusCode::CONFLICT, "runtime_conflict")
            }
            ConsoleError::WitRegistry(WitRegistryError::AlreadyExists(_)) => {
                (StatusCode::CONFLICT, "wit_package_exists")
            }
            ConsoleError::WitRegistry(WitRegistryError::NotFound(_)) => {
                (StatusCode::NOT_FOUND, "wit_package_not_found")
            }
            ConsoleError::Runtime(RuntimeError::DeadlineExceeded { .. }) => {
                (StatusCode::REQUEST_TIMEOUT, "deadline_exceeded")
            }
            ConsoleError::Runtime(
                RuntimeError::InvalidServiceKey
                | RuntimeError::InvalidManifest(_)
                | RuntimeError::ManifestParse(_)
                | RuntimeError::ComponentCompilation(_)
                | RuntimeError::UnsupportedImport(_)
                | RuntimeError::ActorInitialization(_)
                | RuntimeError::ComponentError { .. },
            ) => (StatusCode::UNPROCESSABLE_ENTITY, "invalid_component"),
            ConsoleError::Package(_) => (StatusCode::UNPROCESSABLE_ENTITY, "invalid_package"),
            ConsoleError::WitRegistry(WitRegistryError::InvalidPackage(
                WitPackageError::Source { .. }
                | WitPackageError::Encode(_)
                | WitPackageError::Decode(_)
                | WitPackageError::MissingVersion(_),
            )) => (StatusCode::UNPROCESSABLE_ENTITY, "invalid_wit_package"),
            ConsoleError::Storage { .. }
            | ConsoleError::ManifestSerialize(_)
            | ConsoleError::Database(_)
            | ConsoleError::KvBackendStart(_)
            | ConsoleError::InvalidDatabaseData(_)
            | ConsoleError::Runtime(_)
            | ConsoleError::WitRegistry(WitRegistryError::InvalidStorage(_))
            | ConsoleError::WitRegistry(WitRegistryError::Storage { .. })
            | ConsoleError::WitRegistry(WitRegistryError::LockPoisoned)
            | ConsoleError::Join(_)
            | ConsoleError::LockPoisoned => (StatusCode::INTERNAL_SERVER_ERROR, "internal_error"),
        };

        (
            status,
            Json(json!({
                "error": {
                    "code": code,
                    "message": self.0.to_string(),
                }
            })),
        )
            .into_response()
    }
}
