utoipa-ts 0.1.2

use std::{
    any::TypeId,
    collections::{BTreeMap, HashSet},
    fmt::Write as _,
    path::{Path, PathBuf},
};

pub use ts_rs;
pub use utoipa_ts_macros::path;

pub mod __private {
    pub use inventory;
}

const NOTE: &str = "// This file was generated by utoipa-ts. Do not edit this file manually.\n";
pub const EXPORT_PATH_ENV: &str = "UTOIPA_TS_PATH";
pub const DEFAULT_EXPORT_PATH: &str = "api.ts";
const DEFAULT_EXPORT_FILE_NAME: &str = "api.ts";

pub struct Endpoint {
    pub name: &'static str,
    pub method: &'static str,
    pub path: &'static str,
    pub render: fn(&mut TypeCollector) -> EndpointSpec,
}

inventory::collect!(Endpoint);

#[derive(Debug, Clone)]
pub struct EndpointSpec {
    pub name: &'static str,
    pub method: &'static str,
    pub path: &'static str,
    pub params: Vec<FieldSpec>,
    pub request_body: Option<String>,
    pub responses: Vec<ResponseSpec>,
}

#[derive(Debug, Clone)]
pub struct FieldSpec {
    pub name: String,
    pub ty: String,
    pub required: bool,
}

#[derive(Debug, Clone)]
pub struct ResponseSpec {
    pub status: &'static str,
    pub body: Option<String>,
}

pub struct EndpointRender<'a> {
    collector: &'a mut TypeCollector,
    spec: EndpointSpec,
}

impl<'a> EndpointRender<'a> {
    pub fn new(
        collector: &'a mut TypeCollector,
        name: &'static str,
        method: &'static str,
        path: &'static str,
    ) -> Self {
        Self {
            collector,
            spec: EndpointSpec {
                name,
                method,
                path,
                params: Vec::new(),
                request_body: None,
                responses: Vec::new(),
            },
        }
    }

    pub fn param<T>(&mut self, name: &'static str)
    where
        T: ts_rs::TS + 'static,
    {
        self.spec.params.push(FieldSpec {
            name: name.to_owned(),
            ty: self.collector.type_ref::<T>(),
            required: true,
        });
    }

    pub fn params<T>(&mut self)
    where
        T: utoipa::IntoParams + utoipa::ToSchema,
    {
        self.collector.collect_schema_declarations::<T>();

        self.spec.params.extend(
            T::into_params(|| Some(utoipa::openapi::path::ParameterIn::Query))
                .into_iter()
                .map(|param| {
                    let nullable = param.schema.as_ref().is_some_and(schema_ref_is_nullable);
                    let defaulted = param.schema.as_ref().is_some_and(schema_ref_has_default);
                    let required = matches!(param.required, utoipa::openapi::Required::True)
                        && !nullable
                        && !defaulted;

                    FieldSpec {
                        name: param.name,
                        ty: param
                            .schema
                            .as_ref()
                            .map(schema_ref_to_ts)
                            .unwrap_or_else(|| "unknown".to_owned()),
                        required,
                    }
                }),
        );
    }

    pub fn request_body<T>(&mut self)
    where
        T: ts_rs::TS + 'static,
    {
        self.spec.request_body = Some(self.collector.type_ref::<T>());
    }

    pub fn response<T>(&mut self, status: &'static str)
    where
        T: ts_rs::TS + 'static,
    {
        self.spec.responses.push(ResponseSpec {
            status,
            body: Some(self.collector.type_ref::<T>()),
        });
    }

    pub fn empty_response(&mut self, status: &'static str) {
        self.spec
            .responses
            .push(ResponseSpec { status, body: None });
    }

    pub fn finish(self) -> EndpointSpec {
        self.spec
    }
}

pub struct TypeCollector {
    cfg: ts_rs::Config,
    seen: HashSet<TypeId>,
    declarations: BTreeMap<String, String>,
}

impl TypeCollector {
    pub fn new() -> Self {
        Self {
            cfg: ts_rs::Config::from_env(),
            seen: HashSet::new(),
            declarations: BTreeMap::new(),
        }
    }

    pub fn type_ref<T>(&mut self) -> String
    where
        T: ts_rs::TS + 'static,
    {
        self.collect::<T>();
        T::name(&self.cfg)
    }

    fn collect<T>(&mut self)
    where
        T: ts_rs::TS + 'static,
    {
        if !self.seen.insert(TypeId::of::<T>()) {
            return;
        }

        struct Visitor<'a>(&'a mut TypeCollector);

        impl ts_rs::TypeVisitor for Visitor<'_> {
            fn visit<T>(&mut self)
            where
                T: ts_rs::TS + 'static + ?Sized,
            {
                self.0.collect_unsized::<T>();
            }
        }

        T::visit_dependencies(&mut Visitor(self));
        T::visit_generics(&mut Visitor(self));
        self.insert_declaration::<T>();
    }

    fn collect_unsized<T>(&mut self)
    where
        T: ts_rs::TS + 'static + ?Sized,
    {
        if !self.seen.insert(TypeId::of::<T>()) {
            return;
        }

        struct Visitor<'a>(&'a mut TypeCollector);

        impl ts_rs::TypeVisitor for Visitor<'_> {
            fn visit<T>(&mut self)
            where
                T: ts_rs::TS + 'static + ?Sized,
            {
                self.0.collect_unsized::<T>();
            }
        }

        T::visit_dependencies(&mut Visitor(self));
        T::visit_generics(&mut Visitor(self));
        self.insert_declaration_unsized::<T>();
    }

    fn insert_declaration<T>(&mut self)
    where
        T: ts_rs::TS + 'static,
    {
        self.insert_declaration_unsized::<T>();
    }

    fn insert_declaration_unsized<T>(&mut self)
    where
        T: ts_rs::TS + 'static + ?Sized,
    {
        if T::output_path().is_none() {
            return;
        }

        let ident = T::ident(&self.cfg);
        self.declarations
            .entry(ident)
            .or_insert_with(|| format!("export {}", T::decl(&self.cfg)));
    }

    fn collect_schema_declarations<T>(&mut self)
    where
        T: utoipa::ToSchema,
    {
        let mut schemas = Vec::new();
        T::schemas(&mut schemas);

        for (name, schema) in schemas {
            self.declarations
                .entry(name.clone())
                .or_insert_with(|| render_schema_declaration(&name, &schema));
        }
    }
}

impl Default for TypeCollector {
    fn default() -> Self {
        Self::new()
    }
}

pub fn export_all(path: impl AsRef<Path>) -> std::io::Result<()> {
    let path = resolve_export_path(path);
    let mut collector = TypeCollector::new();
    let mut endpoints = inventory::iter::<Endpoint>
        .into_iter()
        .map(|endpoint| (endpoint.render)(&mut collector))
        .collect::<Vec<_>>();

    endpoints.sort_by_key(|endpoint| (endpoint.method, endpoint.path, endpoint.name));

    if let Some(parent) = path
        .parent()
        .filter(|parent| !parent.as_os_str().is_empty())
    {
        std::fs::create_dir_all(parent)?;
    }

    std::fs::write(path, render_file(&collector.declarations, &endpoints))
}

pub fn export_all_default() -> std::io::Result<()> {
    export_all_from_path_or_env(None::<&Path>)
}

pub fn export_all_from_path_or_env(path: Option<impl AsRef<Path>>) -> std::io::Result<()> {
    let path = path
        .map(|path| path.as_ref().to_path_buf())
        .or_else(|| std::env::var_os(EXPORT_PATH_ENV).map(PathBuf::from))
        .unwrap_or_else(|| PathBuf::from(DEFAULT_EXPORT_PATH));

    export_all(path)
}

fn resolve_export_path(path: impl AsRef<Path>) -> PathBuf {
    let path = path.as_ref();

    if path.is_dir() || path.extension().is_none() {
        path.join(DEFAULT_EXPORT_FILE_NAME)
    } else {
        path.to_path_buf()
    }
}

#[macro_export]
macro_rules! export {
    () => {
        #[test]
        fn export_api() -> ::std::io::Result<()> {
            $crate::export_all_default()
        }
    };

    ($path:expr $(,)?) => {
        #[test]
        fn export_api() -> ::std::io::Result<()> {
            $crate::export_all_from_path_or_env(Some($path))
        }
    };
}

fn render_file(declarations: &BTreeMap<String, String>, endpoints: &[EndpointSpec]) -> String {
    let mut out = String::new();
    out.push_str(NOTE);

    for declaration in declarations.values() {
        out.push('\n');
        out.push_str(declaration);
        out.push('\n');
    }

    out.push_str("\nexport type Api = {\n");

    for endpoint in endpoints {
        let _ = writeln!(
            out,
            "  \"{} {}\": {{",
            endpoint.method,
            endpoint.path.replace('"', "\\\"")
        );

        if !&endpoint.params.is_empty() {
            write_fields_object(&mut out, "params", &endpoint.params, 4);
        }

        if let Some(body) = &endpoint.request_body {
            let _ = writeln!(out, "    body: {};", body);
        }

        out.push_str("    responses: {\n");
        for response in &endpoint.responses {
            let body = response.body.as_deref().unwrap_or("never");
            let _ = writeln!(out, "      {}: {};", ts_key(response.status), body);
        }
        out.push_str("    };\n");

        out.push_str("  };\n");
    }

    out.push_str("};\n");
    out
}

fn write_fields_object(out: &mut String, name: &str, fields: &[FieldSpec], indent: usize) {
    let padding = " ".repeat(indent);
    if fields.is_empty() {
        let _ = writeln!(out, "{padding}{name}: Record<string, never>;");
        return;
    }

    let _ = writeln!(out, "{padding}{name}: {{");
    for field in fields {
        let optional = if field.required { "" } else { "?" };
        let _ = writeln!(
            out,
            "{padding}  {}{}: {};",
            ts_key(&field.name),
            optional,
            field.ty
        );
    }
    let _ = writeln!(out, "{padding}}};");
}

fn schema_ref_to_ts(schema: &utoipa::openapi::RefOr<utoipa::openapi::schema::Schema>) -> String {
    match schema {
        utoipa::openapi::RefOr::T(schema) => schema_to_ts(schema),
        utoipa::openapi::RefOr::Ref(reference) => reference
            .ref_location
            .rsplit('/')
            .next()
            .filter(|name| !name.is_empty())
            .unwrap_or("unknown")
            .to_owned(),
    }
}

fn schema_ref_is_nullable(
    schema: &utoipa::openapi::RefOr<utoipa::openapi::schema::Schema>,
) -> bool {
    match schema {
        utoipa::openapi::RefOr::T(schema) => schema_is_nullable(schema),
        utoipa::openapi::RefOr::Ref(_) => false,
    }
}

fn schema_ref_has_default(
    schema: &utoipa::openapi::RefOr<utoipa::openapi::schema::Schema>,
) -> bool {
    match schema {
        utoipa::openapi::RefOr::T(schema) => schema_has_default(schema),
        utoipa::openapi::RefOr::Ref(_) => false,
    }
}

fn schema_is_nullable(schema: &utoipa::openapi::schema::Schema) -> bool {
    match schema {
        utoipa::openapi::schema::Schema::Object(object) => {
            schema_type_is_nullable(&object.schema_type)
        }
        utoipa::openapi::schema::Schema::OneOf(one_of) => {
            one_of.items.iter().any(schema_ref_is_nullable)
        }
        utoipa::openapi::schema::Schema::AllOf(all_of) => {
            all_of.items.iter().any(schema_ref_is_nullable)
        }
        utoipa::openapi::schema::Schema::AnyOf(any_of) => {
            any_of.items.iter().any(schema_ref_is_nullable)
        }
        _ => false,
    }
}

fn schema_has_default(schema: &utoipa::openapi::schema::Schema) -> bool {
    match schema {
        utoipa::openapi::schema::Schema::Object(object) => object.default.is_some(),
        utoipa::openapi::schema::Schema::Array(array) => array.default.is_some(),
        utoipa::openapi::schema::Schema::OneOf(one_of) => one_of.default.is_some(),
        utoipa::openapi::schema::Schema::AllOf(all_of) => all_of.default.is_some(),
        utoipa::openapi::schema::Schema::AnyOf(any_of) => any_of.default.is_some(),
        _ => false,
    }
}

fn schema_type_is_nullable(schema_type: &utoipa::openapi::schema::SchemaType) -> bool {
    match schema_type {
        utoipa::openapi::schema::SchemaType::Type(utoipa::openapi::schema::Type::Null) => true,
        utoipa::openapi::schema::SchemaType::Array(types) => {
            types.contains(&utoipa::openapi::schema::Type::Null)
        }
        _ => false,
    }
}

fn schema_to_ts(schema: &utoipa::openapi::schema::Schema) -> String {
    match schema {
        utoipa::openapi::schema::Schema::Object(object) => object_to_ts(object),
        utoipa::openapi::schema::Schema::Array(array) => match &array.items {
            utoipa::openapi::schema::ArrayItems::RefOrSchema(item) => {
                format!("{}[]", schema_ref_to_ts(item))
            }
            utoipa::openapi::schema::ArrayItems::False => "never[]".to_owned(),
        },
        utoipa::openapi::schema::Schema::OneOf(one_of) => one_of
            .items
            .iter()
            .map(schema_ref_to_ts)
            .collect::<Vec<_>>()
            .join(" | "),
        utoipa::openapi::schema::Schema::AllOf(all_of) => all_of
            .items
            .iter()
            .map(schema_ref_to_ts)
            .collect::<Vec<_>>()
            .join(" & "),
        utoipa::openapi::schema::Schema::AnyOf(any_of) => any_of
            .items
            .iter()
            .map(schema_ref_to_ts)
            .collect::<Vec<_>>()
            .join(" | "),
        _ => "unknown".to_owned(),
    }
}

fn render_schema_declaration(
    name: &str,
    schema: &utoipa::openapi::RefOr<utoipa::openapi::schema::Schema>,
) -> String {
    format!(
        "export type {} = {};",
        ts_key(name),
        schema_ref_to_ts(schema)
    )
}

fn object_to_ts(object: &utoipa::openapi::schema::Object) -> String {
    if let Some(enum_values) = &object.enum_values {
        if enum_values.is_empty() {
            return "never".to_owned();
        }

        return enum_values
            .iter()
            .map(ToString::to_string)
            .collect::<Vec<_>>()
            .join(" | ");
    }

    if !object.properties.is_empty() {
        let fields = object
            .properties
            .iter()
            .map(|(name, schema)| FieldSpec {
                name: name.clone(),
                ty: schema_ref_to_ts(schema),
                required: object.required.contains(name),
            })
            .collect::<Vec<_>>();
        let mut out = String::new();
        out.push_str("{\n");

        for field in fields {
            let optional = if field.required { "" } else { "?" };
            let _ = writeln!(out, "  {}{}: {};", ts_key(&field.name), optional, field.ty);
        }

        out.push('}');
        return out;
    }

    schema_type_to_ts(&object.schema_type)
}

fn schema_type_to_ts(schema_type: &utoipa::openapi::schema::SchemaType) -> String {
    match schema_type {
        utoipa::openapi::schema::SchemaType::Type(ty) => primitive_type_to_ts(ty).to_owned(),
        utoipa::openapi::schema::SchemaType::Array(types) => types
            .iter()
            .filter(|ty| **ty != utoipa::openapi::schema::Type::Null)
            .map(primitive_type_to_ts)
            .collect::<Vec<_>>()
            .join(" | "),
        utoipa::openapi::schema::SchemaType::AnyValue => "unknown".to_owned(),
    }
}

fn primitive_type_to_ts(ty: &utoipa::openapi::schema::Type) -> &'static str {
    match ty {
        utoipa::openapi::schema::Type::Object => "Record<string, unknown>",
        utoipa::openapi::schema::Type::String => "string",
        utoipa::openapi::schema::Type::Integer | utoipa::openapi::schema::Type::Number => "number",
        utoipa::openapi::schema::Type::Boolean => "boolean",
        utoipa::openapi::schema::Type::Array => "unknown[]",
        utoipa::openapi::schema::Type::Null => "null",
    }
}

fn ts_key(key: &str) -> String {
    if key
        .chars()
        .all(|ch| ch.is_ascii_alphanumeric() || ch == '_')
    {
        key.to_owned()
    } else {
        format!("\"{}\"", key.replace('"', "\\\""))
    }
}