forge_codegen/

parser.rs

//! Rust source code parser for extracting FORGE schema definitions.
//!
//! This module parses Rust source files to extract model, enum, and function
//! definitions without requiring compilation.

use std::path::{Path, PathBuf};

use forge_core::schema::{
    EnumDef, EnumVariant, FieldDef, FunctionArg, FunctionDef, FunctionKind, RustType,
    SchemaRegistry, TableDef,
};
use forge_core::util::to_snake_case;
use quote::ToTokens;
use syn::{Attribute, Expr, Fields, FnArg, Lit, Meta, Pat, ReturnType};

use crate::Error;

fn collect_rs_files(dir: &Path, out: &mut Vec<PathBuf>) {
    let entries = match std::fs::read_dir(dir) {
        Ok(e) => e,
        Err(_) => return,
    };
    for entry in entries.flatten() {
        let path = entry.path();
        if path.is_dir() {
            collect_rs_files(&path, out);
        } else if path.extension().is_some_and(|ext| ext == "rs") {
            out.push(path);
        }
    }
}

/// Parse all Rust source files in a directory and extract schema definitions.
pub fn parse_project(src_dir: &Path) -> Result<SchemaRegistry, Error> {
    let registry = SchemaRegistry::new();

    let mut files = Vec::new();
    collect_rs_files(src_dir, &mut files);
    files.sort();

    for path in &files {
        let content = std::fs::read_to_string(path)?;
        if let Err(e) = parse_file(&content, &registry) {
            tracing::debug!(file = ?path, error = %e, "Failed to parse file");
        }
    }

    Ok(registry)
}

/// Parse a single Rust source file and extract schema definitions.
fn parse_file(content: &str, registry: &SchemaRegistry) -> Result<(), Error> {
    let file = syn::parse_file(content).map_err(|e| Error::Template(e.to_string()))?;

    for item in file.items {
        match item {
            syn::Item::Struct(item_struct) => {
                if has_forge_model_attr(&item_struct.attrs) {
                    if let Some(table) = parse_model(&item_struct) {
                        registry.register_table(table);
                    }
                } else if has_serde_derive(&item_struct.attrs) {
                    // Parse DTO structs (those with Serialize/Deserialize)
                    if let Some(table) = parse_dto_struct(&item_struct) {
                        registry.register_table(table);
                    }
                }
            }
            syn::Item::Enum(item_enum) => {
                if has_forge_enum_attr(&item_enum.attrs) {
                    if let Some(enum_def) = parse_enum(&item_enum) {
                        registry.register_enum(enum_def);
                    }
                } else if has_serde_derive(&item_enum.attrs) {
                    // Parse enums with Serialize/Deserialize
                    if let Some(enum_def) = parse_enum(&item_enum) {
                        registry.register_enum(enum_def);
                    }
                }
            }
            syn::Item::Fn(item_fn) => {
                if let Some(func) = parse_function(&item_fn) {
                    registry.register_function(func);
                }
            }
            _ => {}
        }
    }

    Ok(())
}

/// Check if attributes contain #[forge::model] or #[model].
fn has_forge_model_attr(attrs: &[Attribute]) -> bool {
    attrs.iter().any(|attr| {
        let path = attr.path();
        path.is_ident("model")
            || path.segments.len() == 2
                && path.segments[0].ident == "forge"
                && path.segments[1].ident == "model"
    })
}

/// Check if attributes contain #[forge_enum] or #[forge::enum_type].
fn has_forge_enum_attr(attrs: &[Attribute]) -> bool {
    attrs.iter().any(|attr| {
        let path = attr.path();
        path.is_ident("forge_enum")
            || path.is_ident("enum_type")
            || path.segments.len() == 2
                && path.segments[0].ident == "forge"
                && (path.segments[1].ident == "enum_type" || path.segments[1].ident == "forge_enum")
    })
}

/// Check if attributes contain #[derive(...Serialize...)] or #[derive(...Deserialize...)].
fn has_serde_derive(attrs: &[Attribute]) -> bool {
    attrs.iter().any(|attr| {
        if !attr.path().is_ident("derive") {
            return false;
        }
        let tokens = attr.meta.to_token_stream().to_string();
        tokens.contains("Serialize") || tokens.contains("Deserialize")
    })
}

/// Parse a DTO struct (with Serialize/Deserialize) into a TableDef.
fn parse_dto_struct(item: &syn::ItemStruct) -> Option<TableDef> {
    let struct_name = item.ident.to_string();

    // Use struct name as table name (DTOs don't have SQL tables)
    let mut table = TableDef::new(&struct_name, &struct_name);

    // Mark as DTO (not a database table)
    table.is_dto = true;

    // Extract documentation
    table.doc = get_doc_comment(&item.attrs);

    // Extract fields
    if let Fields::Named(fields) = &item.fields {
        for field in &fields.named {
            if let Some(field_name) = &field.ident {
                let field_def = parse_field(field_name.to_string(), &field.ty, &field.attrs);
                table.fields.push(field_def);
            }
        }
    }

    Some(table)
}

/// Parse a struct with #[model] attribute into a TableDef.
fn parse_model(item: &syn::ItemStruct) -> Option<TableDef> {
    let struct_name = item.ident.to_string();
    let table_name = get_table_name_from_attrs(&item.attrs).unwrap_or_else(|| {
        let snake = to_snake_case(&struct_name);
        pluralize(&snake)
    });

    let mut table = TableDef::new(&table_name, &struct_name);

    // Extract documentation
    table.doc = get_doc_comment(&item.attrs);

    // Extract fields
    if let Fields::Named(fields) = &item.fields {
        for field in &fields.named {
            if let Some(field_name) = &field.ident {
                let field_def = parse_field(field_name.to_string(), &field.ty, &field.attrs);
                table.fields.push(field_def);
            }
        }
    }

    Some(table)
}

/// Parse a field definition.
fn parse_field(name: String, ty: &syn::Type, attrs: &[Attribute]) -> FieldDef {
    let rust_type = type_to_rust_type(ty);
    let mut field = FieldDef::new(&name, rust_type);
    field.column_name = to_snake_case(&name);
    field.doc = get_doc_comment(attrs);
    field
}

/// Parse an enum with #[forge_enum] attribute into an EnumDef.
fn parse_enum(item: &syn::ItemEnum) -> Option<EnumDef> {
    let enum_name = item.ident.to_string();
    let mut enum_def = EnumDef::new(&enum_name);
    enum_def.doc = get_doc_comment(&item.attrs);

    for variant in &item.variants {
        let variant_name = variant.ident.to_string();
        let mut enum_variant = EnumVariant::new(&variant_name);
        enum_variant.doc = get_doc_comment(&variant.attrs);

        // Check for explicit value
        if let Some((_, Expr::Lit(lit))) = &variant.discriminant
            && let Lit::Int(int_lit) = &lit.lit
            && let Ok(value) = int_lit.base10_parse::<i32>()
        {
            enum_variant.int_value = Some(value);
        }

        enum_def.variants.push(enum_variant);
    }

    Some(enum_def)
}

/// Parse a function with #[query] or #[mutation] attribute.
fn parse_function(item: &syn::ItemFn) -> Option<FunctionDef> {
    let kind = get_function_kind(&item.attrs)?;
    let func_name = item.sig.ident.to_string();

    // Get return type
    let return_type = match &item.sig.output {
        ReturnType::Default => RustType::Custom("()".to_string()),
        ReturnType::Type(_, ty) => extract_result_type(ty),
    };

    let mut func = FunctionDef::new(&func_name, kind, return_type);
    func.doc = get_doc_comment(&item.attrs);
    func.is_async = item.sig.asyncness.is_some();

    // Parse arguments (skip first arg which is usually context)
    let mut skip_first = true;
    for arg in &item.sig.inputs {
        if let FnArg::Typed(pat_type) = arg {
            // Skip context argument (usually first)
            if skip_first {
                skip_first = false;
                // Check if it's a context type
                let type_str = quote::quote!(#pat_type.ty).to_string();
                if type_str.contains("Context")
                    || type_str.contains("QueryContext")
                    || type_str.contains("MutationContext")
                {
                    continue;
                }
            }

            // Extract argument name
            if let Pat::Ident(pat_ident) = &*pat_type.pat {
                let arg_name = pat_ident.ident.to_string();
                let arg_type = type_to_rust_type(&pat_type.ty);
                func.args.push(FunctionArg::new(arg_name, arg_type));
            }
        }
    }

    Some(func)
}

/// Get the function kind from attributes.
fn get_function_kind(attrs: &[Attribute]) -> Option<FunctionKind> {
    for attr in attrs {
        let path = attr.path();
        let segments: Vec<_> = path.segments.iter().map(|s| s.ident.to_string()).collect();

        // Check for #[forge::X] or #[X] patterns
        let kind_str = if segments.len() == 2 && segments[0] == "forge" {
            Some(segments[1].as_str())
        } else if segments.len() == 1 {
            Some(segments[0].as_str())
        } else {
            None
        };

        if let Some(kind) = kind_str {
            match kind {
                "query" => return Some(FunctionKind::Query),
                "mutation" => return Some(FunctionKind::Mutation),
                "job" => return Some(FunctionKind::Job),
                "cron" => return Some(FunctionKind::Cron),
                "workflow" => return Some(FunctionKind::Workflow),
                _ => {}
            }
        }
    }
    None
}

/// Extract the inner type from Result<T, E>.
fn extract_result_type(ty: &syn::Type) -> RustType {
    let type_str = quote::quote!(#ty).to_string().replace(' ', "");

    // Check for Result<T, _>
    if let Some(rest) = type_str.strip_prefix("Result<") {
        // Find the inner type before the comma or angle bracket
        let mut depth = 0;
        let mut end_idx = 0;
        for (i, c) in rest.chars().enumerate() {
            match c {
                '<' => depth += 1,
                '>' => {
                    if depth == 0 {
                        end_idx = i;
                        break;
                    }
                    depth -= 1;
                }
                ',' if depth == 0 => {
                    end_idx = i;
                    break;
                }
                _ => {}
            }
        }
        let inner = &rest[..end_idx];
        return type_to_rust_type(
            &syn::parse_str(inner)
                .unwrap_or_else(|_| syn::parse_str::<syn::Type>("String").unwrap()),
        );
    }

    type_to_rust_type(ty)
}

/// Convert a syn::Type to RustType.
fn type_to_rust_type(ty: &syn::Type) -> RustType {
    let type_str = quote::quote!(#ty).to_string().replace(' ', "");

    // Handle common types
    match type_str.as_str() {
        "String" | "&str" => RustType::String,
        "i32" => RustType::I32,
        "i64" => RustType::I64,
        "f32" => RustType::F32,
        "f64" => RustType::F64,
        "bool" => RustType::Bool,
        "Uuid" | "uuid::Uuid" => RustType::Uuid,
        "DateTime<Utc>" | "chrono::DateTime<Utc>" | "chrono::DateTime<chrono::Utc>" => {
            RustType::DateTime
        }
        "NaiveDate" | "chrono::NaiveDate" => RustType::Date,
        "NaiveTime" | "chrono::NaiveTime" => RustType::Custom("NaiveTime".to_string()),
        "serde_json::Value" | "Value" => RustType::Json,
        "Vec<u8>" => RustType::Bytes,
        _ => {
            // Handle Option<T> by recursively resolving the inner type
            if let Some(inner) = type_str
                .strip_prefix("Option<")
                .and_then(|s| s.strip_suffix('>'))
            {
                let inner_ty: syn::Type =
                    syn::parse_str(inner).unwrap_or_else(|_| syn::parse_str("String").unwrap());
                return RustType::Option(Box::new(type_to_rust_type(&inner_ty)));
            }

            // Handle Vec<T> by recursively resolving the inner type
            if let Some(inner) = type_str
                .strip_prefix("Vec<")
                .and_then(|s| s.strip_suffix('>'))
            {
                if inner == "u8" {
                    return RustType::Bytes;
                }
                let inner_ty: syn::Type =
                    syn::parse_str(inner).unwrap_or_else(|_| syn::parse_str("String").unwrap());
                return RustType::Vec(Box::new(type_to_rust_type(&inner_ty)));
            }

            // Default to custom type
            RustType::Custom(type_str)
        }
    }
}

/// Get #[table(name = "...")] value from attributes.
fn get_table_name_from_attrs(attrs: &[Attribute]) -> Option<String> {
    for attr in attrs {
        if attr.path().is_ident("table")
            && let Meta::List(list) = &attr.meta
        {
            let tokens = list.tokens.to_string();
            if let Some(value) = extract_name_value(&tokens) {
                return Some(value);
            }
        }
    }
    None
}

/// Get string value from attribute like #[attr = "value"].
fn get_attribute_string_value(attr: &Attribute) -> Option<String> {
    if let Meta::NameValue(nv) = &attr.meta
        && let Expr::Lit(lit) = &nv.value
        && let Lit::Str(s) = &lit.lit
    {
        return Some(s.value());
    }
    None
}

/// Get documentation comment from attributes.
fn get_doc_comment(attrs: &[Attribute]) -> Option<String> {
    let docs: Vec<String> = attrs
        .iter()
        .filter_map(|attr| {
            if attr.path().is_ident("doc") {
                get_attribute_string_value(attr)
            } else {
                None
            }
        })
        .collect();

    if docs.is_empty() {
        None
    } else {
        Some(
            docs.into_iter()
                .map(|s| s.trim().to_string())
                .collect::<Vec<_>>()
                .join("\n"),
        )
    }
}

/// Extract name value from "name = \"value\"" format.
fn extract_name_value(s: &str) -> Option<String> {
    let parts: Vec<&str> = s.splitn(2, '=').collect();
    if parts.len() == 2 {
        let value = parts[1].trim();
        if let Some(stripped) = value.strip_prefix('"').and_then(|s| s.strip_suffix('"')) {
            return Some(stripped.to_string());
        }
    }
    None
}

/// Simple English pluralization.
fn pluralize(s: &str) -> String {
    if s.ends_with('s')
        || s.ends_with("sh")
        || s.ends_with("ch")
        || s.ends_with('x')
        || s.ends_with('z')
    {
        format!("{}es", s)
    } else if let Some(stem) = s.strip_suffix('y') {
        if !s.ends_with("ay") && !s.ends_with("ey") && !s.ends_with("oy") && !s.ends_with("uy") {
            format!("{}ies", stem)
        } else {
            format!("{}s", s)
        }
    } else {
        format!("{}s", s)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_model_source() {
        let source = r#"
            #[model]
            struct User {
                #[id]
                id: Uuid,
                email: String,
                name: Option<String>,
                #[indexed]
                created_at: DateTime<Utc>,
            }
        "#;

        let registry = SchemaRegistry::new();
        parse_file(source, &registry).unwrap();

        let table = registry.get_table("users").unwrap();
        assert_eq!(table.struct_name, "User");
        assert_eq!(table.fields.len(), 4);
    }

    #[test]
    fn test_parse_enum_source() {
        let source = r#"
            #[forge_enum]
            enum ProjectStatus {
                Draft,
                Active,
                Completed,
            }
        "#;

        let registry = SchemaRegistry::new();
        parse_file(source, &registry).unwrap();

        let enum_def = registry.get_enum("ProjectStatus").unwrap();
        assert_eq!(enum_def.variants.len(), 3);
    }

    #[test]
    fn test_to_snake_case() {
        assert_eq!(to_snake_case("UserProfile"), "user_profile");
        assert_eq!(to_snake_case("ID"), "i_d");
        assert_eq!(to_snake_case("createdAt"), "created_at");
    }

    #[test]
    fn test_pluralize() {
        assert_eq!(pluralize("user"), "users");
        assert_eq!(pluralize("category"), "categories");
        assert_eq!(pluralize("box"), "boxes");
        assert_eq!(pluralize("address"), "addresses");
    }

    #[test]
    fn test_parse_query_function() {
        let source = r#"
            #[query]
            async fn get_user(ctx: QueryContext, id: Uuid) -> Result<User> {
                todo!()
            }
        "#;

        let registry = SchemaRegistry::new();
        parse_file(source, &registry).unwrap();

        let func = registry.get_function("get_user").unwrap();
        assert_eq!(func.name, "get_user");
        assert_eq!(func.kind, FunctionKind::Query);
        assert!(func.is_async);
    }

    #[test]
    fn test_parse_mutation_function() {
        let source = r#"
            #[mutation]
            async fn create_user(ctx: MutationContext, name: String, email: String) -> Result<User> {
                todo!()
            }
        "#;

        let registry = SchemaRegistry::new();
        parse_file(source, &registry).unwrap();

        let func = registry.get_function("create_user").unwrap();
        assert_eq!(func.name, "create_user");
        assert_eq!(func.kind, FunctionKind::Mutation);
        assert_eq!(func.args.len(), 2);
        assert_eq!(func.args[0].name, "name");
        assert_eq!(func.args[1].name, "email");
    }
}