nidus-openapi 1.0.4

use std::collections::BTreeMap;

use axum::{Json, Router, response::Html, routing::get};
use nidus_http::error::RoutePathError;
use nidus_http::router::{OpenApiSchemaRegistrar, RouteMetadata};
use serde_json::{Value, json};
use utoipa::{PartialSchema, ToSchema};

use crate::html::docs_html;
use crate::route::OpenApiRoute;

/// Errors emitted while building an OpenAPI document.
#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
pub enum OpenApiDocumentError {
    /// Two routes attempted to register the same path and HTTP method.
    #[error("duplicate OpenAPI operation `{method}` `{path}`")]
    DuplicateOperation {
        /// Lowercase OpenAPI operation method.
        method: String,
        /// Normalized OpenAPI path.
        path: String,
    },
    /// Route path normalization failed.
    #[error(transparent)]
    RoutePath(#[from] RoutePathError),
    /// Schema registration failed.
    #[error("OpenAPI schema registration failed: {message}")]
    SchemaRegistration {
        /// Safe diagnostic from the schema registration failure.
        message: String,
    },
}

/// Minimal OpenAPI document metadata builder.
#[derive(Clone, Debug)]
pub struct OpenApiDocument {
    title: String,
    version: String,
    routes: Vec<OpenApiRoute>,
    schemas: BTreeMap<String, Value>,
}

impl OpenApiDocument {
    /// Creates an OpenAPI document.
    pub fn new(title: impl Into<String>, version: impl Into<String>) -> Self {
        Self {
            title: title.into(),
            version: version.into(),
            routes: Vec::new(),
            schemas: BTreeMap::new(),
        }
    }

    /// Adds route metadata to the document.
    pub fn route(mut self, route: OpenApiRoute) -> Self {
        self = self
            .try_route(route)
            .unwrap_or_else(|error| panic!("{error}"));
        self
    }

    /// Tries to add route metadata to the document.
    pub fn try_route(mut self, route: OpenApiRoute) -> Result<Self, OpenApiDocumentError> {
        if self
            .routes
            .iter()
            .any(|existing| existing.path() == route.path() && existing.method() == route.method())
        {
            return Err(OpenApiDocumentError::DuplicateOperation {
                method: route.method().to_owned(),
                path: route.path().to_owned(),
            });
        }
        self.routes.push(route);
        Ok(self)
    }

    /// Adds generated route metadata under a controller prefix.
    pub fn controller_routes(self, controller_prefix: &str, routes: &[RouteMetadata]) -> Self {
        self.try_controller_routes(controller_prefix, routes)
            .unwrap_or_else(|error| panic!("{error}"))
    }

    /// Tries to add generated route metadata under a controller prefix.
    pub fn try_controller_routes(
        mut self,
        controller_prefix: &str,
        routes: &[RouteMetadata],
    ) -> Result<Self, OpenApiDocumentError> {
        for route in routes {
            self = self.try_route(OpenApiRoute::try_from_route_metadata_at_path(
                route,
                route.try_full_path(controller_prefix)?,
            )?)?;
        }
        Ok(self)
    }

    /// Registers schemas from route OpenAPI metadata callbacks.
    pub fn schemas_from_route_metadata(mut self, routes: &[RouteMetadata]) -> Self {
        self = self
            .try_schemas_from_route_metadata(routes)
            .unwrap_or_else(|error| panic!("{error}"));
        self
    }

    /// Tries to register schemas from route OpenAPI metadata callbacks.
    pub fn try_schemas_from_route_metadata(
        mut self,
        routes: &[RouteMetadata],
    ) -> Result<Self, OpenApiDocumentError> {
        for route in routes {
            self = self
                .try_with_schema_registrar(route.request_schema_registrar())?
                .try_with_schema_registrar(route.response_schema_registrar())?;
        }
        Ok(self)
    }

    /// Adds a DTO schema generated by `utoipa::ToSchema`.
    pub fn schema<T>(mut self) -> Self
    where
        T: ToSchema,
    {
        self = self
            .try_schema::<T>()
            .unwrap_or_else(|error| panic!("{error}"));
        self
    }

    /// Tries to add a DTO schema generated by `utoipa::ToSchema`.
    pub fn try_schema<T>(mut self) -> Result<Self, OpenApiDocumentError>
    where
        T: ToSchema,
    {
        self.schemas = self.register_schemas(Self::try_collect_openapi_schemas::<T>()?);
        Ok(self)
    }

    fn try_collect_openapi_schemas<T: ToSchema>()
    -> Result<Vec<(String, Value)>, OpenApiDocumentError> {
        let mut openapi_schemas: Vec<(
            String,
            utoipa::openapi::RefOr<utoipa::openapi::schema::Schema>,
        )> = vec![(T::name().to_string(), <T as PartialSchema>::schema())];
        <T as ToSchema>::schemas(&mut openapi_schemas);
        let mut schemas = Vec::new();
        for (name, schema) in openapi_schemas {
            let value = serde_json::to_value(schema).map_err(|error| {
                OpenApiDocumentError::SchemaRegistration {
                    message: format!("schema `{name}`: {error}"),
                }
            })?;
            schemas.push((name, value));
        }
        Ok(schemas)
    }

    /// Creates an OpenAPI document from generated route metadata.
    pub fn from_route_metadata(
        title: impl Into<String>,
        version: impl Into<String>,
        routes: &[RouteMetadata],
    ) -> Self {
        Self::try_from_route_metadata(title, version, routes)
            .unwrap_or_else(|error| panic!("{error}"))
    }

    /// Tries to create an OpenAPI document from generated route metadata.
    pub fn try_from_route_metadata(
        title: impl Into<String>,
        version: impl Into<String>,
        routes: &[RouteMetadata],
    ) -> Result<Self, OpenApiDocumentError> {
        let mut document = Self::new(title, version);
        for route in routes {
            document = document.try_route(OpenApiRoute::try_from_route_metadata(route)?)?;
        }
        Ok(document)
    }

    /// Creates an OpenAPI document from a controller prefix and generated route metadata.
    pub fn from_controller_routes(
        title: impl Into<String>,
        version: impl Into<String>,
        controller_prefix: &str,
        routes: &[RouteMetadata],
    ) -> Self {
        Self::try_from_controller_routes(title, version, controller_prefix, routes)
            .unwrap_or_else(|error| panic!("{error}"))
    }

    /// Tries to create an OpenAPI document from controller route metadata.
    pub fn try_from_controller_routes(
        title: impl Into<String>,
        version: impl Into<String>,
        controller_prefix: &str,
        routes: &[RouteMetadata],
    ) -> Result<Self, OpenApiDocumentError> {
        Self::new(title, version).try_controller_routes(controller_prefix, routes)
    }

    /// Renders the document as JSON.
    pub fn to_json_value(&self) -> Value {
        let mut paths = serde_json::Map::new();
        for route in &self.routes {
            let entry = paths
                .entry(route.path().to_owned())
                .or_insert_with(|| Value::Object(serde_json::Map::new()));
            if let Value::Object(methods) = entry {
                methods.insert(route.method().to_owned(), route.to_json_value());
            }
        }

        let mut document = json!({
            "openapi": "3.1.0",
            "info": {
                "title": self.title,
                "version": self.version,
            },
            "paths": paths,
        });

        if !self.schemas.is_empty() {
            document["components"] = json!({
                "schemas": &self.schemas,
            });
        }

        document
    }

    /// Builds an Axum router serving OpenAPI JSON and Swagger UI docs.
    pub fn into_router(self) -> Router {
        let json = self.to_json_value();
        let docs = docs_html(&self.title, "/openapi.json");

        Router::new()
            .route(
                "/openapi.json",
                get(move || {
                    let json = json.clone();
                    async move { Json(json) }
                }),
            )
            .route(
                "/docs",
                get(move || {
                    let docs = docs.clone();
                    async move { Html(docs) }
                }),
            )
    }

    fn try_with_schema_registrar(
        mut self,
        registrar: Option<OpenApiSchemaRegistrar>,
    ) -> Result<Self, OpenApiDocumentError> {
        let Some(registrar) = registrar else {
            return Ok(self);
        };
        let mut schemas = Vec::new();
        registrar(&mut schemas).map_err(|error| OpenApiDocumentError::SchemaRegistration {
            message: error.to_string(),
        })?;
        self.schemas = self.register_schemas(schemas);
        Ok(self)
    }

    fn register_schemas(&self, schemas: Vec<(String, Value)>) -> BTreeMap<String, Value> {
        let mut registered = self.schemas.clone();
        for (name, schema) in schemas {
            registered.entry(name).or_insert(schema);
        }
        registered
    }
}