vika_cli/generator/
zod_schema.rs

1use crate::error::Result;
2use crate::generator::swagger_parser::{get_schema_name_from_ref, resolve_ref};
3use crate::generator::utils::to_pascal_case;
4use openapiv3::{OpenAPI, ReferenceOr, Schema, SchemaKind, Type};
5use std::collections::HashMap;
6
7pub struct ZodSchema {
8    pub content: String,
9}
10
11pub fn generate_zod_schemas(
12    openapi: &OpenAPI,
13    schemas: &HashMap<String, Schema>,
14    schema_names: &[String],
15) -> Result<Vec<ZodSchema>> {
16    generate_zod_schemas_with_registry(
17        openapi,
18        schemas,
19        schema_names,
20        &mut std::collections::HashMap::new(),
21    )
22}
23
24pub fn generate_zod_schemas_with_registry(
25    openapi: &OpenAPI,
26    schemas: &HashMap<String, Schema>,
27    schema_names: &[String],
28    enum_registry: &mut std::collections::HashMap<String, String>,
29) -> Result<Vec<ZodSchema>> {
30    let mut zod_schemas = Vec::new();
31    let mut processed = std::collections::HashSet::new();
32
33    for schema_name in schema_names {
34        if let Some(schema) = schemas.get(schema_name) {
35            generate_zod_for_schema(
36                openapi,
37                schema_name,
38                schema,
39                &mut zod_schemas,
40                &mut processed,
41                enum_registry,
42                None,
43            )?;
44        }
45    }
46
47    Ok(zod_schemas)
48}
49
50fn generate_zod_for_schema(
51    openapi: &OpenAPI,
52    name: &str,
53    schema: &Schema,
54    zod_schemas: &mut Vec<ZodSchema>,
55    processed: &mut std::collections::HashSet<String>,
56    enum_registry: &mut std::collections::HashMap<String, String>,
57    parent_schema_name: Option<&str>,
58) -> Result<()> {
59    if processed.contains(name) {
60        return Ok(());
61    }
62    processed.insert(name.to_string());
63
64    let schema_name = to_pascal_case(name);
65    let zod_def = schema_to_zod(
66        openapi,
67        schema,
68        zod_schemas,
69        processed,
70        0,
71        enum_registry,
72        None,
73        parent_schema_name,
74    )?;
75
76    // Handle enums at top level (when schema itself is an enum)
77    // Note: For property-level enums, they're handled in schema_to_zod with context
78    if let SchemaKind::Type(Type::String(string_type)) = &schema.schema_kind {
79        if !string_type.enumeration.is_empty() {
80            let mut enum_values: Vec<String> = string_type
81                .enumeration
82                .iter()
83                .filter_map(|v| v.as_ref().cloned())
84                .collect();
85            if !enum_values.is_empty() {
86                enum_values.sort();
87                let enum_key = enum_values.join(",");
88
89                // Check if this enum already exists in registry
90                if enum_registry.get(&enum_key).is_some() {
91                    // Enum already generated, skip
92                    return Ok(());
93                }
94
95                // Use schema name if available, otherwise generate from values
96                let enum_name = if !name.is_empty() {
97                    format!("{}Enum", to_pascal_case(name))
98                } else {
99                    // Generate unique name from values
100                    let value_hash: String = enum_values
101                        .iter()
102                        .take(3)
103                        .map(|v| v.chars().next().unwrap_or('X'))
104                        .collect();
105                    format!("Enum{}", value_hash)
106                };
107
108                // Store in registry
109                enum_registry.insert(enum_key, enum_name.clone());
110
111                let enum_schema = generate_enum_zod(&enum_name, &enum_values);
112                zod_schemas.push(enum_schema);
113                return Ok(());
114            }
115        }
116    }
117
118    // Only create object schema if it's an object type
119    if matches!(&schema.schema_kind, SchemaKind::Type(Type::Object(_))) {
120        zod_schemas.push(ZodSchema {
121            content: format!(
122                "export const {}Schema = z.object({{\n{}\n}});",
123                schema_name, zod_def
124            ),
125        });
126    }
127
128    Ok(())
129}
130
131#[allow(clippy::too_many_arguments)]
132fn schema_to_zod(
133    openapi: &OpenAPI,
134    schema: &Schema,
135    zod_schemas: &mut Vec<ZodSchema>,
136    processed: &mut std::collections::HashSet<String>,
137    indent: usize,
138    enum_registry: &mut std::collections::HashMap<String, String>,
139    context: Option<(&str, &str)>, // (property_name, parent_schema_name)
140    current_schema_name: Option<&str>, // Current schema being processed (for enum naming context)
141) -> Result<String> {
142    // Prevent infinite recursion with a reasonable depth limit
143    if indent > 100 {
144        return Ok(format!("{}z.any()", "  ".repeat(indent)));
145    }
146    let indent_str = "  ".repeat(indent);
147
148    match &schema.schema_kind {
149        SchemaKind::Type(type_) => {
150            match type_ {
151                Type::String(string_type) => {
152                    if !string_type.enumeration.is_empty() {
153                        let mut enum_values: Vec<String> = string_type
154                            .enumeration
155                            .iter()
156                            .filter_map(|v| v.as_ref().cloned())
157                            .collect();
158                        if !enum_values.is_empty() {
159                            // Create a key from sorted enum values to check registry
160                            enum_values.sort();
161                            let enum_key = enum_values.join(",");
162
163                            // For generic property names, include context in the key to avoid conflicts
164                            let context_key = if let Some((prop_name, parent_schema)) = context {
165                                let generic_names = ["status", "type", "state", "kind"];
166                                if generic_names.contains(&prop_name.to_lowercase().as_str())
167                                    && !parent_schema.is_empty()
168                                {
169                                    // Include parent schema in key for generic properties to avoid conflicts
170                                    format!("{}:{}", enum_key, parent_schema)
171                                } else {
172                                    enum_key.clone()
173                                }
174                            } else {
175                                enum_key.clone()
176                            };
177
178                            // Check if this enum already exists in registry
179                            // First check context_key (for context-aware enums)
180                            // Then check base enum_key to deduplicate enums with same values
181                            let existing_enum_name = enum_registry
182                                .get(&context_key)
183                                .or_else(|| enum_registry.get(&enum_key))
184                                .cloned();
185
186                            if let Some(existing_name) = existing_enum_name {
187                                // Store in registry with context_key for future lookups
188                                enum_registry.insert(context_key.clone(), existing_name.clone());
189
190                                // Check if enum schema has already been generated
191                                let schema_already_generated = zod_schemas.iter().any(|s| {
192                                    s.content.contains(&format!("{}Schema", existing_name))
193                                });
194
195                                // If not generated yet, generate it now
196                                if !schema_already_generated {
197                                    let enum_schema =
198                                        generate_enum_zod(&existing_name, &enum_values);
199                                    zod_schemas.push(enum_schema);
200                                }
201
202                                // Enums don't need z.lazy(), use directly
203                                return Ok(format!("{}{}Schema", indent_str, existing_name));
204                            }
205
206                            // Generate meaningful enum name using context (property name + parent schema) or fallback
207                            let enum_name = if let Some((prop_name, parent_schema)) = context {
208                                // Use property name + parent schema for meaningful name to avoid conflicts
209                                // For generic names like "status", use parent schema to differentiate
210                                let prop_pascal = to_pascal_case(prop_name);
211
212                                // If property name is generic (status, type, etc.), use parent schema
213                                let generic_names = ["status", "type", "state", "kind"];
214                                if generic_names.contains(&prop_name.to_lowercase().as_str())
215                                    && !parent_schema.is_empty()
216                                {
217                                    let parent_pascal = to_pascal_case(parent_schema);
218                                    // Remove common suffixes from parent schema name
219                                    let parent_clean = parent_pascal
220                                        .trim_end_matches("ResponseDto")
221                                        .trim_end_matches("Dto")
222                                        .trim_end_matches("Response")
223                                        .to_string();
224
225                                    // Check if parent already contains the property name (e.g., "KycStatus" contains "Status")
226                                    // Use case-insensitive matching and check if property name is a suffix or contained
227                                    let prop_lower = prop_pascal.to_lowercase();
228                                    let parent_lower = parent_clean.to_lowercase();
229
230                                    // Check if parent ends with property name (e.g., "KycStatus" ends with "Status")
231                                    // or if property is contained in parent (case-insensitive)
232                                    if parent_lower.ends_with(&prop_lower)
233                                        || parent_lower.contains(&prop_lower)
234                                    {
235                                        // Parent already contains property name, just use parent + Enum
236                                        format!("{}Enum", parent_clean)
237                                    } else {
238                                        // Combine parent + property
239                                        format!("{}{}Enum", parent_clean, prop_pascal)
240                                    }
241                                } else {
242                                    format!("{}Enum", prop_pascal)
243                                }
244                            } else if !enum_values.is_empty() {
245                                // Fallback: use first value to create name
246                                let first_value = &enum_values[0];
247                                let base_name = first_value
248                                    .chars()
249                                    .take(1)
250                                    .collect::<String>()
251                                    .to_uppercase()
252                                    + &first_value.chars().skip(1).collect::<String>();
253                                format!("{}Enum", to_pascal_case(&base_name))
254                            } else {
255                                "UnknownEnum".to_string()
256                            };
257
258                            // Store in registry using context_key (includes context for generic properties)
259                            // Also store with base enum_key for deduplication
260                            enum_registry.insert(context_key.clone(), enum_name.clone());
261                            enum_registry.insert(enum_key.clone(), enum_name.clone());
262
263                            // Generate enum schema
264                            let enum_schema = generate_enum_zod(&enum_name, &enum_values);
265                            zod_schemas.push(enum_schema);
266
267                            // Enums don't need z.lazy(), use directly
268                            return Ok(format!("{}{}Schema", indent_str, enum_name));
269                        }
270                    }
271                    let mut zod_expr = format!("{}z.string()", indent_str);
272
273                    // Add string constraints
274                    if let Some(min_length) = string_type.min_length {
275                        zod_expr = format!("{}.min({})", zod_expr, min_length);
276                    }
277                    if let Some(max_length) = string_type.max_length {
278                        zod_expr = format!("{}.max({})", zod_expr, max_length);
279                    }
280                    if let Some(pattern) = &string_type.pattern {
281                        // Escape regex pattern for JavaScript
282                        let escaped_pattern = pattern.replace('\\', "\\\\").replace('"', "\\\"");
283                        zod_expr = format!("{}.regex(/^{}$/)", zod_expr, escaped_pattern);
284                    }
285                    // Handle format - it's a VariantOrUnknownOrEmpty
286                    // Try to extract format as string from both Item and Unknown variants
287                    let format_str_opt = match &string_type.format {
288                        openapiv3::VariantOrUnknownOrEmpty::Item(format) => {
289                            // Convert enum variant to string using Debug
290                            let debug_str = format!("{:?}", format);
291                            // Extract format name (e.g., "Email" from "Email" or "StringFormat::Email")
292                            let format_name = debug_str
293                                .split("::")
294                                .last()
295                                .unwrap_or(&debug_str)
296                                .to_lowercase();
297                            Some(format_name)
298                        }
299                        openapiv3::VariantOrUnknownOrEmpty::Unknown(s) => Some(s.clone()),
300                        _ => None,
301                    };
302
303                    if let Some(format_str) = format_str_opt {
304                        match format_str.as_str() {
305                            "email" => zod_expr = format!("{}.email()", zod_expr),
306                            "uri" | "url" => zod_expr = format!("{}.url()", zod_expr),
307                            "uuid" => zod_expr = format!("{}.uuid()", zod_expr),
308                            "date-time" | "datetime" | "date_time" => {
309                                zod_expr = format!("{}.datetime()", zod_expr)
310                            }
311                            _ => {}
312                        }
313                    }
314
315                    Ok(zod_expr)
316                }
317                Type::Number(number_type) => {
318                    let mut zod_expr = format!("{}z.number()", indent_str);
319
320                    // Add number constraints
321                    if let Some(minimum) = number_type.minimum {
322                        zod_expr = format!("{}.min({})", zod_expr, minimum);
323                    }
324                    if let Some(maximum) = number_type.maximum {
325                        zod_expr = format!("{}.max({})", zod_expr, maximum);
326                    }
327                    if let Some(multiple_of) = number_type.multiple_of {
328                        zod_expr = format!("{}.multipleOf({})", zod_expr, multiple_of);
329                    }
330
331                    Ok(zod_expr)
332                }
333                Type::Integer(integer_type) => {
334                    let mut zod_expr = format!("{}z.number().int()", indent_str);
335
336                    // Add integer constraints
337                    if let Some(minimum) = integer_type.minimum {
338                        zod_expr = format!("{}.min({})", zod_expr, minimum);
339                    }
340                    if let Some(maximum) = integer_type.maximum {
341                        zod_expr = format!("{}.max({})", zod_expr, maximum);
342                    }
343                    if let Some(multiple_of) = integer_type.multiple_of {
344                        zod_expr = format!("{}.multipleOf({})", zod_expr, multiple_of);
345                    }
346
347                    Ok(zod_expr)
348                }
349                Type::Boolean(_) => Ok(format!("{}z.boolean()", indent_str)),
350                Type::Array(array_type) => {
351                    let item_zod = if let Some(items) = &array_type.items {
352                        match items {
353                            ReferenceOr::Reference { reference } => {
354                                if let Some(ref_name) = get_schema_name_from_ref(reference) {
355                                    // If already processed, use lazy reference to avoid infinite recursion
356                                    if processed.contains(&ref_name) {
357                                        format!(
358                                            "{}z.lazy(() => {}Schema)",
359                                            indent_str,
360                                            to_pascal_case(&ref_name)
361                                        )
362                                    } else {
363                                        let resolved =
364                                            resolve_ref(openapi, reference).map_err(|e| {
365                                                crate::error::SchemaError::InvalidReference {
366                                                    ref_path: format!("Failed to resolve: {}", e),
367                                                }
368                                            })?;
369                                        if let ReferenceOr::Item(item_schema) = resolved {
370                                            // Check if it's an object that needs to be extracted
371                                            if matches!(
372                                                &item_schema.schema_kind,
373                                                SchemaKind::Type(Type::Object(_))
374                                            ) {
375                                                generate_zod_for_schema(
376                                                    openapi,
377                                                    &ref_name,
378                                                    &item_schema,
379                                                    zod_schemas,
380                                                    processed,
381                                                    enum_registry,
382                                                    None,
383                                                )?;
384                                                format!(
385                                                    "{}z.lazy(() => {}Schema)",
386                                                    indent_str,
387                                                    to_pascal_case(&ref_name)
388                                                )
389                                            } else {
390                                                schema_to_zod(
391                                                    openapi,
392                                                    &item_schema,
393                                                    zod_schemas,
394                                                    processed,
395                                                    indent,
396                                                    enum_registry,
397                                                    None,
398                                                    current_schema_name,
399                                                )?
400                                            }
401                                        } else {
402                                            format!(
403                                                "{}z.lazy(() => {}Schema)",
404                                                indent_str,
405                                                to_pascal_case(&ref_name)
406                                            )
407                                        }
408                                    }
409                                } else {
410                                    format!("{}z.any()", indent_str)
411                                }
412                            }
413                            ReferenceOr::Item(item_schema) => {
414                                // If it's an object, we need to generate it inline or as a separate schema
415                                if matches!(
416                                    &item_schema.schema_kind,
417                                    SchemaKind::Type(Type::Object(_))
418                                ) {
419                                    // Generate object fields
420                                    let object_fields = schema_to_zod(
421                                        openapi,
422                                        item_schema,
423                                        zod_schemas,
424                                        processed,
425                                        indent + 1,
426                                        enum_registry,
427                                        None,
428                                        current_schema_name,
429                                    )?;
430                                    // Wrap in z.object()
431                                    format!(
432                                        "{}z.object({{\n{}\n{}}})",
433                                        indent_str, object_fields, indent_str
434                                    )
435                                } else {
436                                    schema_to_zod(
437                                        openapi,
438                                        item_schema,
439                                        zod_schemas,
440                                        processed,
441                                        indent,
442                                        enum_registry,
443                                        None,
444                                        current_schema_name,
445                                    )?
446                                }
447                            }
448                        }
449                    } else {
450                        format!("{}z.any()", indent_str)
451                    };
452
453                    let mut array_zod = format!("{}z.array({})", indent_str, item_zod.trim_start());
454
455                    // Add array constraints
456                    if let Some(min_items) = array_type.min_items {
457                        array_zod = format!("{}.min({})", array_zod, min_items);
458                    }
459                    if let Some(max_items) = array_type.max_items {
460                        array_zod = format!("{}.max({})", array_zod, max_items);
461                    }
462
463                    Ok(array_zod)
464                }
465                Type::Object(object_type) => {
466                    if !object_type.properties.is_empty() {
467                        let mut fields = Vec::new();
468                        // Get parent schema name from context if available, otherwise use current_schema_name parameter
469                        // For object properties, use the current schema name as parent
470                        let parent_schema_for_props = if let Some((_, parent)) = context {
471                            if !parent.is_empty() {
472                                parent.to_string()
473                            } else if let Some(current) = current_schema_name {
474                                current.to_string()
475                            } else {
476                                String::new()
477                            }
478                        } else if let Some(current) = current_schema_name {
479                            current.to_string()
480                        } else {
481                            String::new()
482                        };
483
484                        for (prop_name, prop_schema_ref) in object_type.properties.iter() {
485                            let prop_zod = match prop_schema_ref {
486                                ReferenceOr::Reference { reference } => {
487                                    // For $ref properties, always use the schema name (don't inline)
488                                    if let Some(ref_name) = get_schema_name_from_ref(reference) {
489                                        // Generate the referenced schema if not already processed
490                                        if !processed.contains(&ref_name) {
491                                            if let Ok(ReferenceOr::Item(ref_schema)) =
492                                                resolve_ref(openapi, reference)
493                                            {
494                                                if matches!(
495                                                    &ref_schema.schema_kind,
496                                                    SchemaKind::Type(Type::Object(_))
497                                                ) {
498                                                    generate_zod_for_schema(
499                                                        openapi,
500                                                        &ref_name,
501                                                        &ref_schema,
502                                                        zod_schemas,
503                                                        processed,
504                                                        enum_registry,
505                                                        Some(&parent_schema_for_props),
506                                                    )?;
507                                                }
508                                            }
509                                        }
510                                        format!(
511                                            "{}z.lazy(() => {}Schema)",
512                                            indent_str,
513                                            to_pascal_case(&ref_name)
514                                        )
515                                    } else {
516                                        format!("{}z.any()", indent_str)
517                                    }
518                                }
519                                ReferenceOr::Item(prop_schema) => schema_to_zod(
520                                    openapi,
521                                    prop_schema,
522                                    zod_schemas,
523                                    processed,
524                                    indent + 1,
525                                    enum_registry,
526                                    Some((prop_name, &parent_schema_for_props)),
527                                    current_schema_name,
528                                )?,
529                            };
530
531                            let required = object_type.required.contains(prop_name);
532
533                            let nullable = prop_schema_ref
534                                .as_item()
535                                .map(|s| s.schema_data.nullable)
536                                .unwrap_or(false);
537
538                            let mut zod_expr = prop_zod.trim_start().to_string();
539                            if nullable {
540                                zod_expr = format!("{}.nullable()", zod_expr);
541                            }
542                            if !required {
543                                zod_expr = format!("{}.optional()", zod_expr);
544                            }
545
546                            fields.push(format!(
547                                "{}{}: {},",
548                                "  ".repeat(indent + 1),
549                                prop_name,
550                                zod_expr
551                            ));
552                        }
553                        Ok(fields.join("\n"))
554                    } else {
555                        Ok(format!("{}z.record(z.string(), z.any())", indent_str))
556                    }
557                }
558            }
559        }
560        SchemaKind::Any(_) => Ok(format!("{}z.any()", indent_str)),
561        SchemaKind::OneOf { one_of, .. } => {
562            let mut variant_schemas = Vec::new();
563            for item in one_of {
564                match item {
565                    ReferenceOr::Reference { reference } => {
566                        if let Some(ref_name) = get_schema_name_from_ref(reference) {
567                            variant_schemas.push(format!(
568                                "{}z.lazy(() => {}Schema)",
569                                indent_str,
570                                crate::generator::utils::to_pascal_case(&ref_name)
571                            ));
572                        } else {
573                            variant_schemas.push(format!("{}z.any()", indent_str));
574                        }
575                    }
576                    ReferenceOr::Item(item_schema) => {
577                        let item_zod = schema_to_zod(
578                            openapi,
579                            item_schema,
580                            zod_schemas,
581                            processed,
582                            indent,
583                            enum_registry,
584                            None,
585                            current_schema_name,
586                        )?;
587                        variant_schemas.push(item_zod);
588                    }
589                }
590            }
591            if variant_schemas.is_empty() {
592                Ok(format!("{}z.any()", indent_str))
593            } else {
594                Ok(format!(
595                    "{}z.union([{}])",
596                    indent_str,
597                    variant_schemas.join(", ")
598                ))
599            }
600        }
601        SchemaKind::AllOf { all_of, .. } => {
602            let mut all_schemas = Vec::new();
603            for item in all_of {
604                match item {
605                    ReferenceOr::Reference { reference } => {
606                        if let Some(ref_name) = get_schema_name_from_ref(reference) {
607                            all_schemas.push(format!(
608                                "{}z.lazy(() => {}Schema)",
609                                indent_str,
610                                crate::generator::utils::to_pascal_case(&ref_name)
611                            ));
612                        } else {
613                            all_schemas.push(format!("{}z.any()", indent_str));
614                        }
615                    }
616                    ReferenceOr::Item(item_schema) => {
617                        let item_zod = schema_to_zod(
618                            openapi,
619                            item_schema,
620                            zod_schemas,
621                            processed,
622                            indent,
623                            enum_registry,
624                            None,
625                            current_schema_name,
626                        )?;
627                        all_schemas.push(item_zod);
628                    }
629                }
630            }
631            if all_schemas.is_empty() {
632                Ok(format!("{}z.any()", indent_str))
633            } else if all_schemas.len() == 1 {
634                Ok(all_schemas[0].clone())
635            } else {
636                // AllOf represents intersection: all schemas must be satisfied
637                // Zod uses .and() for intersection. Chain them: schema1.and(schema2).and(schema3)
638                let mut result = all_schemas[0].clone();
639                for schema in all_schemas.iter().skip(1) {
640                    result = format!("{}.and({})", result.trim(), schema.trim());
641                }
642                Ok(result)
643            }
644        }
645        SchemaKind::AnyOf { any_of, .. } => {
646            // AnyOf is treated same as OneOf (union)
647            let mut variant_schemas = Vec::new();
648            for item in any_of {
649                match item {
650                    ReferenceOr::Reference { reference } => {
651                        if let Some(ref_name) = get_schema_name_from_ref(reference) {
652                            variant_schemas.push(format!(
653                                "{}z.lazy(() => {}Schema)",
654                                indent_str,
655                                crate::generator::utils::to_pascal_case(&ref_name)
656                            ));
657                        } else {
658                            variant_schemas.push(format!("{}z.any()", indent_str));
659                        }
660                    }
661                    ReferenceOr::Item(item_schema) => {
662                        let item_zod = schema_to_zod(
663                            openapi,
664                            item_schema,
665                            zod_schemas,
666                            processed,
667                            indent,
668                            enum_registry,
669                            None,
670                            current_schema_name,
671                        )?;
672                        variant_schemas.push(item_zod);
673                    }
674                }
675            }
676            if variant_schemas.is_empty() {
677                Ok(format!("{}z.any()", indent_str))
678            } else {
679                Ok(format!(
680                    "{}z.union([{}])",
681                    indent_str,
682                    variant_schemas.join(", ")
683                ))
684            }
685        }
686        SchemaKind::Not { .. } => Ok(format!("{}z.any()", indent_str)),
687    }
688}
689
690fn generate_enum_zod(name: &str, values: &[String]) -> ZodSchema {
691    let enum_name = to_pascal_case(name);
692    let enum_values = values
693        .iter()
694        .map(|v| format!("\"{}\"", v))
695        .collect::<Vec<_>>()
696        .join(", ");
697
698    ZodSchema {
699        content: format!(
700            "export const {}Schema = z.enum([{}]);",
701            enum_name, enum_values
702        ),
703    }
704}