drizzle-migrations 0.1.11

//! `SQLite` schema code generation
//!
//! This module generates Rust source code from introspected DDL entities.
//! The generated code uses the lowercase attribute syntax (e.g., `primary` instead of `PRIMARY`)
//! that is the current recommended style.

use super::collection::SQLiteDDL;
use super::ddl::{Column, ForeignKey, Index, Table, View};
use crate::utils::escape_for_rust_literal;
use drizzle_types::sqlite::SQLTypeCategory;
use heck::{ToLowerCamelCase, ToPascalCase, ToSnakeCase};
use std::collections::{HashMap, HashSet};
use std::fmt::Write;

/// Result of code generation
#[derive(Debug, Clone, Default)]
pub struct GeneratedSchema {
    /// The generated Rust source code
    pub code: String,
    /// Tables that were generated
    pub tables: Vec<String>,
    /// Indexes that were generated
    pub indexes: Vec<String>,
    /// Views that were generated
    pub views: Vec<String>,
    /// Any warnings during generation
    pub warnings: Vec<String>,
}

/// Options for code generation
#[derive(Debug, Clone, Default)]
pub struct CodegenOptions {
    /// Module documentation
    pub module_doc: Option<String>,
    /// Whether to include a schema struct
    pub include_schema: bool,
    /// Schema struct name
    pub schema_name: String,
    /// Whether to use public visibility
    pub use_pub: bool,
    /// Field naming style for generated Rust members
    pub field_casing: FieldCasing,
}

/// Casing strategy for generated Rust field names.
#[derive(Debug, Clone, Copy, Default)]
pub enum FieldCasing {
    /// `snake_case` (default)
    #[default]
    Snake,
    /// `camelCase`
    Camel,
    /// Preserve source casing as much as possible
    Preserve,
}

fn sanitize_rust_identifier(name: &str) -> String {
    let mut out = String::with_capacity(name.len());
    for (idx, ch) in name.chars().enumerate() {
        let valid = if idx == 0 {
            ch == '_' || ch.is_ascii_alphabetic()
        } else {
            ch == '_' || ch.is_ascii_alphanumeric()
        };

        if valid {
            out.push(ch);
        } else {
            out.push('_');
        }
    }

    if out.is_empty() { "_".to_string() } else { out }
}

fn apply_field_casing(name: &str, casing: FieldCasing) -> String {
    match casing {
        FieldCasing::Snake => name.to_snake_case(),
        FieldCasing::Camel => name.to_lower_camel_case(),
        FieldCasing::Preserve => sanitize_rust_identifier(name),
    }
}

struct SchemaMaps<'a> {
    table_columns: HashMap<String, Vec<&'a Column>>,
    table_pks: HashMap<String, HashSet<String>>,
    table_uniques: HashMap<String, HashSet<String>>,
    fk_map: HashMap<(String, String), (&'a ForeignKey, usize)>,
}

fn build_schema_maps(ddl: &SQLiteDDL) -> SchemaMaps<'_> {
    let mut table_columns: HashMap<String, Vec<&Column>> = HashMap::new();
    for column in ddl.columns.list() {
        table_columns
            .entry(column.table.to_string())
            .or_default()
            .push(column);
    }

    let mut table_pks: HashMap<String, HashSet<String>> = HashMap::new();
    for pk in ddl.pks.list() {
        for col in pk.columns.iter() {
            table_pks
                .entry(pk.table.to_string())
                .or_default()
                .insert(col.to_string());
        }
    }

    let mut table_uniques: HashMap<String, HashSet<String>> = HashMap::new();
    for unique in ddl.uniques.list() {
        if unique.columns.len() == 1 {
            table_uniques
                .entry(unique.table.to_string())
                .or_default()
                .insert(unique.columns[0].to_string());
        }
    }

    let mut fk_map: HashMap<(String, String), (&ForeignKey, usize)> = HashMap::new();
    for fk in ddl.fks.list() {
        for (idx, col) in fk.columns.iter().enumerate() {
            fk_map.insert((fk.table.to_string(), col.to_string()), (fk, idx));
        }
    }

    SchemaMaps {
        table_columns,
        table_pks,
        table_uniques,
        fk_map,
    }
}

fn write_module_header(code: &mut String, options: &CodegenOptions) {
    code.push_str("//! Auto-generated SQLite schema from introspection\n");
    code.push_str("//!\n");
    if let Some(doc) = &options.module_doc {
        for line in doc.lines() {
            code.push_str("//! ");
            code.push_str(line);
            code.push('\n');
        }
    }
    code.push('\n');
    code.push_str("use drizzle::sqlite::prelude::*;\n\n");
}

/// Generate Rust schema code from DDL
#[must_use]
pub fn generate_rust_schema(ddl: &SQLiteDDL, options: &CodegenOptions) -> GeneratedSchema {
    let mut result = GeneratedSchema::default();
    let mut code = String::new();

    write_module_header(&mut code, options);

    let SchemaMaps {
        table_columns,
        table_pks,
        table_uniques,
        fk_map,
    } = build_schema_maps(ddl);

    // Generate table structs
    for table in ddl.tables.list() {
        let table_name = table.name.to_string();
        let columns = table_columns
            .get(&table_name)
            .map_or(&[][..], std::vec::Vec::as_slice);

        // Preserve DB/introspection order when available (cid -> ordinal_position).
        let mut columns_sorted: Vec<&Column> = columns.to_vec();
        columns_sorted.sort_by(|a, b| {
            let ao = a.ordinal_position.unwrap_or(i32::MAX);
            let bo = b.ordinal_position.unwrap_or(i32::MAX);
            ao.cmp(&bo).then_with(|| a.name.cmp(&b.name))
        });
        let pk_columns = table_pks.get(&table_name);
        let unique_columns = table_uniques.get(&table_name);
        let is_composite_pk = pk_columns.is_some_and(|pks| pks.len() > 1);

        let ctx = TableGenContext {
            table,
            columns: &columns_sorted,
            pk_columns,
            unique_columns,
            is_composite_pk,
            fk_map: &fk_map,
            use_pub: options.use_pub,
            field_casing: options.field_casing,
        };

        let table_code = generate_table_struct(&ctx);

        code.push_str(&table_code);
        code.push('\n');
        result.tables.push(table_name);
    }

    // Generate index structs
    for index in ddl.indexes.list() {
        let index_code = generate_index_struct(index, options.use_pub, options.field_casing);
        code.push_str(&index_code);
        code.push('\n');
        result.indexes.push(index.name.to_string());
    }

    // Generate view structs
    for view in ddl.views.list() {
        // Skip existing views (not managed by drizzle)
        if view.is_existing {
            continue;
        }
        let view_name = view.name.to_string();
        let columns = table_columns
            .get(&view_name)
            .map_or(&[][..], std::vec::Vec::as_slice);
        let view_code = generate_view_struct(view, columns, options.use_pub, options.field_casing);
        code.push_str(&view_code);
        code.push('\n');
        result.views.push(view_name);
    }

    // Generate schema struct if requested
    if options.include_schema {
        let schema_code = generate_schema_struct(
            &options.schema_name,
            &result.tables,
            &result.indexes,
            options.use_pub,
            options.field_casing,
        );
        code.push_str(&schema_code);
    }

    result.code = code;
    result
}

/// Generate a single table struct
struct TableGenContext<'a> {
    table: &'a Table,
    columns: &'a [&'a Column],
    pk_columns: Option<&'a HashSet<String>>,
    unique_columns: Option<&'a HashSet<String>>,
    is_composite_pk: bool,
    fk_map: &'a HashMap<(String, String), (&'a ForeignKey, usize)>,
    use_pub: bool,
    field_casing: FieldCasing,
}

/// Generate a single table struct
fn generate_table_struct(ctx: &TableGenContext<'_>) -> String {
    let mut code = String::new();
    let vis = if ctx.use_pub { "pub " } else { "" };

    // Struct name is PascalCase of table name
    let struct_name = ctx.table.name.to_pascal_case();

    // Check if table name differs from struct name
    let needs_name_attr = apply_field_casing(&struct_name, ctx.field_casing) != ctx.table.name;

    // Build table attribute options
    let mut table_attrs = Vec::new();
    if needs_name_attr {
        table_attrs.push(format!("name = \"{}\"", ctx.table.name));
    }
    if ctx.table.strict {
        table_attrs.push("strict".to_string());
    }
    if ctx.table.without_rowid {
        table_attrs.push("without_rowid".to_string());
    }

    // Table attribute
    if table_attrs.is_empty() {
        code.push_str("#[SQLiteTable]\n");
    } else {
        let _ = writeln!(code, "#[SQLiteTable({})]", table_attrs.join(", "));
    }

    // Struct definition
    let _ = writeln!(code, "{vis}struct {struct_name} {{");

    // Fields
    for column in ctx.columns {
        let field_code = generate_column_field(
            column,
            ctx.pk_columns,
            ctx.unique_columns,
            ctx.is_composite_pk,
            ctx.fk_map,
            ctx.use_pub,
            ctx.field_casing,
        );
        code.push_str(&field_code);
    }

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

/// Generate a single column as a struct field
fn generate_column_field(
    column: &Column,
    pk_columns: Option<&HashSet<String>>,
    unique_columns: Option<&HashSet<String>>,
    is_composite_pk: bool,
    fk_map: &HashMap<(String, String), (&ForeignKey, usize)>,
    use_pub: bool,
    field_casing: FieldCasing,
) -> String {
    let vis = if use_pub { "pub " } else { "" };

    // Determine column attributes
    let mut attrs = Vec::new();
    let column_name = column.name.to_string();

    // Check if primary key
    let is_pk = pk_columns.is_some_and(|pks| pks.contains(&column_name));

    // Only add primary if it's a single-column PK (not composite)
    if is_pk && !is_composite_pk {
        attrs.push("primary".to_string());
    }

    // Check autoincrement
    if column.autoincrement == Some(true) {
        attrs.push("autoincrement".to_string());
    }

    // Check unique (only for single-column constraints)
    let is_unique = unique_columns.is_some_and(|uniques| uniques.contains(&column_name));
    if is_unique {
        attrs.push("unique".to_string());
    }

    // Check default value
    if let Some(default) = &column.default {
        // Format default value for Rust
        let default_str = format_default_value(default, &column.sql_type);
        if let Some(d) = default_str {
            attrs.push(format!("default = {d}"));
        }
    }

    // Check foreign key
    if let Some((fk, idx)) = fk_map.get(&(column.table.to_string(), column_name))
        && let Some(ref_col) = fk.columns_to.get(*idx)
    {
        let ref_table_struct = fk.table_to.to_pascal_case();
        attrs.push(format!("references = {ref_table_struct}::{ref_col}"));

        // Add ON DELETE if specified
        if let Some(on_delete) = &fk.on_delete
            && !on_delete.eq_ignore_ascii_case("NO ACTION")
        {
            let action = on_delete.replace(' ', "_").to_lowercase();
            attrs.push(format!("on_delete = {action}"));
        }

        // Add ON UPDATE if specified
        if let Some(on_update) = &fk.on_update
            && !on_update.eq_ignore_ascii_case("NO ACTION")
        {
            let action = on_update.replace(' ', "_").to_lowercase();
            attrs.push(format!("on_update = {action}"));
        }
    }

    // Build the #[column(...)] attribute if there are any modifiers
    let attr_str = if attrs.is_empty() {
        String::new()
    } else {
        format!("    #[column({})]\n", attrs.join(", "))
    };

    // Determine if column is effectively NOT NULL:
    // Per SQLite docs (https://sqlite.org/lang_createtable.html):
    // - Explicit NOT NULL constraint
    // - INTEGER PRIMARY KEY is implicitly NOT NULL (special case)
    // - Other PRIMARY KEY types can technically be NULL due to SQLite legacy bug
    let is_integer_pk =
        is_pk && SQLTypeCategory::from_sql_type(&column.sql_type) == SQLTypeCategory::Integer;
    let is_not_null = column.not_null || is_integer_pk;

    // Determine Rust type from SQL type
    let rust_type = sql_type_to_rust_type(&column.sql_type, is_not_null);

    // Field name (snake_case)
    let field_name = apply_field_casing(column.name.as_ref(), field_casing);

    format!("{attr_str}    {vis}{field_name}: {rust_type},\n")
}

/// Format a default value for Rust syntax
fn format_default_value(default: &str, sql_type: &str) -> Option<String> {
    let category = SQLTypeCategory::from_sql_type(sql_type);

    // Skip function calls or complex expressions - these need default_fn
    if default.contains('(') && default.contains(')') {
        // Could be a function like CURRENT_TIMESTAMP - we'll return None
        // and add a warning instead
        return None;
    }

    match category {
        SQLTypeCategory::Integer => {
            // Boolean defaults
            if default == "0" || default == "1" {
                return Some(default.to_string());
            }
            // Integer defaults
            default.parse::<i64>().ok().map(|v| v.to_string())
        }
        SQLTypeCategory::Real => default.parse::<f64>().ok().map(|v| v.to_string()),
        SQLTypeCategory::Text | SQLTypeCategory::Blob => {
            // Remove surrounding quotes if present
            let trimmed = default.trim_matches(|c| c == '\'' || c == '"');
            Some(format!("\"{trimmed}\""))
        }
        SQLTypeCategory::Numeric => default
            .parse::<i64>()
            .map(|v| v.to_string())
            .ok()
            .or_else(|| default.parse::<f64>().map(|v| v.to_string()).ok()),
    }
}

/// Convert SQL type to Rust type
fn sql_type_to_rust_type(sql_type: &str, not_null: bool) -> String {
    // Handle boolean specifically before the category match
    if sql_type.eq_ignore_ascii_case("boolean") {
        return if not_null {
            "bool".to_string()
        } else {
            "Option<bool>".to_string()
        };
    }

    let category = SQLTypeCategory::from_sql_type(sql_type);

    let base_type = match category {
        SQLTypeCategory::Integer | SQLTypeCategory::Numeric => "i64",
        SQLTypeCategory::Real => "f64",
        SQLTypeCategory::Text => "String",
        SQLTypeCategory::Blob => "Vec<u8>",
    };

    if not_null {
        base_type.to_string()
    } else {
        format!("Option<{base_type}>")
    }
}

/// Generate an index struct
fn generate_index_struct(index: &Index, use_pub: bool, field_casing: FieldCasing) -> String {
    let mut code = String::new();
    let vis = if use_pub { "pub " } else { "" };

    // Index struct name is PascalCase
    let struct_name = index.name.to_pascal_case();
    let table_struct = index.table.to_pascal_case();

    // Build column references
    let columns: Vec<String> = index
        .columns
        .iter()
        .map(|c| format!("{}::{}", table_struct, c.value))
        .map(|s| {
            if let Some((table, col)) = s.split_once("::") {
                format!("{}::{}", table, apply_field_casing(col, field_casing))
            } else {
                s
            }
        })
        .collect();

    // Index attribute
    if index.is_unique {
        code.push_str("#[SQLiteIndex(unique)]\n");
    } else {
        code.push_str("#[SQLiteIndex]\n");
    }

    // Struct definition (tuple struct with column references)
    let _ = writeln!(
        code,
        "{}struct {}({});",
        vis,
        struct_name,
        columns.join(", ")
    );

    code
}

/// Generate a view struct
fn generate_view_struct(
    view: &View,
    columns: &[&Column],
    use_pub: bool,
    field_casing: FieldCasing,
) -> String {
    let struct_name = view.name.to_pascal_case();
    let vis = if use_pub { "pub " } else { "" };

    let mut code = String::new();

    // Build view attributes
    let mut attrs = Vec::new();

    // Check if view name differs from struct name (snake_case version)
    if apply_field_casing(&struct_name, field_casing) != view.name.as_ref() {
        attrs.push(format!("name = \"{}\"", view.name));
    }

    // Add definition
    if let Some(def) = &view.definition {
        let escaped_def = escape_for_rust_literal(def);
        attrs.push(format!("definition = \"{escaped_def}\""));
    }

    // Build the attribute line
    if attrs.is_empty() {
        code.push_str("#[SQLiteView]\n");
    } else {
        let _ = writeln!(code, "#[SQLiteView({})]", attrs.join(", "));
    }

    // Struct definition with column fields
    let _ = writeln!(code, "{vis}struct {struct_name} {{");

    // Sort columns by ordinal position
    let mut sorted_columns: Vec<&&Column> = columns.iter().collect();
    sorted_columns.sort_by(|a, b| {
        let ao = a.ordinal_position.unwrap_or(i32::MAX);
        let bo = b.ordinal_position.unwrap_or(i32::MAX);
        ao.cmp(&bo).then_with(|| a.name.cmp(&b.name))
    });

    // Generate fields for each column
    for column in sorted_columns {
        let field_name = apply_field_casing(column.name.as_ref(), field_casing);
        let rust_type = sql_type_to_rust_type(&column.sql_type, column.not_null);
        let _ = writeln!(code, "    {vis}{field_name}: {rust_type},");
    }

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

/// Generate a schema struct
fn generate_schema_struct(
    schema_name: &str,
    tables: &[String],
    indexes: &[String],
    use_pub: bool,
    field_casing: FieldCasing,
) -> String {
    let mut code = String::new();
    let vis = if use_pub { "pub " } else { "" };

    code.push_str("#[derive(SQLiteSchema)]\n");
    let _ = writeln!(code, "{vis}struct {schema_name} {{");

    // Add tables
    for table in tables {
        let field_name = apply_field_casing(table, field_casing);
        let type_name = table.to_pascal_case();
        let _ = writeln!(code, "    {vis}{field_name}: {type_name},");
    }

    // Add indexes
    for index in indexes {
        let field_name = apply_field_casing(index, field_casing);
        let type_name = index.to_pascal_case();
        let _ = writeln!(code, "    {vis}{field_name}: {type_name},");
    }

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

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

    #[test]
    fn test_generate_simple_table() {
        let mut ddl = SQLiteDDL::new();
        ddl.tables.push(Table::new("users"));
        ddl.columns.push(
            Column::new("users", "id", "integer")
                .not_null()
                .autoincrement(),
        );
        ddl.columns
            .push(Column::new("users", "name", "text").not_null());
        ddl.columns.push(Column::new("users", "email", "text"));
        ddl.pks.push(PrimaryKey::from_strings(
            "users".to_string(),
            "users_pk".to_string(),
            vec!["id".to_string()],
        ));

        let options = CodegenOptions {
            include_schema: false,
            schema_name: "AppSchema".to_string(),
            use_pub: true,
            ..Default::default()
        };

        let result = generate_rust_schema(&ddl, &options);

        assert_eq!(
            result.code,
            "\
//! Auto-generated SQLite schema from introspection
//!

use drizzle::sqlite::prelude::*;

#[SQLiteTable]
pub struct Users {
    pub email: Option<String>,
    #[column(primary, autoincrement)]
    pub id: i64,
    pub name: String,
}

"
        );
        assert_eq!(result.tables, vec!["users"]);
    }

    #[test]
    fn test_generate_table_with_unique() {
        let mut ddl = SQLiteDDL::new();
        ddl.tables.push(Table::new("accounts"));
        ddl.columns
            .push(Column::new("accounts", "id", "integer").not_null());
        ddl.columns
            .push(Column::new("accounts", "email", "text").not_null());
        ddl.uniques.push(UniqueConstraint::from_strings(
            "accounts".to_string(),
            "accounts_email_unique".to_string(),
            vec!["email".to_string()],
        ));

        let options = CodegenOptions::default();
        let result = generate_rust_schema(&ddl, &options);

        assert_eq!(
            result.code,
            "\
//! Auto-generated SQLite schema from introspection
//!

use drizzle::sqlite::prelude::*;

#[SQLiteTable]
struct Accounts {
    #[column(unique)]
    email: String,
    id: i64,
}

"
        );
    }

    #[test]
    fn test_generate_table_with_foreign_key() {
        let mut ddl = SQLiteDDL::new();
        ddl.tables.push(Table::new("posts"));
        ddl.columns
            .push(Column::new("posts", "id", "integer").not_null());
        ddl.columns
            .push(Column::new("posts", "author_id", "integer").not_null());

        let fk = ForeignKey::from_strings(
            "posts".to_string(),
            "fk_posts_author".to_string(),
            vec!["author_id".to_string()],
            "users".to_string(),
            vec!["id".to_string()],
        );
        ddl.fks.push(fk);

        let options = CodegenOptions::default();
        let result = generate_rust_schema(&ddl, &options);

        assert_eq!(
            result.code,
            "\
//! Auto-generated SQLite schema from introspection
//!

use drizzle::sqlite::prelude::*;

#[SQLiteTable]
struct Posts {
    #[column(references = Users::id)]
    author_id: i64,
    id: i64,
}

"
        );
    }

    #[test]
    fn test_generate_index() {
        let mut ddl = SQLiteDDL::new();
        ddl.tables.push(Table::new("users"));
        ddl.columns
            .push(Column::new("users", "email", "text").not_null());

        ddl.indexes.push(
            Index::new(
                "users",
                "users_email_idx",
                vec![IndexColumn {
                    value: "email".into(),
                    is_expression: false,
                }],
            )
            .unique(),
        );

        let options = CodegenOptions::default();
        let result = generate_rust_schema(&ddl, &options);

        assert_eq!(
            result.code,
            "\
//! Auto-generated SQLite schema from introspection
//!

use drizzle::sqlite::prelude::*;

#[SQLiteTable]
struct Users {
    email: String,
}

#[SQLiteIndex(unique)]
struct UsersEmailIdx(Users::email);

"
        );
    }

    #[test]
    fn test_generate_schema_struct() {
        let mut ddl = SQLiteDDL::new();
        ddl.tables.push(Table::new("users"));
        ddl.tables.push(Table::new("posts"));

        let options = CodegenOptions {
            include_schema: true,
            schema_name: "AppSchema".to_string(),
            use_pub: true,
            ..Default::default()
        };

        let result = generate_rust_schema(&ddl, &options);

        assert_eq!(
            result.code,
            "\
//! Auto-generated SQLite schema from introspection
//!

use drizzle::sqlite::prelude::*;

#[SQLiteTable]
pub struct Users {
}

#[SQLiteTable]
pub struct Posts {
}

#[derive(SQLiteSchema)]
pub struct AppSchema {
    pub users: Users,
    pub posts: Posts,
}
"
        );
    }

    #[test]
    fn test_sql_type_to_rust_type() {
        assert_eq!(sql_type_to_rust_type("integer", true), "i64");
        assert_eq!(sql_type_to_rust_type("integer", false), "Option<i64>");
        assert_eq!(sql_type_to_rust_type("text", true), "String");
        assert_eq!(sql_type_to_rust_type("text", false), "Option<String>");
        assert_eq!(sql_type_to_rust_type("real", true), "f64");
        assert_eq!(sql_type_to_rust_type("blob", true), "Vec<u8>");
        assert_eq!(sql_type_to_rust_type("boolean", true), "bool");
        assert_eq!(sql_type_to_rust_type("boolean", false), "Option<bool>");
        assert_eq!(sql_type_to_rust_type("BOOLEAN", true), "bool");
    }
}