convex-typegen 0.2.0

use std::io::{Seek, SeekFrom, Write};

use serde_json::Value as JsonValue;

use crate::convex::{ConvexFunction, ConvexFunctions, ConvexSchema, ConvexTable};
use crate::errors::ConvexTypeGeneratorError;

pub(crate) fn generate_code(path: &str, data: (ConvexSchema, ConvexFunctions)) -> Result<(), ConvexTypeGeneratorError>
{
    let mut file = std::fs::File::create(path)?;

    // Clear the file
    file.set_len(0)?;
    file.seek(SeekFrom::Start(0))?;

    let file_header = r#"// This file is generated by convex-typegen. Do not modify directly.
// You can find more information about convex-typegen at https://github.com/JamalLyons/convex-typegen

#![allow(non_camel_case_types)]
#![allow(non_snake_case)]

use serde::{Serialize, Deserialize};

"#;

    file.write_all(file_header.as_bytes())?;

    // A buffer to hold the generated code
    let mut code = String::new();

    // First generate all enums from the tables
    for table in &data.0.tables {
        code.push_str(&generate_table_enums(table));
    }

    // Then generate the table structs
    for table in data.0.tables {
        code.push_str(&generate_table_code(table));
    }

    // Generate function argument types
    for function in data.1 {
        code.push_str(&generate_function_code(function));
    }

    file.write_all(code.as_bytes())?;

    Ok(())
}

/// Generate enums for a table's union types
fn generate_table_enums(table: &ConvexTable) -> String
{
    let mut code = String::new();

    for column in &table.columns {
        // Handle regular unions
        if let Some("union") = column.data_type["type"].as_str() {
            let enum_name = format!(
                "{}{}",
                capitalize_first_letter(&table.name),
                capitalize_first_letter(&column.name)
            );
            code.push_str("#[derive(Debug, Clone)]\n");
            code.push_str(&format!("pub enum {} {{\n", enum_name));

            if let Some(variants) = column.data_type["variants"].as_array() {
                for variant in variants {
                    match variant["type"].as_str() {
                        Some("literal") => {
                            if let Some(value) = variant["value"]["value"].as_str() {
                                code.push_str(&format!("    {},\n", to_pascal_case(value)));
                            }
                        }
                        Some(type_name) => {
                            let rust_type = convex_type_to_rust_type(variant, Some(&table.name), Some(&column.name));
                            code.push_str(&format!("    {}({}),\n", to_pascal_case(type_name), rust_type));
                        }
                        None => continue,
                    }
                }
            }

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

        // Handle optional unions
        if let Some("optional") = column.data_type["type"].as_str() {
            if let Some("union") = column.data_type["inner"]["type"].as_str() {
                let enum_name = format!(
                    "{}Optional{}",
                    capitalize_first_letter(&table.name),
                    capitalize_first_letter(&column.name)
                );
                code.push_str("#[derive(Debug, Clone)]\n");
                code.push_str(&format!("pub enum {} {{\n", enum_name));

                if let Some(variants) = column.data_type["inner"]["variants"].as_array() {
                    for variant in variants {
                        match variant["type"].as_str() {
                            Some("literal") => {
                                if let Some(value) = variant["value"]["value"].as_str() {
                                    code.push_str(&format!("    {},\n", to_pascal_case(value)));
                                }
                            }
                            Some(type_name) => {
                                let rust_type = convex_type_to_rust_type(variant, Some(&table.name), Some(&column.name));
                                code.push_str(&format!("    {}({}),\n", to_pascal_case(type_name), rust_type));
                            }
                            None => continue,
                        }
                    }
                }

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

    code
}

/// Generate the code for a table.
fn generate_table_code(table: ConvexTable) -> String
{
    let mut code = String::new();

    let table_struct_name = format!("{}Table", capitalize_first_letter(&table.name));

    code.push_str("#[derive(Debug, Clone)]\n");
    code.push_str(&format!("pub struct {} {{\n", table_struct_name));

    // Generate fields for each column
    for column in table.columns {
        let rust_type = if column.data_type["type"].as_str() == Some("union") {
            format!(
                "{}{}",
                capitalize_first_letter(&table.name),
                capitalize_first_letter(&column.name)
            )
        } else {
            convex_type_to_rust_type(&column.data_type, Some(&table.name), Some(&column.name))
        };
        code.push_str(&format!("    pub {}: {},\n", column.name, rust_type));
    }

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

/// Convert a Convex type to its corresponding Rust type
fn convex_type_to_rust_type(data_type: &JsonValue, table_name: Option<&str>, field_name: Option<&str>) -> String
{
    // Get the base type from the "type" field
    let type_str = data_type["type"].as_str().unwrap_or("unknown");

    match type_str {
        "string" => "String".to_string(),
        "number" => "f64".to_string(),
        "boolean" => "bool".to_string(),
        "null" => "()".to_string(),
        "int64" => "i64".to_string(),
        "bytes" => "Vec<u8>".to_string(),
        "any" => "serde_json::Value".to_string(),

        "array" => {
            let element_type = convex_type_to_rust_type(&data_type["elements"], None, None);
            format!("Vec<{}>", element_type)
        }

        "object" => {
            if let Some(props) = data_type["properties"].as_object() {
                let value_type = props
                    .values()
                    .next()
                    .map(|v| convex_type_to_rust_type(v, None, None))
                    .unwrap_or_else(|| "serde_json::Value".to_string());
                format!("std::collections::BTreeMap<String, {}>", value_type)
            } else {
                "serde_json::Value".to_string()
            }
        }

        "record" => {
            let key_type = convex_type_to_rust_type(&data_type["keyType"], None, None);
            let value_type = convex_type_to_rust_type(&data_type["valueType"], None, None);
            format!("std::collections::HashMap<{}, {}>", key_type, value_type)
        }

        "optional" => {
            let inner_type = match data_type["inner"]["type"].as_str() {
                Some("union") => {
                    if let (Some(table), Some(field)) = (table_name, field_name) {
                        format!("{}Optional{}", capitalize_first_letter(table), capitalize_first_letter(field))
                    } else {
                        "serde_json::Value".to_string()
                    }
                }
                _ => convex_type_to_rust_type(&data_type["inner"], None, None),
            };
            format!("Option<{}>", inner_type)
        }

        "literal" => {
            // Handle literal types
            if let Some(value) = data_type["value"]["value"].as_str() {
                format!("\"{}\"", value)
            } else {
                "String".to_string()
            }
        }

        "id" => "String".to_string(),

        _ => "serde_json::Value".to_string(), // fallback for unknown types
    }
}

/// Generate the code for a function.
fn generate_function_code(function: ConvexFunction) -> String
{
    let mut code = String::new();

    // Generate the args struct name
    let struct_name = format!("{}Args", capitalize_first_letter(&function.name));

    // Generate struct with derive macros
    code.push_str("#[derive(Debug, Clone, Serialize, Deserialize)]\n");
    code.push_str(&format!("pub struct {} {{\n", struct_name));

    // Generate fields for each parameter
    for param in &function.params {
        let rust_type = convex_type_to_rust_type(&param.data_type, None, None);
        code.push_str(&format!("    pub {}: {},\n", param.name, rust_type));
    }

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

    // Add implementation block with static FUNCTION_PATH method
    code.push_str(&format!("impl {} {{\n", struct_name));
    code.push_str("    pub const FUNCTION_PATH: &'static str = ");
    code.push_str(&format!("\"{}:{}\";\n", function.file_name, function.name));
    code.push_str("}\n\n");

    // Generate From implementation to convert to BTreeMap
    code.push_str(&format!(
        "impl From<{}> for std::collections::BTreeMap<String, serde_json::Value> {{\n",
        struct_name
    ));
    code.push_str(&format!("    fn from(_args: {}) -> Self {{\n", struct_name));

    // Only create map and insert values if there are parameters
    if function.params.is_empty() {
        code.push_str("        std::collections::BTreeMap::new()\n");
    } else {
        code.push_str("        let mut map = std::collections::BTreeMap::new();\n");
        // Convert each field to a serde_json::Value and insert into map
        for param in &function.params {
            code.push_str(&format!(
                "        map.insert(\"{}\".to_string(), serde_json::to_value(_args.{}).unwrap());\n",
                param.name, param.name
            ));
        }
        code.push_str("        map\n");
    }

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

    code
}

/// Capitalize the first letter of a string
fn capitalize_first_letter(s: &str) -> String
{
    // If the string is empty, return an empty string
    if s.is_empty() {
        return String::new();
    }

    // Split the string into the first character and the rest of the string
    let mut chars = s.chars();
    let first_char = chars.next().expect("Expected a character but got none");
    let rest = chars.collect::<String>();

    // Capitalize the first character and concatenate with the rest of the string
    first_char.to_uppercase().to_string() + &rest
}

fn to_pascal_case(s: &str) -> String
{
    s.split(|c: char| !c.is_alphanumeric())
        .filter(|s| !s.is_empty())
        .map(|word| {
            let mut chars = word.chars();
            match chars.next() {
                None => String::new(),
                Some(first) => first.to_uppercase().collect::<String>() + &chars.collect::<String>().to_lowercase(),
            }
        })
        .collect()
}