rohas_codegen/

typescript.rs

use crate::error::Result;
use crate::templates;
use rohas_parser::{Api, Event, FieldType, Model, Schema, WebSocket};
use std::fs;
use std::path::Path;

pub fn generate_models(schema: &Schema, output_dir: &Path) -> Result<()> {
    let models_dir = output_dir.join("generated/models");

    for model in &schema.models {
        let content = generate_model_content(model);
        let file_name = format!("{}.ts", templates::to_snake_case(&model.name));
        fs::write(models_dir.join(file_name), content)?;
    }

    Ok(())
}

fn generate_model_content(model: &Model) -> String {
    let mut content = String::new();

    content.push_str("import { z } from 'zod';\n\n");

    content.push_str(&format!("export interface {} {{\n", model.name));

    for field in &model.fields {
        let ts_type = field.field_type.to_typescript();
        let optional = if field.optional { "?" } else { "" };
        content.push_str(&format!("  {}{}: {};\n", field.name, optional, ts_type));
    }

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

    // Generate zod schema
    content.push_str(&format!(
        "export const {}Schema = z.object({{\n",
        model.name
    ));
    for field in &model.fields {
        let zod_type = field_type_to_zod(&field.field_type, field.optional);
        content.push_str(&format!("  {}: {},\n", field.name, zod_type));
    }
    content.push_str("});\n\n");

    content.push_str(&format!(
        "export function is{}(obj: any): obj is {} {{\n",
        model.name, model.name
    ));
    content.push_str(&format!(
        "  return {}Schema.safeParse(obj).success;\n",
        model.name
    ));
    content.push_str("}\n");

    content
}

fn field_type_to_zod(field_type: &rohas_parser::FieldType, optional: bool) -> String {
    use rohas_parser::FieldType;

    let zod_type = match field_type {
        FieldType::Int | FieldType::Float => "z.number()".to_string(),
        FieldType::String => "z.string()".to_string(),
        FieldType::Boolean => "z.boolean()".to_string(),
        FieldType::DateTime => "z.date()".to_string(),
        FieldType::Json => "z.any()".to_string(),
        FieldType::Custom(name) => format!("{}Schema", name),
        FieldType::Array(inner) => {
            let inner_zod = field_type_to_zod(inner, false);
            format!("z.array({})", inner_zod)
        }
    };

    if optional {
        format!("{}.optional()", zod_type)
    } else {
        zod_type
    }
}

pub fn generate_dtos(schema: &Schema, output_dir: &Path) -> Result<()> {
    let dto_dir = output_dir.join("generated/dto");

    for input in &schema.inputs {
        let content = generate_model_content(&rohas_parser::Model {
            name: input.name.clone(),
            fields: input.fields.clone(),
            attributes: vec![],
        });
        let file_name = format!("{}.ts", templates::to_snake_case(&input.name));
        fs::write(dto_dir.join(file_name), content)?;
    }

    Ok(())
}

pub fn generate_apis(schema: &Schema, output_dir: &Path) -> Result<()> {
    let api_dir = output_dir.join("generated/api");

    for api in &schema.apis {
        let content = generate_api_content(api);
        let file_name = format!("{}.ts", templates::to_snake_case(&api.name));
        fs::write(api_dir.join(file_name), content)?;
    }

    let handlers_dir = output_dir.join("handlers/api");
    for api in &schema.apis {
        let file_name = format!("{}.ts", &api.name);
        let handler_path = handlers_dir.join(&file_name);

        if !handler_path.exists() {
            let content = generate_api_handler_stub(api);
            fs::write(handler_path, content)?;
        }
    }

    Ok(())
}

fn generate_api_content(api: &Api) -> String {
    let mut content = String::new();

    content.push_str("import { z } from 'zod';\n");

    let request_type = format!("{}Request", api.name);
    let response_type = format!("{}Response", api.name);
    let handler_type = format!("{}Handler", api.name);

    let response_is_primitive = is_primitive_type(&api.response);

    if !response_is_primitive {
        content.push_str(&format!(
            "import {{ {}, {}Schema }} from '@generated/models/{}';\n",
            api.response,
            api.response,
            templates::to_snake_case(&api.response)
        ));
    }

    if let Some(body) = &api.body {
        let body_is_primitive = is_primitive_type(body);
        if !body_is_primitive {
            if body.ends_with("Input") {
                content.push_str(&format!(
                    "import {{ {}, {}Schema }} from '@generated/dto/{}';\n",
                    body,
                    body,
                    templates::to_snake_case(body)
                ));
            } else {
                content.push_str(&format!(
                    "import {{ {}, {}Schema }} from '@generated/models/{}';\n",
                    body,
                    body,
                    templates::to_snake_case(body)
                ));
            }
        }
    }

    if !content.is_empty() {
        content.push_str("\n");
    }

    let path_params = extract_path_params(&api.path);

    content.push_str(&format!("export interface {} {{\n", request_type));

    for param in &path_params {
        content.push_str(&format!("  {}: string;\n", param));
    }

    if let Some(body) = &api.body {
        let ts_type = if is_primitive_type(body) {
            primitive_to_typescript(body)
        } else {
            body.to_string()
        };
        content.push_str(&format!("  body: {};\n", ts_type));
    }

    content.push_str("  queryParams?: Record<string, string>;\n");

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

    // Generate zod schema for request
    content.push_str(&format!(
        "export const {}Schema = z.object({{\n",
        request_type
    ));
    for param in &path_params {
        content.push_str(&format!("  {}: z.string(),\n", param));
    }
    if let Some(body) = &api.body {
        let body_is_primitive = is_primitive_type(body);
        if body_is_primitive {
            let zod_type = match body.as_str() {
                "String" => "z.string()",
                "Int" | "Float" => "z.number()",
                "Boolean" => "z.boolean()",
                "DateTime" | "Date" => "z.date()",
                _ => "z.any()",
            };
            content.push_str(&format!("  body: {},\n", zod_type));
        } else {
            if body.ends_with("Input") {
                content.push_str(&format!("  body: {}Schema,\n", body));
            } else {
                content.push_str(&format!("  body: {}Schema,\n", body));
            }
        }
    }
    content.push_str("  queryParams: z.record(z.string()).optional(),\n");
    content.push_str("});\n\n");

    let response_ts_type = if response_is_primitive {
        primitive_to_typescript(&api.response)
    } else {
        api.response.clone()
    };

    content.push_str(&format!("export interface {} {{\n", response_type));
    content.push_str(&format!("  data: {};\n", response_ts_type));
    content.push_str("}\n\n");

    // Generate zod schema for response
    let response_zod_type = if response_is_primitive {
        match api.response.as_str() {
            "String" => "z.string()".to_string(),
            "Int" | "Float" => "z.number()".to_string(),
            "Boolean" => "z.boolean()".to_string(),
            "DateTime" | "Date" => "z.date()".to_string(),
            _ => "z.any()".to_string(),
        }
    } else {
        format!("{}Schema", api.response)
    };
    content.push_str(&format!(
        "export const {}Schema = z.object({{\n",
        response_type
    ));
    content.push_str(&format!("  data: {},\n", response_zod_type));
    content.push_str("});\n\n");

    content.push_str(&format!(
        "export type {} = (req: {}) => Promise<{}>;\n",
        handler_type, request_type, response_type
    ));

    content
}

fn is_primitive_type(type_name: &str) -> bool {
    matches!(
        type_name,
        "String" | "Int" | "Float" | "Boolean" | "DateTime" | "Date"
    )
}

fn primitive_to_typescript(type_name: &str) -> String {
    match type_name {
        "String" => "string".to_string(),
        "Int" | "Float" => "number".to_string(),
        "Boolean" => "boolean".to_string(),
        "DateTime" | "Date" => "Date".to_string(),
        _ => type_name.to_string(),
    }
}

/// Extract path parameters from a path string
/// e.g., "/test/{name}" -> ["name"]
/// e.g., "/users/{id}/posts/{postId}" -> ["id", "postId"]
fn extract_path_params(path: &str) -> Vec<String> {
    let mut params = Vec::new();
    let mut in_param = false;
    let mut current_param = String::new();

    for ch in path.chars() {
        match ch {
            '{' => {
                in_param = true;
                current_param.clear();
            }
            '}' => {
                if in_param && !current_param.is_empty() {
                    params.push(current_param.clone());
                }
                in_param = false;
            }
            _ if in_param => {
                current_param.push(ch);
            }
            _ => {}
        }
    }

    params
}

fn generate_api_handler_stub(api: &Api) -> String {
    let mut content = String::new();

    let request_type = format!("{}Request", api.name);
    let response_type = format!("{}Response", api.name);
    let handler_name = format!("handle{}", api.name);

    content.push_str(&format!(
        "import {{ {}, {} }} from '@generated/api/{}';\n",
        request_type,
        response_type,
        templates::to_snake_case(&api.name)
    ));
    content.push_str("import { State } from '@generated/state';\n\n");

    content.push_str(&format!(
        "export async function {}(req: {}, state: State): Promise<{}> {{\n",
        handler_name, request_type, response_type
    ));
    content.push_str("  // TODO: Implement handler logic\n");
    content.push_str("  // For auto-triggers (defined in schema triggers): use state.setPayload('EventName', {...})\n");
    content.push_str("  // For manual triggers: use state.triggerEvent('EventName', {...})\n");
    content.push_str("  throw new Error('Not implemented');\n");
    content.push_str("}\n");

    content
}

pub fn generate_events(schema: &Schema, output_dir: &Path) -> Result<()> {
    let events_dir = output_dir.join("generated/events");

    for event in &schema.events {
        let content = generate_event_content(event);
        let file_name = format!("{}.ts", templates::to_snake_case(&event.name));
        fs::write(events_dir.join(file_name), content)?;
    }

    // Generate handler stubs
    let handlers_dir = output_dir.join("handlers/events");
    for event in &schema.events {
        for handler in &event.handlers {
            let file_name = format!("{}.ts", handler);
            let handler_path = handlers_dir.join(&file_name);

            if !handler_path.exists() {
                let content = generate_event_handler_stub(event, handler);
                fs::write(handler_path, content)?;
            }
        }
    }

    Ok(())
}

fn payload_type_to_zod(type_name: &str) -> String {
    match type_name {
        "String" => "z.string()".to_string(),
        "Int" | "Float" => "z.number()".to_string(),
        "Boolean" | "Bool" => "z.boolean()".to_string(),
        "DateTime" | "Date" => "z.date()".to_string(),
        _ => format!("{}Schema", type_name),
    }
}

fn generate_event_content(event: &Event) -> String {
    let mut content = String::new();

    content.push_str("import { z } from 'zod';\n");

    let payload_is_primitive = is_primitive_type(&event.payload);

    if payload_is_primitive {
        content.push_str("\n");
    } else {
        content.push_str(&format!(
            "import {{ {}, {}Schema }} from '@generated/models/{}';\n\n",
            event.payload,
            event.payload,
            templates::to_snake_case(&event.payload)
        ));
    }

    let payload_ts_type = if payload_is_primitive {
        primitive_to_typescript(&event.payload)
    } else {
        event.payload.clone()
    };

    content.push_str(&format!("export interface {} {{\n", event.name));
    content.push_str(&format!("  payload: {};\n", payload_ts_type));
    content.push_str("  timestamp: Date;\n");
    content.push_str("}\n\n");

    // Generate zod schema for event
    let payload_zod_type = payload_type_to_zod(&event.payload);
    content.push_str(&format!(
        "export const {}Schema = z.object({{\n",
        event.name
    ));
    content.push_str(&format!("  payload: {},\n", payload_zod_type));
    content.push_str("  timestamp: z.date(),\n");
    content.push_str("});\n\n");

    content.push_str(&format!(
        "export type {}Handler = (event: {}) => Promise<void>;\n",
        event.name, event.name
    ));

    content
}

fn generate_event_handler_stub(event: &Event, handler_name: &str) -> String {
    let mut content = String::new();

    content.push_str(&format!(
        "import {{ {} }} from '@generated/events/{}';\n\n",
        event.name,
        templates::to_snake_case(&event.name)
    ));

    content.push_str(&format!(
        "export async function {}(event: {}): Promise<void> {{\n",
        handler_name, event.name
    ));
    content.push_str("  // TODO: Implement event handler\n");
    content.push_str(&format!("  console.log('Handling event:', event);\n"));
    content.push_str("}\n");

    content
}

pub fn generate_crons(schema: &Schema, output_dir: &Path) -> Result<()> {
    let cron_dir = output_dir.join("generated/cron");

    for cron in &schema.crons {
        let content = format!(
            "export interface {} {{\n  schedule: string;\n}}\n",
            cron.name
        );
        let file_name = format!("{}.ts", templates::to_snake_case(&cron.name));
        fs::write(cron_dir.join(file_name), content)?;
    }

    // Generate handler stubs
    let handlers_dir = output_dir.join("handlers/cron");
    for cron in &schema.crons {
        let file_name = format!("{}.ts", templates::to_snake_case(&cron.name));
        let handler_path = handlers_dir.join(&file_name);

        if !handler_path.exists() {
            let content = format!(
                "export async function handle{}(): Promise<void> {{\n  // TODO: Implement cron job\n  console.log('Running cron: {}');\n}}\n",
                cron.name, cron.name
            );
            fs::write(handler_path, content)?;
        }
    }

    Ok(())
}

pub fn generate_websockets(schema: &Schema, output_dir: &Path) -> Result<()> {
    let ws_dir = output_dir.join("generated/websockets");

    for ws in &schema.websockets {
        let content = generate_websocket_content(ws);
        let file_name = format!("{}.ts", templates::to_snake_case(&ws.name));
        fs::write(ws_dir.join(file_name), content)?;
    }

    let handlers_dir = output_dir.join("handlers/websockets");
    for ws in &schema.websockets {
        if !ws.on_connect.is_empty() {
            for handler in &ws.on_connect {
                let file_name = format!("{}.ts", handler);
                let handler_path = handlers_dir.join(&file_name);
                if !handler_path.exists() {
                    let content = generate_websocket_handler_stub(ws, "onConnect", handler);
                    fs::write(handler_path, content)?;
                }
            }
        }
        if !ws.on_message.is_empty() {
            for handler in &ws.on_message {
                let file_name = format!("{}.ts", handler);
                let handler_path = handlers_dir.join(&file_name);
                if !handler_path.exists() {
                    let content = generate_websocket_handler_stub(ws, "onMessage", handler);
                    fs::write(handler_path, content)?;
                }
            }
        }
        if !ws.on_disconnect.is_empty() {
            for handler in &ws.on_disconnect {
                let file_name = format!("{}.ts", handler);
                let handler_path = handlers_dir.join(&file_name);
                if !handler_path.exists() {
                    let content = generate_websocket_handler_stub(ws, "onDisconnect", handler);
                    fs::write(handler_path, content)?;
                }
            }
        }
    }

    Ok(())
}

pub fn generate_middlewares(schema: &Schema, output_dir: &Path) -> Result<()> {
    use std::collections::HashSet;
    
    let mut middleware_names = HashSet::new();
    
    for api in &schema.apis {
        for middleware in &api.middlewares {
            middleware_names.insert(middleware.clone());
        }
    }
    
    for ws in &schema.websockets {
        for middleware in &ws.middlewares {
            middleware_names.insert(middleware.clone());
        }
    }
    
    if middleware_names.is_empty() {
        return Ok(());
    }
    
    let middlewares_dir = output_dir.join("middlewares");
    for middleware_name in middleware_names {
        let file_name = format!("{}.ts", middleware_name);
        let middleware_path = middlewares_dir.join(&file_name);
        
        if !middleware_path.exists() {
            let content = generate_middleware_stub(&middleware_name);
            fs::write(middleware_path, content)?;
        }
    }
    
    Ok(())
}

fn generate_middleware_stub(middleware_name: &str) -> String {
    let mut content = String::new();
    
    content.push_str("import { State } from '@generated/state';\n\n");
    
    content.push_str("export interface MiddlewareContext {\n");
    content.push_str("  payload?: any;\n");
    content.push_str("  query_params?: Record<string, string>;\n");
    content.push_str("  connection?: any;\n");
    content.push_str("  websocket_name?: string;\n");
    content.push_str("  api_name?: string;\n");
    content.push_str("  trace_id?: string;\n");
    content.push_str("}\n\n");
    
    content.push_str(&format!(
        "export async function {}Middleware(\n",
        middleware_name
    ));
    content.push_str("  context: MiddlewareContext,\n");
    content.push_str("  state: State\n");
    content.push_str("): Promise<MiddlewareContext | null> {\n");
    content.push_str("  /**\n");
    content.push_str(&format!("   * Middleware function for {}.\n", middleware_name));
    content.push_str("   * \n");
    content.push_str("   * @param context - Request context containing:\n");
    content.push_str("   *   - payload: Request payload (for APIs)\n");
    content.push_str("   *   - query_params: Query parameters (for APIs)\n");
    content.push_str("   *   - connection: WebSocket connection info (for WebSockets)\n");
    content.push_str("   *   - websocket_name: WebSocket name (for WebSockets)\n");
    content.push_str("   *   - api_name: API name (for APIs)\n");
    content.push_str("   *   - trace_id: Trace ID\n");
    content.push_str("   * @param state - State object for logging and triggering events\n");
    content.push_str("   * @returns Modified context with 'payload' and/or 'query_params' keys,\n");
    content.push_str("   *   or null to pass through unchanged. Throw an error to reject the request.\n");
    content.push_str("   * \n");
    content.push_str("   * To reject the request, throw an error:\n");
    content.push_str("   *   throw new Error('Access denied');\n");
    content.push_str("   * \n");
    content.push_str("   * To modify the request:\n");
    content.push_str("   *   return {\n");
    content.push_str("   *     ...context,\n");
    content.push_str("   *     payload: modifiedPayload,\n");
    content.push_str("   *     query_params: modifiedQueryParams\n");
    content.push_str("   *   };\n");
    content.push_str("   */\n");
    content.push_str("  // TODO: Implement middleware logic\n");
    content.push_str("  // Example: Validate authentication\n");
    content.push_str("  // Example: Rate limiting\n");
    content.push_str("  // Example: Logging\n");
    content.push_str("  // Example: Modify payload/query_params\n");
    content.push_str("  \n");
    content.push_str("  // Pass through unchanged\n");
    content.push_str("  return null;\n");
    content.push_str("}\n");
    
    content
}

fn generate_websocket_content(ws: &WebSocket) -> String {
    let mut content = String::new();

    content.push_str("import { z } from 'zod';\n");

    if let Some(message_type) = &ws.message {
        let message_field_type = FieldType::from_str(message_type);
        let is_custom_type = matches!(message_field_type, FieldType::Custom(_));
        if is_custom_type {
            content.push_str(&format!(
                "import {{ {}, {}Schema }} from '@generated/dto/{}';\n",
                message_type,
                message_type,
                templates::to_snake_case(message_type)
            ));
        }
    }

    content.push_str("\n");

    if let Some(message_type) = &ws.message {
        let message_field_type = FieldType::from_str(message_type);
        let ts_type = message_field_type.to_typescript();
        content.push_str(&format!("export interface {}Message {{\n", ws.name));
        content.push_str(&format!("  data: {};\n", ts_type));
        content.push_str("  timestamp: Date;\n");
        content.push_str("}\n\n");

        // Generate zod schema
        let zod_type = if matches!(message_field_type, FieldType::Custom(_)) {
            format!("{}Schema", message_type)
        } else {
            field_type_to_zod(&message_field_type, false)
        };
        content.push_str(&format!(
            "export const {}MessageSchema = z.object({{\n",
            ws.name
        ));
        content.push_str(&format!("  data: {},\n", zod_type));
        content.push_str("  timestamp: z.date(),\n");
        content.push_str("});\n\n");
    } else {
        content.push_str(&format!("export interface {}Message {{\n", ws.name));
        content.push_str("  data: any;\n");
        content.push_str("  timestamp: Date;\n");
        content.push_str("}\n\n");

        content.push_str(&format!(
            "export const {}MessageSchema = z.object({{\n",
            ws.name
        ));
        content.push_str("  data: z.any(),\n");
        content.push_str("  timestamp: z.date(),\n");
        content.push_str("});\n\n");
    }

    // Generate connection type
    content.push_str(&format!("export interface {}Connection {{\n", ws.name));
    content.push_str("  connectionId: string;\n");
    content.push_str("  path: string;\n");
    content.push_str("  connectedAt: Date;\n");
    content.push_str("}\n\n");

    content.push_str(&format!(
        "export const {}ConnectionSchema = z.object({{\n",
        ws.name
    ));
    content.push_str("  connectionId: z.string(),\n");
    content.push_str("  path: z.string(),\n");
    content.push_str("  connectedAt: z.date(),\n");
    content.push_str("});\n");

    content
}

fn generate_websocket_handler_stub(
    ws: &WebSocket,
    handler_type: &str,
    handler_name: &str,
) -> String {
    let mut content = String::new();

    content.push_str(&format!(
        "import {{ {}Message, {}Connection }} from '@generated/websockets/{}';\n",
        ws.name,
        ws.name,
        templates::to_snake_case(&ws.name)
    ));
    content.push_str("import { State } from '@generated/state';\n\n");

    match handler_type {
        "onConnect" => {
            content.push_str(&format!(
                "export async function {}(connection: {}Connection, state: State): Promise<{}Message | null> {{\n",
                handler_name,
                ws.name,
                ws.name
            ));
            content.push_str("  // TODO: Implement onConnect handler\n");
            content
                .push_str("  // Return a message to send to the client on connection, or null\n");
            content.push_str(&format!(
                "  console.log('Client connected:', connection.connectionId);\n"
            ));
            content.push_str("  return null;\n");
            content.push_str("}\n");
        }
        "onMessage" => {
            content.push_str(&format!(
                "export async function {}(message: {}Message, connection: {}Connection, state: State): Promise<{}Message | null> {{\n",
                handler_name,
                ws.name,
                ws.name,
                ws.name
            ));
            content.push_str("  // TODO: Implement onMessage handler\n");
            content.push_str("  // Return a message to send back to the client, or null\n");
            content.push_str(&format!(
                "  console.log('Received message:', message.data);\n"
            ));
            content.push_str("  // For auto-triggers (defined in schema triggers): use state.setPayload('EventName', {...})\n");
            content
                .push_str("  // For manual triggers: use state.triggerEvent('EventName', {...})\n");
            content.push_str("  return null;\n");
            content.push_str("}\n");
        }
        "onDisconnect" => {
            content.push_str(&format!(
                "export async function {}(connection: {}Connection, state: State): Promise<void> {{\n",
                handler_name,
                ws.name
            ));
            content.push_str("  // TODO: Implement onDisconnect handler\n");
            content.push_str(&format!(
                "  console.log('Client disconnected:', connection.connectionId);\n"
            ));
            content.push_str("}\n");
        }
        _ => {}
    }

    content
}

pub fn generate_state(output_dir: &Path) -> Result<()> {
    let generated_dir = output_dir.join("generated");
    let content = r#"export interface TriggeredEvent {
  eventName: string;
  payload: any;
}

/**
 * Logger for handlers to emit structured logs.
 */
export class Logger {
  private handlerName: string;
  private logFn?: (level: string, handler: string, message: string, fields: Record<string, any>) => void;

  constructor(handlerName: string, logFn?: (level: string, handler: string, message: string, fields: Record<string, any>) => void) {
    this.handlerName = handlerName;
    this.logFn = logFn;
  }

  /**
   * Log an info message.
   * 
   * @param message - Log message
   * @param fields - Additional fields to include in the log
   */
  info(message: string, fields?: Record<string, any>): void {
    if (this.logFn) {
      this.logFn("info", this.handlerName, message, fields || {});
    } else {
      console.log(`[${this.handlerName}] ${message}`, fields || {});
    }
  }

  /**
   * Log an error message.
   * 
   * @param message - Log message
   * @param fields - Additional fields to include in the log
   */
  error(message: string, fields?: Record<string, any>): void {
    if (this.logFn) {
      this.logFn("error", this.handlerName, message, fields || {});
    } else {
      console.error(`[${this.handlerName}] ${message}`, fields || {});
    }
  }

  /**
   * Log a warning message.
   * 
   * @param message - Log message
   * @param fields - Additional fields to include in the log
   */
  warning(message: string, fields?: Record<string, any>): void {
    if (this.logFn) {
      this.logFn("warn", this.handlerName, message, fields || {});
    } else {
      console.warn(`[${this.handlerName}] ${message}`, fields || {});
    }
  }

  /**
   * Log a warning message (alias for warning).
   * 
   * @param message - Log message
   * @param fields - Additional fields to include in the log
   */
  warn(message: string, fields?: Record<string, any>): void {
    this.warning(message, fields);
  }

  /**
   * Log a debug message.
   * 
   * @param message - Log message
   * @param fields - Additional fields to include in the log
   */
  debug(message: string, fields?: Record<string, any>): void {
    if (this.logFn) {
      this.logFn("debug", this.handlerName, message, fields || {});
    } else {
      console.debug(`[${this.handlerName}] ${message}`, fields || {});
    }
  }

  /**
   * Log a trace message.
   * 
   * @param message - Log message
   * @param fields - Additional fields to include in the log
   */
  trace(message: string, fields?: Record<string, any>): void {
    if (this.logFn) {
      this.logFn("trace", this.handlerName, message, fields || {});
    } else {
      console.trace(`[${this.handlerName}] ${message}`, fields || {});
    }
  }
}

/**
 * Context object for handlers to trigger events and access runtime state.
 */
export class State {
  private triggers: TriggeredEvent[] = [];
  private autoTriggerPayloads: Map<string, any> = new Map();
  public logger: Logger;

  constructor(handlerName?: string, logFn?: (level: string, handler: string, message: string, fields: Record<string, any>) => void) {
    this.logger = new Logger(handlerName || "unknown", logFn);
  }

  /**
   * Manually trigger an event with the given payload.
   * 
   * Use this for events that are NOT defined in the schema's triggers list.
   * 
   * @param eventName - Name of the event to trigger
   * @param payload - Event payload data (will be serialized to JSON)
   */
  triggerEvent(eventName: string, payload: any): void {
    this.triggers.push({
      eventName,
      payload,
    });
  }

  /**
   * Set the payload for an auto-triggered event.
   * 
   * Use this for events that ARE defined in the schema's triggers list.
   * The event will be automatically triggered after the handler completes,
   * using the payload you set here.
   * 
   * @param eventName - Name of the event (must match a trigger in schema)
   * @param payload - Event payload data (will be serialized to JSON)
   */
  setPayload(eventName: string, payload: any): void {
    this.autoTriggerPayloads.set(eventName, payload);
  }

  /**
   * Get all manually triggered events. Used internally by the runtime.
   */
  getTriggers(): TriggeredEvent[] {
    return [...this.triggers];
  }

  /**
   * Get payload for an auto-triggered event. Used internally by the runtime.
   */
  getAutoTriggerPayload(eventName: string): any | undefined {
    return this.autoTriggerPayloads.get(eventName);
  }

  /**
   * Get all auto-trigger payloads. Used internally by the runtime.
   */
  getAllAutoTriggerPayloads(): Map<string, any> {
    return new Map(this.autoTriggerPayloads);
  }
}
"#;

    fs::write(generated_dir.join("state.ts"), content)?;
    Ok(())
}

pub fn generate_index(schema: &Schema, output_dir: &Path) -> Result<()> {
    let mut content = String::new();

    content.push_str("export * from './state';\n\n");

    content.push_str("// Models\n");
    for model in &schema.models {
        content.push_str(&format!(
            "export * from './models/{}';\n",
            templates::to_snake_case(&model.name)
        ));
    }

    content.push_str("\n// DTOs\n");
    for input in &schema.inputs {
        content.push_str(&format!(
            "export * from './dto/{}';\n",
            templates::to_snake_case(&input.name)
        ));
    }

    content.push_str("\n// APIs\n");
    for api in &schema.apis {
        content.push_str(&format!(
            "export * from './api/{}';\n",
            templates::to_snake_case(&api.name)
        ));
    }

    content.push_str("\n// Events\n");
    for event in &schema.events {
        content.push_str(&format!(
            "export * from './events/{}';\n",
            templates::to_snake_case(&event.name)
        ));
    }

    content.push_str("\n// WebSockets\n");
    for ws in &schema.websockets {
        content.push_str(&format!(
            "export * from './websockets/{}';\n",
            templates::to_snake_case(&ws.name)
        ));
    }

    fs::write(output_dir.join("generated/index.ts"), content)?;

    Ok(())
}