spikard-cli 0.15.5

//! TypeScript code generation from `OpenAPI` schemas

use super::NodeDtoStyle;
use anyhow::Result;
use heck::{ToPascalCase, ToSnakeCase};
use openapiv3::{
    OpenAPI, Operation, Parameter, ParameterSchemaOrContent, ReferenceOr, Schema, SchemaKind, StringFormat, Type,
    VariantOrUnknownOrEmpty,
};
use std::collections::{HashMap, HashSet, VecDeque};

pub struct TypeScriptGenerator {
    spec: OpenAPI,
    dto: NodeDtoStyle,
}

impl TypeScriptGenerator {
    #[must_use]
    pub const fn new(spec: OpenAPI, dto: NodeDtoStyle) -> Self {
        Self { spec, dto }
    }

    pub fn generate(&self) -> Result<String> {
        let mut output = String::new();

        output.push_str(&self.generate_header());

        output.push_str(&self.generate_schemas()?);

        output.push_str(&self.generate_routes()?);

        output.push_str(&self.generate_main());

        Ok(output)
    }

    fn generate_header(&self) -> String {
        match self.dto {
            NodeDtoStyle::Zod => format!(
                r#"// Generated by Spikard OpenAPI code generator
// OpenAPI Version: {}
// Title: {}
// DO NOT EDIT - regenerate from OpenAPI schema

import {{ route }} from "spikard";
import type {{ Body, Path, Query, Request }} from "spikard";
import {{ z }} from "zod";

"#,
                self.spec.openapi, self.spec.info.title
            ),
        }
    }

    fn generate_schemas(&self) -> Result<String> {
        let mut output = String::new();
        output.push_str("// Zod Schemas\n\n");

        if let Some(components) = &self.spec.components {
            // Convert ReferenceOr::Item into a plain HashMap for sorting
            let mut schemas_map: HashMap<String, Schema> = HashMap::new();
            let mut schema_refs: HashMap<String, ReferenceOr<Schema>> = HashMap::new();

            for (name, schema_ref) in &components.schemas {
                schema_refs.insert(name.clone(), schema_ref.clone());
                if let ReferenceOr::Item(schema) = schema_ref {
                    schemas_map.insert(name.clone(), schema.clone());
                }
            }

            // Topologically sort the schemas
            let sorted_names = topological_sort_schemas(&schemas_map);

            // Generate schemas in topologically sorted order
            for name in sorted_names {
                if let Some(ReferenceOr::Item(schema)) = schema_refs.get(&name) {
                    output.push_str(&self.generate_zod_schema(&name, schema)?);
                    output.push('\n');
                }
            }
        }

        Ok(output)
    }

    fn generate_zod_schema(&self, name: &str, schema: &Schema) -> Result<String> {
        let schema_name = format!("{}Schema", name.to_pascal_case());
        let type_name = name.to_pascal_case();
        let mut output = String::new();

        if let Some(description) = &schema.schema_data.description {
            output.push_str(&format!("/** {description} */\n"));
        }

        let mut schema_expr = match &schema.schema_kind {
            SchemaKind::Type(Type::Object(obj)) => {
                let mut expr = String::from("z.object({\n");

                for (prop_name, prop_schema_ref) in &obj.properties {
                    let is_required = obj.required.contains(prop_name);
                    let field_name = prop_name.to_snake_case();

                    let zod_type = match prop_schema_ref {
                        ReferenceOr::Item(prop_schema) => Self::schema_to_zod_type(prop_schema, !is_required),
                        ReferenceOr::Reference { reference } => {
                            let ref_name = reference.split('/').next_back().unwrap();
                            let ref_schema = format!("{}Schema", ref_name.to_pascal_case());
                            if is_required {
                                ref_schema
                            } else {
                                format!("{ref_schema}.optional()")
                            }
                        }
                    };

                    expr.push_str(&format!("\t{field_name}: {zod_type},\n"));
                }

                expr.push_str("})");
                expr
            }
            _ => "z.unknown()".to_string(),
        };

        if schema.schema_data.nullable {
            schema_expr.push_str(".nullable()");
        }

        output.push_str(&format!("export const {schema_name} = {schema_expr};\n"));

        output.push_str(&format!("\nexport type {type_name} = z.infer<typeof {schema_name}>;\n"));

        Ok(output)
    }

    /// Extract type name from a schema reference or inline schema
    fn extract_type_from_schema_ref(&self, schema_ref: &ReferenceOr<Schema>) -> String {
        match schema_ref {
            ReferenceOr::Reference { reference } => {
                let ref_name = reference.split('/').next_back().unwrap();
                ref_name.to_pascal_case()
            }
            ReferenceOr::Item(schema) => Self::schema_to_typescript_type(schema, false),
        }
    }

    /// Extract TypeScript type from a schema reference or inline schema
    fn extract_typescript_type_from_ref(&self, schema_ref: &ReferenceOr<Schema>) -> String {
        match schema_ref {
            ReferenceOr::Reference { reference } => {
                let ref_name = reference.split('/').next_back().unwrap();
                ref_name.to_pascal_case()
            }
            ReferenceOr::Item(schema) => Self::schema_to_typescript_type(schema, false),
        }
    }

    /// Extract request body TypeScript type from operation
    fn extract_request_body_type(&self, operation: &Operation) -> Option<String> {
        operation.request_body.as_ref().and_then(|body_ref| match body_ref {
            ReferenceOr::Item(request_body) => request_body.content.get("application/json").and_then(|media_type| {
                media_type
                    .schema
                    .as_ref()
                    .map(|schema_ref| self.extract_typescript_type_from_ref(schema_ref))
            }),
            ReferenceOr::Reference { reference } => {
                let ref_name = reference.split('/').next_back().unwrap();
                Some(ref_name.to_pascal_case())
            }
        })
    }

    /// Extract response type from operation (looks for 200/201 responses)
    fn extract_response_type(&self, operation: &Operation) -> String {
        use openapiv3::StatusCode;

        let response = operation
            .responses
            .responses
            .get(&StatusCode::Code(200))
            .or_else(|| operation.responses.responses.get(&StatusCode::Code(201)))
            .or_else(|| operation.responses.responses.get(&StatusCode::Range(2)));

        if let Some(response_ref) = response {
            match response_ref {
                ReferenceOr::Item(response) => {
                    if let Some(content) = response.content.get("application/json")
                        && let Some(schema_ref) = &content.schema
                    {
                        return self.extract_type_from_schema_ref(schema_ref);
                    }
                }
                ReferenceOr::Reference { reference } => {
                    let ref_name = reference.split('/').next_back().unwrap();
                    return ref_name.to_pascal_case();
                }
            }
        }

        "Record<string, unknown>".to_string()
    }

    fn schema_to_zod_type(schema: &Schema, optional: bool) -> String {
        let mut base_type = match &schema.schema_kind {
            SchemaKind::Type(Type::String(string_type)) => {
                let enum_values = string_type
                    .enumeration
                    .iter()
                    .flatten()
                    .map(|value| format!("{value:?}"))
                    .collect::<Vec<_>>();
                if !enum_values.is_empty() {
                    format!("z.enum([{}])", enum_values.join(", "))
                } else {
                    match &string_type.format {
                        VariantOrUnknownOrEmpty::Item(StringFormat::Date) => "z.string()".to_string(),
                        VariantOrUnknownOrEmpty::Item(StringFormat::DateTime) => "z.string().datetime()".to_string(),
                        VariantOrUnknownOrEmpty::Unknown(format) if format == "email" => {
                            "z.string().email()".to_string()
                        }
                        VariantOrUnknownOrEmpty::Unknown(format) if format == "uuid" => "z.string().uuid()".to_string(),
                        _ => "z.string()".to_string(),
                    }
                }
            }
            SchemaKind::Type(Type::Number(_)) => "z.number()".to_string(),
            SchemaKind::Type(Type::Integer(_)) => "z.number().int()".to_string(),
            SchemaKind::Type(Type::Boolean(_)) => "z.boolean()".to_string(),
            SchemaKind::Type(Type::Array(arr)) => {
                let item_type = match &arr.items {
                    Some(ReferenceOr::Item(item_schema)) => Self::schema_to_zod_type(item_schema, false),
                    Some(ReferenceOr::Reference { reference }) => {
                        let ref_name = reference.split('/').next_back().unwrap();
                        format!("{}Schema", ref_name.to_pascal_case())
                    }
                    None => "z.record(z.string(), z.unknown())".to_string(),
                };
                format!("z.array({item_type})")
            }
            SchemaKind::Type(Type::Object(obj)) => {
                if obj.properties.is_empty() {
                    "z.record(z.string(), z.unknown())".to_string()
                } else {
                    let mut fields = String::from("z.object({\n");
                    for (prop_name, prop_schema_ref) in &obj.properties {
                        let is_required = obj.required.contains(prop_name);
                        let field_name = prop_name.to_snake_case();
                        let zod_type = match prop_schema_ref {
                            ReferenceOr::Item(prop_schema) => Self::schema_to_zod_type(prop_schema, !is_required),
                            ReferenceOr::Reference { reference } => {
                                let ref_name = reference.split('/').next_back().unwrap();
                                let ref_schema = format!("{}Schema", ref_name.to_pascal_case());
                                if is_required {
                                    ref_schema
                                } else {
                                    format!("{ref_schema}.optional()")
                                }
                            }
                        };
                        fields.push_str(&format!("\t{field_name}: {zod_type},\n"));
                    }
                    fields.push_str("})");
                    fields
                }
            }
            SchemaKind::OneOf { one_of } => {
                let members = one_of
                    .iter()
                    .map(|schema_ref| match schema_ref {
                        ReferenceOr::Item(item_schema) => Self::schema_to_zod_type(item_schema, false),
                        ReferenceOr::Reference { reference } => {
                            let ref_name = reference.split('/').next_back().unwrap();
                            format!("{}Schema", ref_name.to_pascal_case())
                        }
                    })
                    .collect::<Vec<_>>();
                format!("z.union([{}])", members.join(", "))
            }
            SchemaKind::AnyOf { any_of } => {
                let members = any_of
                    .iter()
                    .map(|schema_ref| match schema_ref {
                        ReferenceOr::Item(item_schema) => Self::schema_to_zod_type(item_schema, false),
                        ReferenceOr::Reference { reference } => {
                            let ref_name = reference.split('/').next_back().unwrap();
                            format!("{}Schema", ref_name.to_pascal_case())
                        }
                    })
                    .collect::<Vec<_>>();
                format!("z.union([{}])", members.join(", "))
            }
            SchemaKind::AllOf { all_of } => {
                let members = all_of
                    .iter()
                    .map(|schema_ref| match schema_ref {
                        ReferenceOr::Item(item_schema) => Self::schema_to_zod_type(item_schema, false),
                        ReferenceOr::Reference { reference } => {
                            let ref_name = reference.split('/').next_back().unwrap();
                            format!("{}Schema", ref_name.to_pascal_case())
                        }
                    })
                    .collect::<Vec<_>>();
                match members.split_first() {
                    Some((first, rest)) => rest
                        .iter()
                        .fold(first.clone(), |acc, member| format!("{acc}.and({member})")),
                    None => "z.unknown()".to_string(),
                }
            }
            _ => "z.unknown()".to_string(),
        };

        if schema.schema_data.nullable {
            base_type.push_str(".nullable()");
        }

        if optional {
            base_type.push_str(".optional()");
        }

        base_type
    }

    fn schema_to_typescript_type(schema: &Schema, optional: bool) -> String {
        let mut base_type = match &schema.schema_kind {
            SchemaKind::Type(Type::String(string_type)) => {
                let enum_values = string_type
                    .enumeration
                    .iter()
                    .flatten()
                    .map(|value| format!("{value:?}"))
                    .collect::<Vec<_>>();
                if enum_values.is_empty() {
                    "string".to_string()
                } else {
                    enum_values.join(" | ")
                }
            }
            SchemaKind::Type(Type::Number(_) | Type::Integer(_)) => "number".to_string(),
            SchemaKind::Type(Type::Boolean(_)) => "boolean".to_string(),
            SchemaKind::Type(Type::Array(arr)) => {
                let item_type = match &arr.items {
                    Some(ReferenceOr::Item(item_schema)) => Self::schema_to_typescript_type(item_schema, false),
                    Some(ReferenceOr::Reference { reference }) => {
                        let ref_name = reference.split('/').next_back().unwrap();
                        ref_name.to_pascal_case()
                    }
                    None => "Record<string, unknown>".to_string(),
                };
                let item_type = if item_type.contains(" | ") {
                    format!("({item_type})")
                } else {
                    item_type
                };
                format!("{item_type}[]")
            }
            SchemaKind::Type(Type::Object(obj)) => {
                if obj.properties.is_empty() {
                    "Record<string, unknown>".to_string()
                } else {
                    let fields = obj
                        .properties
                        .iter()
                        .map(|(prop_name, prop_schema_ref)| {
                            let optional_marker = if obj.required.contains(prop_name) { "" } else { "?" };
                            let prop_type = match prop_schema_ref {
                                ReferenceOr::Item(prop_schema) => Self::schema_to_typescript_type(prop_schema, false),
                                ReferenceOr::Reference { reference } => {
                                    let ref_name = reference.split('/').next_back().unwrap();
                                    ref_name.to_pascal_case()
                                }
                            };
                            format!("{prop_name}{optional_marker}: {prop_type}")
                        })
                        .collect::<Vec<_>>()
                        .join("; ");
                    format!("{{ {fields} }}")
                }
            }
            SchemaKind::OneOf { one_of } | SchemaKind::AnyOf { any_of: one_of } => one_of
                .iter()
                .map(|schema_ref| match schema_ref {
                    ReferenceOr::Item(item_schema) => Self::schema_to_typescript_type(item_schema, false),
                    ReferenceOr::Reference { reference } => {
                        let ref_name = reference.split('/').next_back().unwrap();
                        ref_name.to_pascal_case()
                    }
                })
                .collect::<Vec<_>>()
                .join(" | "),
            SchemaKind::AllOf { all_of } => all_of
                .iter()
                .map(|schema_ref| match schema_ref {
                    ReferenceOr::Item(item_schema) => Self::schema_to_typescript_type(item_schema, false),
                    ReferenceOr::Reference { reference } => {
                        let ref_name = reference.split('/').next_back().unwrap();
                        ref_name.to_pascal_case()
                    }
                })
                .collect::<Vec<_>>()
                .join(" & "),
            _ => "unknown".to_string(),
        };

        if schema.schema_data.nullable {
            base_type.push_str(" | null");
        }

        if optional {
            base_type.push_str(" | undefined");
        }

        base_type
    }

    fn generate_routes(&self) -> Result<String> {
        let mut output = String::new();
        output.push_str("\n// Route Handlers\n\n");

        for (path, path_item_ref) in &self.spec.paths.paths {
            let path_item = match path_item_ref {
                ReferenceOr::Item(item) => item,
                ReferenceOr::Reference { .. } => continue,
            };

            if let Some(op) = &path_item.get {
                output.push_str(&self.generate_route_handler(path, "get", op)?);
            }
            if let Some(op) = &path_item.post {
                output.push_str(&self.generate_route_handler(path, "post", op)?);
            }
            if let Some(op) = &path_item.put {
                output.push_str(&self.generate_route_handler(path, "put", op)?);
            }
            if let Some(op) = &path_item.delete {
                output.push_str(&self.generate_route_handler(path, "delete", op)?);
            }
            if let Some(op) = &path_item.patch {
                output.push_str(&self.generate_route_handler(path, "patch", op)?);
            }
        }

        Ok(output)
    }

    fn generate_route_handler(&self, path: &str, method: &str, operation: &Operation) -> Result<String> {
        let mut output = String::new();

        if let Some(summary) = &operation.summary {
            output.push_str(&format!("/**\n * {summary}\n"));
        } else {
            output.push_str("/**\n");
        }
        output.push_str(&format!(" * Route: {} {}\n", method.to_uppercase(), path));
        output.push_str(" */\n");

        let func_name = operation
            .operation_id
            .as_ref()
            .map(|id| id.to_snake_case())
            .unwrap_or_else(|| {
                format!(
                    "{}_{}",
                    method,
                    path.replace('/', "_").replace(['{', '}'], "").trim_matches('_')
                )
            });

        let mut path_params = Vec::new();
        let mut query_params = Vec::new();

        for param_ref in &operation.parameters {
            if let ReferenceOr::Item(param) = param_ref {
                match param {
                    Parameter::Path { parameter_data, .. } => {
                        let type_hint = Self::parameter_typescript_type(parameter_data);
                        path_params.push((parameter_data.name.clone(), type_hint));
                    }
                    Parameter::Query { parameter_data, .. } => {
                        let type_hint = Self::parameter_typescript_type(parameter_data);
                        query_params.push((parameter_data.name.clone(), type_hint, parameter_data.required));
                    }
                    _ => {}
                }
            }
        }

        let body_type = self.extract_request_body_type(operation);

        let return_type = self.extract_response_type(operation);

        output.push_str(&format!("export function {func_name}(_request: Request"));

        for (param_name, param_type) in &path_params {
            output.push_str(&format!(", _{}: Path<{}>", param_name.to_snake_case(), param_type));
        }

        for (param_name, param_type, required) in &query_params {
            if *required {
                output.push_str(&format!(", _{}: Query<{}>", param_name.to_snake_case(), param_type));
            } else {
                output.push_str(&format!(
                    ", _{}: Query<{} | undefined>",
                    param_name.to_snake_case(),
                    param_type
                ));
            }
        }

        if let Some(body_type) = &body_type {
            output.push_str(&format!(", _body: Body<{body_type}>"));
        }

        output.push_str(&format!("): {return_type} {{\n"));

        if let Some(desc) = &operation.description {
            output.push_str(&format!("\t/**\n\t * {desc}\n\t */\n"));
        }

        output.push_str("\tthrow new Error(\"TODO: Implement this endpoint\");\n");
        output.push_str("}\n");

        output.push_str(&format!(
            "route(\"{}\", {{ methods: [\"{}\"] }})({});\n\n",
            path,
            method.to_uppercase(),
            func_name
        ));

        Ok(output)
    }

    fn generate_main(&self) -> String {
        r"
// Run the application
// Note: Actual server setup depends on your runtime configuration
"
        .to_string()
    }

    fn parameter_typescript_type(parameter_data: &openapiv3::ParameterData) -> String {
        match &parameter_data.format {
            ParameterSchemaOrContent::Schema(schema_ref) => match schema_ref {
                ReferenceOr::Item(schema) => Self::schema_to_typescript_type(schema, false),
                ReferenceOr::Reference { reference } => {
                    let ref_name = reference.split('/').next_back().unwrap();
                    ref_name.to_pascal_case()
                }
            },
            ParameterSchemaOrContent::Content(_) => "unknown".to_string(),
        }
    }
}

/// Extract all schema names that are referenced in a given schema
fn extract_schema_dependencies(schema: &Schema) -> HashSet<String> {
    let mut dependencies = HashSet::new();
    extract_dependencies_recursive(schema, &mut dependencies);
    dependencies
}

/// Recursively extract all $ref schema names from a schema
fn extract_dependencies_recursive(schema: &Schema, deps: &mut HashSet<String>) {
    match &schema.schema_kind {
        SchemaKind::Type(Type::Object(obj)) => {
            // Extract dependencies from properties
            for (_prop_name, prop_schema_ref) in &obj.properties {
                match prop_schema_ref {
                    ReferenceOr::Reference { reference } => {
                        if let Some(ref_name) = reference.split('/').next_back() {
                            deps.insert(ref_name.to_string());
                        }
                    }
                    ReferenceOr::Item(prop_schema) => {
                        extract_dependencies_recursive(prop_schema, deps);
                    }
                }
            }
        }
        SchemaKind::Type(Type::Array(arr)) => {
            // Extract dependencies from array items
            if let Some(items) = &arr.items {
                match items {
                    ReferenceOr::Reference { reference } => {
                        if let Some(ref_name) = reference.split('/').next_back() {
                            deps.insert(ref_name.to_string());
                        }
                    }
                    ReferenceOr::Item(item_schema) => {
                        extract_dependencies_recursive(item_schema, deps);
                    }
                }
            }
        }
        _ => {}
    }
}

/// Topologically sort schemas by their dependencies using Kahn's algorithm
fn topological_sort_schemas(schemas: &HashMap<String, Schema>) -> Vec<String> {
    let mut in_degree: HashMap<String, usize> = HashMap::new();
    let mut graph: HashMap<String, Vec<String>> = HashMap::new();

    // Initialize all schemas with zero in-degree
    for schema_name in schemas.keys() {
        in_degree.insert(schema_name.clone(), 0);
        graph.insert(schema_name.clone(), Vec::new());
    }

    // Build dependency graph
    for (schema_name, schema) in schemas {
        let deps = extract_schema_dependencies(schema);
        for dep in deps {
            // Only track dependencies that are defined in the schema set
            if schemas.contains_key(&dep) {
                // Add edge from dependency to current schema
                graph.entry(dep).or_default().push(schema_name.clone());
                // Increment in-degree for current schema
                *in_degree.get_mut(schema_name).unwrap() += 1;
            }
        }
    }

    // Find all schemas with no incoming edges
    let mut queue: VecDeque<String> = in_degree
        .iter()
        .filter(|(_, deg)| **deg == 0)
        .map(|(name, _)| name.clone())
        .collect();

    let mut result = Vec::new();

    // Process schemas in topological order
    while let Some(node) = queue.pop_front() {
        result.push(node.clone());

        // For each neighbor of the current node
        if let Some(neighbors) = graph.get(&node) {
            for neighbor in neighbors {
                let deg = in_degree.get_mut(neighbor).unwrap();
                *deg -= 1;
                if *deg == 0 {
                    queue.push_back(neighbor.clone());
                }
            }
        }
    }

    result
}