prvng_platform/crates/extension-registry/src/handlers.rs

use std::sync::Arc;

use axum::{
    extract::{Path, State},
    http::StatusCode,
    response::IntoResponse,
    routing::{get, head, post, put},
    Json, Router,
};
use extension_registry::service::{
    ExtensionMetadata, ExtensionRegistry, ImageManifest, PushBlobRequest,
};
use serde_json::json;

/// Application state
#[derive(Clone)]
pub struct AppState {
    registry: Arc<ExtensionRegistry>,
}

impl AppState {
    pub fn new(registry: Arc<ExtensionRegistry>) -> Self {
        Self { registry }
    }
}

/// Error response
#[derive(Debug)]
pub struct RegistryError {
    status: StatusCode,
    message: String,
}

impl IntoResponse for RegistryError {
    fn into_response(self) -> axum::response::Response {
        let body = Json(json!({
            "errors": [{
                "code": "REGISTRY_ERROR",
                "message": self.message,
            }]
        }));

        (self.status, body).into_response()
    }
}

impl RegistryError {
    pub fn not_found(msg: impl Into<String>) -> Self {
        Self {
            status: StatusCode::NOT_FOUND,
            message: msg.into(),
        }
    }

    pub fn internal(msg: impl Into<String>) -> Self {
        Self {
            status: StatusCode::INTERNAL_SERVER_ERROR,
            message: msg.into(),
        }
    }

    #[allow(dead_code)]
    pub fn bad_request(msg: impl Into<String>) -> Self {
        Self {
            status: StatusCode::BAD_REQUEST,
            message: msg.into(),
        }
    }
}

/// Check if blob exists (HEAD /v2/<name>/blobs/<digest>)
async fn blob_exists(
    Path((_name, digest)): Path<(String, String)>,
    State(state): State<AppState>,
) -> Result<StatusCode, RegistryError> {
    state
        .registry
        .blob_exists(&digest)
        .await
        .map_err(|e: anyhow::Error| RegistryError::internal(e.to_string()))?;

    Ok(StatusCode::OK)
}

/// Pull blob (GET /v2/<name>/blobs/<digest>)
async fn pull_blob(
    Path((_name, digest)): Path<(String, String)>,
    State(state): State<AppState>,
) -> Result<Vec<u8>, RegistryError> {
    state
        .registry
        .pull_blob(&digest)
        .await
        .map_err(|e: anyhow::Error| {
            if e.to_string().contains("not found") {
                RegistryError::not_found(e.to_string())
            } else {
                RegistryError::internal(e.to_string())
            }
        })
}

/// Push blob (POST /v2/<name>/blobs/uploads/)
async fn push_blob(
    Path(_name): Path<String>,
    State(state): State<AppState>,
    Json(req): Json<PushBlobRequest>,
) -> Result<(StatusCode, Json<serde_json::Value>), RegistryError> {
    let response = state
        .registry
        .push_blob(req)
        .await
        .map_err(|e: anyhow::Error| RegistryError::internal(e.to_string()))?;

    Ok((
        StatusCode::CREATED,
        Json(json!({
            "digest": response.digest,
            "size": response.size,
        })),
    ))
}

/// Pull manifest (GET /v2/<name>/manifests/<reference>)
async fn pull_manifest(
    Path((name, reference)): Path<(String, String)>,
    State(state): State<AppState>,
) -> Result<(StatusCode, Json<serde_json::Value>), RegistryError> {
    let response = state
        .registry
        .get_manifest(&format!("{}:{}", name, reference))
        .await
        .map_err(|e: anyhow::Error| {
            if e.to_string().contains("not found") {
                RegistryError::not_found(e.to_string())
            } else {
                RegistryError::internal(e.to_string())
            }
        })?;

    Ok((
        StatusCode::OK,
        Json(json!({
            "manifest": response.manifest,
            "digest": response.digest,
            "content_type": response.content_type,
        })),
    ))
}

/// Push manifest (PUT /v2/<name>/manifests/<reference>)
async fn push_manifest(
    Path((name, reference)): Path<(String, String)>,
    State(state): State<AppState>,
    Json(manifest): Json<ImageManifest>,
) -> Result<(StatusCode, Json<serde_json::Value>), RegistryError> {
    let digest = state
        .registry
        .put_manifest(format!("{}:{}", name, reference), manifest)
        .await
        .map_err(|e: anyhow::Error| RegistryError::internal(e.to_string()))?;

    Ok((
        StatusCode::CREATED,
        Json(json!({
            "digest": digest,
        })),
    ))
}

/// List repositories catalog (GET /v2/_catalog)
async fn list_catalog(
    State(state): State<AppState>,
) -> Result<Json<serde_json::Value>, RegistryError> {
    let extensions = state
        .registry
        .list_extensions()
        .await
        .map_err(|e| RegistryError::internal(e.to_string()))?;

    let repositories: Vec<String> = extensions.iter().map(|e| e.name.clone()).collect();

    Ok(Json(json!({
        "repositories": repositories,
    })))
}

/// List extension tags (GET /v2/<name>/tags/list)
async fn list_tags(
    Path(name): Path<String>,
    State(state): State<AppState>,
) -> Result<Json<serde_json::Value>, RegistryError> {
    let extension = state
        .registry
        .get_extension(&name)
        .await
        .map_err(|e: anyhow::Error| {
            if e.to_string().contains("not found") {
                RegistryError::not_found(e.to_string())
            } else {
                RegistryError::internal(e.to_string())
            }
        })?;

    Ok(Json(json!({
        "name": name,
        "tags": [extension.version],
    })))
}

/// Register extension metadata (POST /extensions)
async fn register_extension(
    State(state): State<AppState>,
    Json(metadata): Json<ExtensionMetadata>,
) -> Result<(StatusCode, Json<serde_json::Value>), RegistryError> {
    state
        .registry
        .register_extension(metadata.clone())
        .await
        .map_err(|e: anyhow::Error| RegistryError::internal(e.to_string()))?;

    Ok((
        StatusCode::CREATED,
        Json(json!({
            "name": metadata.name,
            "version": metadata.version,
        })),
    ))
}

/// Get extension metadata (GET /extensions/:name)
async fn get_extension(
    Path(name): Path<String>,
    State(state): State<AppState>,
) -> Result<Json<ExtensionMetadata>, RegistryError> {
    state
        .registry
        .get_extension(&name)
        .await
        .map_err(|e: anyhow::Error| {
            if e.to_string().contains("not found") {
                RegistryError::not_found(e.to_string())
            } else {
                RegistryError::internal(e.to_string())
            }
        })
        .map(Json)
}

/// List all extensions (GET /extensions)
async fn list_extensions(
    State(state): State<AppState>,
) -> Result<Json<Vec<ExtensionMetadata>>, RegistryError> {
    state
        .registry
        .list_extensions()
        .await
        .map_err(|e: anyhow::Error| RegistryError::internal(e.to_string()))
        .map(Json)
}

/// Health check (GET /health)
async fn health(State(state): State<AppState>) -> Result<StatusCode, RegistryError> {
    state
        .registry
        .health_check()
        .await
        .map_err(|e: anyhow::Error| RegistryError::internal(e.to_string()))?;

    Ok(StatusCode::OK)
}

/// Build API router
pub fn routes(state: AppState) -> Router {
    Router::new()
        // OCI v2 API
        .route("/v2/:name/blobs/:digest", head(blob_exists))
        .route("/v2/:name/blobs/:digest", get(pull_blob))
        .route("/v2/:name/blobs/uploads", post(push_blob))
        .route("/v2/:name/manifests/:reference", get(pull_manifest))
        .route("/v2/:name/manifests/:reference", put(push_manifest))
        .route("/v2/_catalog", get(list_catalog))
        .route("/v2/:name/tags/list", get(list_tags))
        // Extensions API
        .route("/extensions", post(register_extension))
        .route("/extensions", get(list_extensions))
        .route("/extensions/:name", get(get_extension))
        // Health
        .route("/health", get(health))
        .with_state(state)
}