#include "parser.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>

// Helper function to allocate memory and handle errors
static void *safe_malloc(size_t size) {
    void *ptr = malloc(size);
    if (!ptr) {
        fprintf(stderr, "Memory allocation failed\n");
        exit(1);
    }
    return ptr;
}

// Helper function to create a new node
static ExprNode *expr_node_create(const char *token, int line) {
    ExprNode *node = (ExprNode *)safe_malloc(sizeof(ExprNode));
    node->token = strdup(token ? token : "");
    node->children = NULL;
    node->child_count = 0;
    node->line = line;
    return node;
}

// Forward declaration
static ExprNode *parse_expression(const char **ptr, int line);

// Skip whitespace characters
static const char *skip_whitespace(const char *ptr) {
    while (*ptr && isspace(*ptr)) {
        ptr++;
    }
    return ptr;
}

// Parse a token (atom)
static char *parse_token(const char **ptr) {
    const char *start = *ptr;
    
    // Skip leading whitespace
    start = skip_whitespace(start);
    if (!*start) return NULL;
    
    const char *end = start;
    
    // Handle parentheses specially
    if (*end == '(' || *end == ')') {
        end++;
    } else {
        // Read until whitespace or parentheses
        while (*end && !isspace(*end) && *end != '(' && *end != ')') {
            end++;
        }
    }
    
    if (end == start) return NULL;
    
    size_t len = end - start;
    char *token = (char *)safe_malloc(len + 1);
    memcpy(token, start, len);
    token[len] = '\0';
    
    *ptr = end;
    return token;
}

// Parse a list (expression starting with '(')
static ExprNode *parse_list(const char **ptr, int line) {
    // Skip the opening parenthesis
    (*ptr)++;
    
    // Parse the operator (first token)
    *ptr = skip_whitespace(*ptr);
    if (**ptr == ')') {
        // Empty list
        (*ptr)++;
        return expr_node_create("nil", line);
    }
    
    // Parse the operator/first element
    char *op_token = parse_token(ptr);
    if (!op_token) {
        fprintf(stderr, "Error: Expected operator at line %d\n", line);
        return NULL;
    }
    
    ExprNode *node = expr_node_create(op_token, line);
    free(op_token);
    
    // Parse children
    while (**ptr && **ptr != ')') {
        *ptr = skip_whitespace(*ptr);
        if (**ptr == ')') break;
        
        ExprNode *child = parse_expression(ptr, line);
        if (child) {
            // Resize children array properly
            ExprNode **new_children = (ExprNode **)safe_malloc(sizeof(ExprNode *) * (node->child_count + 1));
            
            // Copy existing children
            for (size_t i = 0; i < node->child_count; i++) {
                new_children[i] = node->children[i];
            }
            
            // Add new child
            new_children[node->child_count] = child;
            
            // Free old array and update
            free(node->children);
            node->children = new_children;
            node->child_count++;
        }
        
        *ptr = skip_whitespace(*ptr);
    }
    
    if (**ptr == ')') {
        (*ptr)++; // Skip closing parenthesis
    } else {
        fprintf(stderr, "Error: Missing closing parenthesis at line %d\n", line);
    }
    
    return node;
}

// Parse an expression (either atom or list)
static ExprNode *parse_expression(const char **ptr, int line) {
    *ptr = skip_whitespace(*ptr);
    
    if (!**ptr) return NULL;
    
    if (**ptr == '(') {
        return parse_list(ptr, line);
    } else {
        // Parse atom
        char *token = parse_token(ptr);
        if (token) {
            ExprNode *node = expr_node_create(token, line);
            free(token);
            return node;
        }
        return NULL;
    }
}

// Main parsing function
ExprNode *expr_parse(const char *source, size_t source_len) {
    if (!source || source_len == 0) return NULL;
    
    const char *ptr = source;
    int line = 1;
    
    ptr = skip_whitespace(ptr);
    if (!*ptr) return NULL;
    
    return parse_expression(&ptr, line);
}

// Free an Expr AST (and all children)
void expr_free(ExprNode *node) {
    if (!node) return;
    
    free(node->token);
    
    for (size_t i = 0; i < node->child_count; i++) {
        expr_free(node->children[i]);
    }
    free(node->children);
    free(node);
}

// Debug: print AST (for dev)
void expr_print(ExprNode *node, int indent) {
    if (!node) return;
    
    for (int i = 0; i < indent; i++) {
        printf("  ");
    }
    
    if (node->child_count == 0) {
        // Atom
        printf("Atom: '%s' (line %d)\n", node->token, node->line);
    } else {
        // List
        printf("List: '%s' (line %d) [%zu children]\n", 
               node->token, node->line, node->child_count);
        
        for (size_t i = 0; i < node->child_count; i++) {
            expr_print(node->children[i], indent + 1);
        }
    }
}