/**
 * @file table.c
 * @brief Table implementation for Baba Yaga
 * @author eli_oat
 * @version 0.0.1
 * @date 2025
 * 
 * This file implements the table data structure for the Baba Yaga language.
 * Tables are immutable hash tables that support both string keys and numeric indices.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

#include "baba_yaga.h"

/* ============================================================================
 * Hash Table Implementation
 * ============================================================================ */

#define TABLE_INITIAL_CAPACITY 16
#define TABLE_LOAD_FACTOR 0.75

/**
 * @brief Hash table entry
 */
typedef struct TableEntry {
    char* key;           /**< String key */
    Value value;         /**< Associated value */
    struct TableEntry* next; /**< Next entry in chain */
} TableEntry;

/**
 * @brief Hash table structure
 */
typedef struct {
    TableEntry** buckets;    /**< Array of bucket chains */
    size_t capacity;         /**< Number of buckets */
    size_t size;             /**< Number of entries */
    Value* array_values;     /**< Array for numeric indices */
    size_t array_size;       /**< Size of array */
    size_t array_capacity;   /**< Capacity of array */
} HashTable;

/**
 * @brief Table value structure
 */
typedef struct {
    HashTable* hash_table;   /**< Hash table for string keys */
    int ref_count;           /**< Reference count for memory management */
} TableValue;

/* ============================================================================
 * Hash Function
 * ============================================================================ */

/**
 * @brief Simple hash function for strings
 * 
 * @param str String to hash
 * @return Hash value
 */
static unsigned int hash_string(const char* str) {
    unsigned int hash = 5381;
    int c;
    
    while ((c = *str++)) {
        hash = ((hash << 5) + hash) + c; /* hash * 33 + c */
    }
    
    return hash;
}

/* ============================================================================
 * Memory Management
 * ============================================================================ */

/**
 * @brief Create a new hash table
 * 
 * @return New hash table, or NULL on failure
 */
static HashTable* hash_table_create(void) {
    HashTable* table = malloc(sizeof(HashTable));
    if (table == NULL) {
        return NULL;
    }
    
    table->capacity = TABLE_INITIAL_CAPACITY;
    table->size = 0;
    table->buckets = calloc(table->capacity, sizeof(TableEntry*));
    if (table->buckets == NULL) {
        free(table);
        return NULL;
    }
    
    table->array_capacity = TABLE_INITIAL_CAPACITY;
    table->array_size = 0;
    table->array_values = calloc(table->array_capacity, sizeof(Value));
    if (table->array_values == NULL) {
        free(table->buckets);
        free(table);
        return NULL;
    }
    
    return table;
}

/**
 * @brief Destroy a hash table
 * 
 * @param table Hash table to destroy
 */
static void hash_table_destroy(HashTable* table) {
    if (table == NULL) {
        return;
    }
    
    /* Free all entries */
    for (size_t i = 0; i < table->capacity; i++) {
        TableEntry* entry = table->buckets[i];
        while (entry != NULL) {
            TableEntry* next = entry->next;
            free(entry->key);
            baba_yaga_value_destroy(&entry->value);
            free(entry);
            entry = next;
        }
    }
    
    /* Free array values */
    for (size_t i = 0; i < table->array_size; i++) {
        baba_yaga_value_destroy(&table->array_values[i]);
    }
    
    free(table->buckets);
    free(table->array_values);
    free(table);
}

/**
 * @brief Resize hash table
 * 
 * @param table Hash table to resize
 * @return true on success, false on failure
 */
static bool hash_table_resize(HashTable* table) {
    size_t old_capacity = table->capacity;
    TableEntry** old_buckets = table->buckets;
    
    table->capacity *= 2;
    table->buckets = calloc(table->capacity, sizeof(TableEntry*));
    if (table->buckets == NULL) {
        table->capacity = old_capacity;
        table->buckets = old_buckets;
        return false;
    }
    
    /* Rehash all entries */
    for (size_t i = 0; i < old_capacity; i++) {
        TableEntry* entry = old_buckets[i];
        while (entry != NULL) {
            TableEntry* next = entry->next;
            unsigned int hash = hash_string(entry->key) % table->capacity;
            entry->next = table->buckets[hash];
            table->buckets[hash] = entry;
            entry = next;
        }
    }
    
    free(old_buckets);
    return true;
}

/**
 * @brief Resize array part of table
 * 
 * @param table Hash table to resize
 * @return true on success, false on failure
 */
static bool hash_table_resize_array(HashTable* table) {
    size_t new_capacity = table->array_capacity * 2;
    Value* new_array = realloc(table->array_values, new_capacity * sizeof(Value));
    if (new_array == NULL) {
        return false;
    }
    
    table->array_values = new_array;
    table->array_capacity = new_capacity;
    return true;
}

/* ============================================================================
 * Table Operations
 * ============================================================================ */

/**
 * @brief Get entry from hash table by key
 * 
 * @param table Hash table
 * @param key String key
 * @return Table entry, or NULL if not found
 */
static TableEntry* hash_table_get_entry(const HashTable* table, const char* key) {
    if (table == NULL || key == NULL) {
        return NULL;
    }
    
    unsigned int hash = hash_string(key) % table->capacity;
    TableEntry* entry = table->buckets[hash];
    
    while (entry != NULL) {
        if (strcmp(entry->key, key) == 0) {
            return entry;
        }
        entry = entry->next;
    }
    
    return NULL;
}

/**
 * @brief Set value in hash table
 * 
 * @param table Hash table
 * @param key String key
 * @param value Value to set
 * @return true on success, false on failure
 */
static bool hash_table_set(HashTable* table, const char* key, const Value* value) {
    if (table == NULL || key == NULL) {
        return false;
    }
    
    /* Check if we need to resize */
    if ((double)table->size / table->capacity >= TABLE_LOAD_FACTOR) {
        if (!hash_table_resize(table)) {
            return false;
        }
    }
    
    unsigned int hash = hash_string(key) % table->capacity;
    TableEntry* entry = table->buckets[hash];
    
    /* Look for existing entry */
    while (entry != NULL) {
        if (strcmp(entry->key, key) == 0) {
            /* Update existing entry */
            baba_yaga_value_destroy(&entry->value);
            entry->value = baba_yaga_value_copy(value);
            return true;
        }
        entry = entry->next;
    }
    
    /* Create new entry */
    entry = malloc(sizeof(TableEntry));
    if (entry == NULL) {
        return false;
    }
    
    entry->key = strdup(key);
    if (entry->key == NULL) {
        free(entry);
        return false;
    }
    
    entry->value = baba_yaga_value_copy(value);
    entry->next = table->buckets[hash];
    table->buckets[hash] = entry;
    table->size++;
    
    return true;
}

/* ============================================================================
 * Public Table API
 * ============================================================================ */

Value baba_yaga_value_table(void) {
    Value value;
    value.type = VAL_TABLE;
    
    TableValue* table_value = malloc(sizeof(TableValue));
    if (table_value == NULL) {
        value.type = VAL_NIL;
        return value;
    }
    
    table_value->hash_table = hash_table_create();
    if (table_value->hash_table == NULL) {
        free(table_value);
        value.type = VAL_NIL;
        return value;
    }
    
    table_value->ref_count = 1;
    value.data.table = table_value;
    
    return value;
}

Value baba_yaga_table_get(const Value* table, const char* key) {
    if (table == NULL || table->type != VAL_TABLE || key == NULL) {
        DEBUG_ERROR("Table get: invalid parameters");
        return baba_yaga_value_nil();
    }
    
    TableValue* table_value = (TableValue*)table->data.table;
    DEBUG_DEBUG("Table get: looking for key '%s' in table with %zu entries", key, table_value->hash_table->size);
    
    TableEntry* entry = hash_table_get_entry(table_value->hash_table, key);
    
    if (entry != NULL) {
        DEBUG_DEBUG("Table get: found key '%s', returning value type %d", key, entry->value.type);
        return baba_yaga_value_copy(&entry->value);
    }
    
    DEBUG_DEBUG("Table get: key '%s' not found", key);
    return baba_yaga_value_nil();
}

Value baba_yaga_table_set(const Value* table, const char* key, const Value* value) {
    if (table == NULL || table->type != VAL_TABLE || key == NULL || value == NULL) {
        DEBUG_ERROR("Table set: invalid parameters");
        return baba_yaga_value_nil();
    }
    
    DEBUG_DEBUG("Table set: setting key '%s' to value type %d", key, value->type);
    
    /* Create new table */
    Value new_table = baba_yaga_value_table();
    if (new_table.type != VAL_TABLE) {
        DEBUG_ERROR("Table set: failed to create new table");
        return baba_yaga_value_nil();
    }
    
    TableValue* new_table_value = (TableValue*)new_table.data.table;
    TableValue* old_table_value = (TableValue*)table->data.table;
    
    DEBUG_DEBUG("Table set: copying %zu entries from old table", old_table_value->hash_table->size);
    
    /* Copy all entries from old table */
    for (size_t i = 0; i < old_table_value->hash_table->capacity; i++) {
        TableEntry* entry = old_table_value->hash_table->buckets[i];
        while (entry != NULL) {
            hash_table_set(new_table_value->hash_table, entry->key, &entry->value);
            entry = entry->next;
        }
    }
    
    /* Copy array values */
    for (size_t i = 0; i < old_table_value->hash_table->array_size; i++) {
        if (i >= new_table_value->hash_table->array_capacity) {
            if (!hash_table_resize_array(new_table_value->hash_table)) {
                baba_yaga_value_destroy(&new_table);
                return baba_yaga_value_nil();
            }
        }
        new_table_value->hash_table->array_values[i] = 
            baba_yaga_value_copy(&old_table_value->hash_table->array_values[i]);
    }
    new_table_value->hash_table->array_size = old_table_value->hash_table->array_size;
    
    /* Set the new value */
    if (!hash_table_set(new_table_value->hash_table, key, value)) {
        DEBUG_ERROR("Table set: failed to set key '%s'", key);
        baba_yaga_value_destroy(&new_table);
        return baba_yaga_value_nil();
    }
    
    DEBUG_DEBUG("Table set: new table has %zu entries", new_table_value->hash_table->size);
    return new_table;
}

Value baba_yaga_table_get_index(const Value* table, int index) {
    if (table == NULL || table->type != VAL_TABLE || index <= 0) {
        return baba_yaga_value_nil();
    }
    
    TableValue* table_value = (TableValue*)table->data.table;
    size_t idx = (size_t)(index - 1);
    
    if (idx < table_value->hash_table->array_size) {
        return baba_yaga_value_copy(&table_value->hash_table->array_values[idx]);
    }
    
    return baba_yaga_value_nil();
}

Value baba_yaga_table_set_index(const Value* table, int index, const Value* value) {
    if (table == NULL || table->type != VAL_TABLE || index <= 0 || value == NULL) {
        return baba_yaga_value_nil();
    }
    
    /* Create new table */
    Value new_table = baba_yaga_value_table();
    if (new_table.type != VAL_TABLE) {
        return baba_yaga_value_nil();
    }
    
    TableValue* new_table_value = (TableValue*)new_table.data.table;
    TableValue* old_table_value = (TableValue*)table->data.table;
    
    /* Copy all entries from old table */
    for (size_t i = 0; i < old_table_value->hash_table->capacity; i++) {
        TableEntry* entry = old_table_value->hash_table->buckets[i];
        while (entry != NULL) {
            hash_table_set(new_table_value->hash_table, entry->key, &entry->value);
            entry = entry->next;
        }
    }
    
    /* Copy array values */
    size_t idx = (size_t)(index - 1);
    size_t new_size = (idx >= old_table_value->hash_table->array_size) ? 
                     idx + 1 : old_table_value->hash_table->array_size;
    
    /* Ensure capacity */
    while (new_size >= new_table_value->hash_table->array_capacity) {
        if (!hash_table_resize_array(new_table_value->hash_table)) {
            baba_yaga_value_destroy(&new_table);
            return baba_yaga_value_nil();
        }
    }
    
    /* Copy existing values */
    for (size_t i = 0; i < old_table_value->hash_table->array_size; i++) {
        new_table_value->hash_table->array_values[i] = 
            baba_yaga_value_copy(&old_table_value->hash_table->array_values[i]);
    }
    
    /* Set the new value */
    new_table_value->hash_table->array_values[idx] = baba_yaga_value_copy(value);
    new_table_value->hash_table->array_size = new_size;
    
    return new_table;
}

size_t baba_yaga_table_size(const Value* table) {
    if (table == NULL || table->type != VAL_TABLE) {
        return 0;
    }
    
    TableValue* table_value = (TableValue*)table->data.table;
    return table_value->hash_table->size + table_value->hash_table->array_size;
}

bool baba_yaga_table_has_key(const Value* table, const char* key) {
    if (table == NULL || table->type != VAL_TABLE || key == NULL) {
        return false;
    }
    
    TableValue* table_value = (TableValue*)table->data.table;
    return hash_table_get_entry(table_value->hash_table, key) != NULL;
}

/**
 * @brief Get all keys from a table
 * 
 * @param table Table value
 * @param keys Array to store keys (caller must free)
 * @param max_keys Maximum number of keys to retrieve
 * @return Number of keys retrieved
 */
size_t baba_yaga_table_get_keys(const Value* table, char** keys, size_t max_keys) {
    if (table == NULL || table->type != VAL_TABLE || keys == NULL || max_keys == 0) {
        return 0;
    }
    
    TableValue* table_value = (TableValue*)table->data.table;
    HashTable* hash_table = table_value->hash_table;
    
    size_t key_count = 0;
    
    /* Get string keys */
    for (size_t i = 0; i < hash_table->capacity && key_count < max_keys; i++) {
        TableEntry* entry = hash_table->buckets[i];
        while (entry != NULL && key_count < max_keys) {
            keys[key_count] = strdup(entry->key);
            key_count++;
            entry = entry->next;
        }
    }
    
    /* Get numeric keys (array indices) */
    for (size_t i = 0; i < hash_table->array_size && key_count < max_keys; i++) {
        char* num_key = malloc(32); /* Enough for large numbers */
        if (num_key != NULL) {
            snprintf(num_key, 32, "%zu", i + 1); /* 1-based indexing */
            keys[key_count] = num_key;
            key_count++;
        }
    }
    
    return key_count;
}

/**
 * @brief Get a value from table by key (supports both string and numeric keys)
 * 
 * @param table Table value
 * @param key Key (string or numeric as string)
 * @return Value at key, or nil if not found
 */
Value baba_yaga_table_get_by_key(const Value* table, const char* key) {
    if (table == NULL || table->type != VAL_TABLE || key == NULL) {
        return baba_yaga_value_nil();
    }
    
    /* Try as string key first */
    Value result = baba_yaga_table_get(table, key);
    if (result.type != VAL_NIL) {
        return result;
    }
    
    /* Try as numeric key */
    char* endptr;
    long index = strtol(key, &endptr, 10);
    if (*endptr == '\0' && index > 0) {
        return baba_yaga_table_get_index(table, (int)index);
    }
    
    return baba_yaga_value_nil();
}

/* ============================================================================
 * Internal Table Management
 * ============================================================================ */

/**
 * @brief Increment reference count for a table
 * 
 * @param table Table value
 */
void table_increment_ref(Value* table) {
    if (table != NULL && table->type == VAL_TABLE) {
        TableValue* table_value = (TableValue*)table->data.table;
        table_value->ref_count++;
    }
}

/**
 * @brief Decrement reference count for a table
 * 
 * @param table Table value
 */
void table_decrement_ref(Value* table) {
    if (table != NULL && table->type == VAL_TABLE) {
        TableValue* table_value = (TableValue*)table->data.table;
        table_value->ref_count--;
        
        if (table_value->ref_count <= 0) {
            hash_table_destroy(table_value->hash_table);
            free(table_value);
        }
    }
}