clicktype-macros 0.2.0

//! Implementation of #[derive(ClickTable)] macro

use proc_macro::TokenStream;
use quote::{quote, format_ident};
use syn::{
    parse_macro_input, Data, DeriveInput, Fields, Attribute, Lit,
};

struct TableAttrs {
    name: String,
    _engine: Option<String>,
    _partition_by: Option<String>,
    _order_by: Option<String>,
    primary_key: Option<String>,
    _settings: Vec<(String, String)>,
}

struct FieldAttrs {
    name: Option<String>,
    primary_key: bool,
    skip: bool,
    codec: Option<String>,
    default: Option<String>,
    materialized: Option<String>,
    alias_expr: Option<String>,
}

pub fn derive(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let struct_name = &input.ident;

    // Parse table-level attributes
    let table_attrs = parse_table_attrs(&input.attrs);
    let table_name = table_attrs.name;

    // Parse fields
    let fields = match &input.data {
        Data::Struct(data) => match &data.fields {
            Fields::Named(fields) => &fields.named,
            _ => panic!("#[derive(ClickTable)] only supports structs with named fields"),
        },
        _ => panic!("#[derive(ClickTable)] only supports structs"),
    };

    // Generate column types and implementations
    let mut column_structs = vec![];
    let mut column_impl = vec![];
    let mut column_const_defs = vec![];
    let mut column_field_defs = vec![];
    let mut column_names = vec![];
    let mut insertable_columns = vec![];
    let mut primary_key_columns = vec![];
    let mut write_row_exprs = vec![];
    let mut read_row_exprs = vec![];
    let mut static_const_names = Vec::new();
    let mut static_const_types = Vec::new();
    let mut ddl_columns = vec![];
    let mut schema_entries = vec![];

    for field in fields {
        let field_name = field.ident.as_ref().unwrap();
        let field_type = &field.ty;
        let attrs = parse_field_attrs(&field.attrs);

        if attrs.skip {
            continue;
        }

        let column_name = attrs.name.unwrap_or_else(|| field_name.to_string());
        let column_struct_name = format_ident!("{}", to_pascal_case(&column_name));
        let column_const_name = format_ident!("{}", field_name.to_string().to_uppercase());

        column_names.push(column_name.clone());

        // Determine if column is insertable (exclude MATERIALIZED and ALIAS)
        if attrs.materialized.is_none() && attrs.alias_expr.is_none() {
            insertable_columns.push(column_name.clone());
            write_row_exprs.push(quote! {
                self.#field_name.write_binary(writer)?;
            });
        }

        if attrs.primary_key {
            primary_key_columns.push(column_name.clone());
        }

        // Read row expression (reads all columns)
        read_row_exprs.push(quote! {
            #field_name: clicktype_core::traits::ClickHouseType::read_binary(reader)?,
        });

        // Generate column struct
        column_structs.push(quote! {
            #[derive(Debug, Clone, Copy)]
            pub struct #column_struct_name;
        });

        // Implement TypedColumn
        let qualified_name = format!("{}.{}", table_name, column_name);
        let is_pk = attrs.primary_key;

        column_impl.push(quote! {
            impl clicktype_core::traits::TypedColumn for #column_struct_name {
                type Table = super::#struct_name;
                type Type = #field_type;

                fn name() -> &'static str {
                    #column_name
                }

                fn qualified_name() -> &'static str {
                    #qualified_name
                }

                fn is_primary_key() -> bool {
                    #is_pk
                }
            }
        });

        // Column const definition
        column_const_defs.push(quote! {
            #field_name: #column_struct_name,
        });

        // Column field for Columns struct
        column_field_defs.push(quote! {
            pub #field_name: #column_struct_name,
        });

        // Static const definition for Event::USER_ID syntax
        static_const_names.push(column_const_name);
        static_const_types.push(column_struct_name.clone());

        // Store field type for DDL generation
        ddl_columns.push(field_type);

        // Store schema entry (column_name, type_name)
        schema_entries.push(quote! {
            (#column_name, <#field_type as clicktype_core::traits::ClickHouseType>::type_name().to_string())
        });
    }

    // Process table-level primary_key attribute
    if let Some(pk_str) = table_attrs.primary_key {
        // Parse comma-separated list of column names
        for col_name in pk_str.split(',') {
            let trimmed = col_name.trim();
            if !trimmed.is_empty() && !primary_key_columns.contains(&trimmed.to_string()) {
                primary_key_columns.push(trimmed.to_string());
            }
        }
    }

    // Generate columns module name
    let columns_mod_name = format_ident!("{}_columns", struct_name.to_string().to_lowercase());

    let expanded = quote! {
        // Column types module
        pub mod #columns_mod_name {
            use super::*;
            use clicktype_core::traits::TypedColumn;

            #(#column_structs)*

            #(#column_impl)*

            /// Struct that groups all columns for ergonomic access
            #[derive(Debug, Clone, Copy)]
            pub struct Columns {
                #(#column_field_defs)*
            }

            impl Columns {
                pub const fn new() -> Self {
                    Self {
                        #(#column_const_defs)*
                    }
                }
            }

            impl Default for Columns {
                fn default() -> Self {
                    Self::new()
                }
            }
        }

        // Implement ClickTable trait
        impl clicktype_core::traits::ClickTable for #struct_name {
            type Columns = #columns_mod_name::Columns;

            fn table_name() -> &'static str {
                #table_name
            }

            fn columns() -> Self::Columns {
                #columns_mod_name::Columns::default()
            }

            fn column_names() -> &'static [&'static str] {
                &[#(#column_names),*]
            }

            fn insertable_columns() -> &'static [&'static str] {
                &[#(#insertable_columns),*]
            }

            fn primary_key_columns() -> &'static [&'static str] {
                &[#(#primary_key_columns),*]
            }

            fn create_table_ddl() -> String {
                let mut ddl = format!("CREATE TABLE {} (\n", #table_name);

                // Add columns with their types
                let columns: Vec<(&str, String)> = vec![
                    #(
                        (#column_names, <#ddl_columns as clicktype_core::traits::ClickHouseType>::type_name().to_string()),
                    )*
                ];

                for (i, (name, type_name)) in columns.iter().enumerate() {
                    if i > 0 {
                        ddl.push_str(",\n");
                    }
                    ddl.push_str(&format!("    {} {}", name, type_name));
                }

                ddl.push_str("\n) ENGINE = MergeTree()");

                // Add primary key if exists
                let pk_cols: &[&str] = &[#(#primary_key_columns),*];
                if !pk_cols.is_empty() {
                    ddl.push_str(&format!("\nPRIMARY KEY ({})", pk_cols.join(", ")));
                }

                ddl.push_str("\nORDER BY ");
                if !pk_cols.is_empty() {
                    ddl.push_str(&format!("({})", pk_cols.join(", ")));
                } else {
                    ddl.push_str("tuple()");
                }

                ddl
            }

            fn schema() -> Vec<(&'static str, String)> {
                vec![
                    #(#schema_entries),*
                ]
            }
        }

        // Implement ClickInsertable
        impl clicktype_core::traits::ClickInsertable for #struct_name {
            fn write_row<W: std::io::Write>(&self, writer: &mut W) -> Result<(), clicktype_core::error::WriteError> {
                use clicktype_core::traits::ClickHouseType;
                #(#write_row_exprs)*
                Ok(())
            }
        }

        // Implement ClickReadable
        impl clicktype_core::traits::ClickReadable for #struct_name {
            fn read_row<R: std::io::Read>(reader: &mut R) -> Result<Self, clicktype_core::error::ReadError> {
                use clicktype_core::traits::ClickHouseType;
                Ok(Self {
                    #(#read_row_exprs)*
                })
            }
        }

        // Static column accessors
        impl #struct_name {
            #(pub const #static_const_names: #columns_mod_name::#static_const_types = #columns_mod_name::#static_const_types;)*
        }
    };

    TokenStream::from(expanded)
}

fn parse_table_attrs(attrs: &[Attribute]) -> TableAttrs {
    let mut name = None;
    let mut engine = None;
    let mut partition_by = None;
    let mut order_by = None;
    let mut primary_key = None;
    let settings = Vec::new();

    for attr in attrs {
        if !attr.path().is_ident("click_table") {
            continue;
        }

        let _ = attr.parse_nested_meta(|meta| {
            if meta.path.is_ident("name") {
                let value: Lit = meta.value()?.parse()?;
                if let Lit::Str(s) = value {
                    name = Some(s.value());
                }
            } else if meta.path.is_ident("engine") {
                let value: Lit = meta.value()?.parse()?;
                if let Lit::Str(s) = value {
                    engine = Some(s.value());
                }
            } else if meta.path.is_ident("partition_by") {
                let value: Lit = meta.value()?.parse()?;
                if let Lit::Str(s) = value {
                    partition_by = Some(s.value());
                }
            } else if meta.path.is_ident("order_by") {
                let value: Lit = meta.value()?.parse()?;
                if let Lit::Str(s) = value {
                    order_by = Some(s.value());
                }
            } else if meta.path.is_ident("primary_key") {
                let value: Lit = meta.value()?.parse()?;
                if let Lit::Str(s) = value {
                    primary_key = Some(s.value());
                }
            }
            Ok(())
        });
    }

    TableAttrs {
        name: name.expect("click_table attribute requires 'name' parameter"),
        _engine: engine,
        _partition_by: partition_by,
        _order_by: order_by,
        primary_key,
        _settings: settings,
    }
}

fn parse_field_attrs(attrs: &[Attribute]) -> FieldAttrs {
    let mut field_attrs = FieldAttrs {
        name: None,
        primary_key: false,
        skip: false,
        codec: None,
        default: None,
        materialized: None,
        alias_expr: None,
    };

    for attr in attrs {
        if !attr.path().is_ident("click_column") {
            continue;
        }

        let _ = attr.parse_nested_meta(|meta| {
            if meta.path.is_ident("primary_key") {
                field_attrs.primary_key = true;
            } else if meta.path.is_ident("skip") {
                field_attrs.skip = true;
            } else if meta.path.is_ident("rename") {
                let value: Lit = meta.value()?.parse()?;
                if let Lit::Str(s) = value {
                    field_attrs.name = Some(s.value());
                }
            } else if meta.path.is_ident("codec") {
                let value: Lit = meta.value()?.parse()?;
                if let Lit::Str(s) = value {
                    field_attrs.codec = Some(s.value());
                }
            } else if meta.path.is_ident("default") {
                let value: Lit = meta.value()?.parse()?;
                if let Lit::Str(s) = value {
                    field_attrs.default = Some(s.value());
                }
            } else if meta.path.is_ident("materialized") {
                let value: Lit = meta.value()?.parse()?;
                if let Lit::Str(s) = value {
                    field_attrs.materialized = Some(s.value());
                }
            } else if meta.path.is_ident("alias") {
                let value: Lit = meta.value()?.parse()?;
                if let Lit::Str(s) = value {
                    field_attrs.alias_expr = Some(s.value());
                }
            }
            Ok(())
        });
    }

    field_attrs
}

fn to_pascal_case(s: &str) -> String {
    s.split('_')
        .map(|word| {
            let mut chars = word.chars();
            match chars.next() {
                None => String::new(),
                Some(first) => {
                    first.to_uppercase().chain(chars).collect()
                }
            }
        })
        .collect()
}