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

typedef struct {
    TokenStream* ts;
    Token token;
} Parser;

/**
 * Reads a new token into p->token.
 */
static void parser_next_token(Parser* p) {
    p->token = tokenstream_next(p->ts);
}

/**
 * Reads a new token if the current token is equal to the expected token.
 * 
 * If they are equal, it continues to the next token.
 * 
 * @param p
 * @param token The expected token.
 * @returns `true` if the current token matches the expected, `false` if it does not.
 */
static bool parser_accept(Parser* p, TokenType token) {
    if (p->token.token == token) {
        parser_next_token(p);
        return true;
    }
    return false;
}

/**
 * Consumes the expected token if present, otherwise logs an error.
 * 
 * @param p The parser state.
 * @param token The token type to expect.
 * @param msg The error message if the token is not found.
 * @return true if the token was consumed, false otherwise.
 */
static bool parser_expect(Parser* p, TokenType token, const char* msg) {
    if (parser_accept(p, token)) {
        return true;
    }
    log_on_line(&p->token.location, msg);
    return false;
}

/**
 * Checks if the current token matches the expected token type without consuming it.
 * 
 * @param p The parser state.
 * @param token The token type to check.
 * @return true if the current token matches, false otherwise.
 */
static bool parser_peek(Parser* p, TokenType token) {
    if (p->token.token == token) {
        return true;
    }
    return false;
}

/**
 * Checks if the current token matches the expected token type without consuming it.
 * 
 * @param p The parser state.
 * @param token The token type to check.
 * @return true if the current token matches, false otherwise.
 */
static bool parser_require(Parser* p, TokenType token, const char* msg) {
    if (parser_peek(p, token)) {
        return true;
    }
    log_on_line(&p->token.location, msg);
    return false;
}

/**
 * Returns the text of the current token and advances the parser to the next token.
 * 
 * @param p The parser state.
 * @return A newly allocated string containing the current token's text.
 */
static char* parser_to_text(Parser* p) {
    char* str = string_copy(p->token.text);
    parser_next_token(p);
    return str;
}

/**
 * Parses the "module" keyword
 */
static bool parse_module_declaration(Parser* p, Module* module) {
    if (!parser_expect(p, TOKEN_MODULE, "expected keyword 'module'")) {
        return false;
    }

    if (!parser_require(p, TOKEN_IDENTIFIER, "expected module identifier")) {
        return false;
    }
    module->name = parser_to_text(p);

    return parser_expect(p, TOKEN_SEMICOLON, "expected ';' after module name");
}

/**
 * Parses an import declaration.
 * 
 * @param p The parser state.
 * @param module The module to add the import to.
 * @param is_public Whether the import is public.
 * @return true if successful, false otherwise.
 */
static bool parse_import_declaration(Parser* p, Module* module, bool is_public) {
    module->import_count++;
    module->imports = realloc(module->imports, sizeof(ImportDeclaration) * module->import_count);

    ImportDeclaration* import = &module->imports[module->import_count - 1];
    memset(import, 0, sizeof(ImportDeclaration));
    import->is_public = is_public;

    if (!parser_require(p, TOKEN_IDENTIFIER, "expected module identifier")) {
        return false;
    }

    import->module_name = parser_to_text(p);

    if (!parser_expect(p, TOKEN_SEMICOLON, "expected ';' after import")) {
        return false;
    }

    return true;
}

/**
 * Parses a primitive type expression.
 * 
 * @param p The parser state.
 * @param expr The type expression to populate.
 * @return true if successful, false otherwise.
 */
static bool parse_primitive_type_expression(Parser* p, TypeExpression* expr) {
    if (parser_accept(p, TOKEN_U8)) {
        expr->tag = TYPE_EXPRESSION_BUILTIN;
        expr->builtin.bitSize = 8;
        expr->builtin.isSigned = false;
        return true;
    } else if (parser_accept(p, TOKEN_U16)) {
        expr->tag = TYPE_EXPRESSION_BUILTIN;
        expr->builtin.bitSize = 16;
        expr->builtin.isSigned = false;
        return true;
    } else if (parser_accept(p, TOKEN_U32)) {
        expr->tag = TYPE_EXPRESSION_BUILTIN;
        expr->builtin.bitSize = 32;
        expr->builtin.isSigned = false;
        return true;
    } else if (parser_accept(p, TOKEN_U64)) {
        expr->tag = TYPE_EXPRESSION_BUILTIN;
        expr->builtin.bitSize = 64;
        expr->builtin.isSigned = false;
        return true;
    } else if (parser_accept(p, TOKEN_I8)) {
        expr->tag = TYPE_EXPRESSION_BUILTIN;
        expr->builtin.bitSize = 8;
        expr->builtin.isSigned = true;
        return true;
    } else if (parser_accept(p, TOKEN_I16)) {
        expr->tag = TYPE_EXPRESSION_BUILTIN;
        expr->builtin.bitSize = 16;
        expr->builtin.isSigned = true;
        return true;
    } else if (parser_accept(p, TOKEN_I32)) {
        expr->tag = TYPE_EXPRESSION_BUILTIN;
        expr->builtin.bitSize = 32;
        expr->builtin.isSigned = true;
        return true;
    } else if (parser_accept(p, TOKEN_I64)) {
        expr->tag = TYPE_EXPRESSION_BUILTIN;
        expr->builtin.bitSize = 64;
        expr->builtin.isSigned = true;
        return true;
    } else {
        log_on_line(&p->token.location, "expected type expression");
        return false;
    }
}

/**
 * Parses an array type expression.
 * 
 * @param p The parser state.
 * @param expr The type expression to populate.
 * @return true if successful, false otherwise.
 */
static bool parse_array_type_expression(Parser* p, TypeExpression* expr) {
    TypeExpression elementType;
    if (!parse_primitive_type_expression(p, &elementType)) {
        return false;
    }

    if (parser_accept(p, TOKEN_BRACKET_OPEN)) {
        expr->tag = TYPE_EXPRESSION_ARRAY;
        expr->array.array = malloc(sizeof(TypeExpression));
        *expr->array.array = elementType;

        if (!parser_expect(p, TOKEN_BRACKET_CLOSE, "expected ']' to end array type")) {
            return false;
        }
    } else {
         *expr = elementType;
         return true;
    }
    return true;
}

/**
 * Parses a type expression.
 * 
 * @param p The parser state.
 * @param expr The type expression to populate.
 * @return true if successful, false otherwise.
 */
static bool parse_type_expression(Parser* p, TypeExpression* expr) {
    return parse_array_type_expression(p, expr);
}

/**
 * Parses an alias declaration.
 * 
 * @param p The parser state.
 * @param module The module to add the alias to.
 * @param is_public Whether the alias is public.
 * @return true if successful, false otherwise.
 */
static bool parse_alias_declaration(Parser* p, Module* module, bool is_public) {
    module->alias_count++;
    module->aliases = realloc(module->aliases, sizeof(AliasDeclaration) * module->alias_count);

    AliasDeclaration* alias = &module->aliases[module->alias_count - 1];
    memset(alias, 0, sizeof(AliasDeclaration));
    alias->is_public = is_public;

    if (!parser_require(p, TOKEN_IDENTIFIER, "expected alias identifier")) {
        return false;
    }
    alias->name = parser_to_text(p);

    if (!parser_expect(p, TOKEN_ASSIGN, "expected '=' after alias name")) {
        return false;
    }

    if (!parse_type_expression(p, &alias->value)) {
        return false;
    }

    if (!parser_expect(p, TOKEN_SEMICOLON, "expected ';' after alias declaration")) {
        return false;
    }

    return true;
}

/**
 * Parses a variable declaration.
 * 
 * @param p The parser state.
 * @param module The module to add the variable to.
 * @param is_public Whether the variable is public.
 * @param is_static Whether the variable is static.
 * @param is_const Whether the variable is constant.
 * @return true if successful, false otherwise.
 */
static bool parse_variable_declaration(Parser* p, Module* module, bool is_public, bool is_static, bool is_const) {
    module->variable_count++;
    module->variables = realloc(module->variables, sizeof(VariableDeclaration) * module->variable_count);

    VariableDeclaration* var = &module->variables[module->variable_count - 1];
    memset(var, 0, sizeof(VariableDeclaration));
    var->is_public = is_public;
    var->is_static = is_static;
    var->is_const = is_const;

    if (!parse_type_expression(p, &var->type)) {
        return false;
    }

    if (!parser_require(p, TOKEN_IDENTIFIER, "expected variable identifier")) {
        return false;
    }
    var->name = parser_to_text(p);

    if (!parser_expect(p, TOKEN_SEMICOLON, "expected ';' after variable declaration")) {
        return false;
    }

    return true;
}

/**
 * Checks if the current token is a primitive type.
 * 
 * The token will not be consumed by this function, so the caller can decide how to handle it if it is a primitive type.
 * 
 * @param p The parser state.
 * @return true if the current token is a primitive type, false otherwise.
 */
static bool parser_accept_primitive(Parser* p) {
    return parser_peek(p, TOKEN_I8) || parser_peek(p, TOKEN_I16) || 
           parser_peek(p, TOKEN_I32) || parser_peek(p, TOKEN_I64) ||
           parser_peek(p, TOKEN_U8) || parser_peek(p, TOKEN_U16) || 
           parser_peek(p, TOKEN_U32) || parser_peek(p, TOKEN_U64);
}

Module* parser_parse(TokenStream* ts) {
    Parser* p = malloc(sizeof(Parser));
    p->ts = ts;
    parser_next_token(p);

    Module* module = malloc(sizeof(Module));
    memset(module, 0, sizeof(Module));
    if (!parse_module_declaration(p, module)) {
        goto fail;
    }

    while (!parser_peek(p, TOKEN_EOF)) {
        bool is_public = false;
        bool is_static = false;
        bool is_const = false;
        bool terminal = false;
        
        while (!terminal) {
            if (parser_accept(p, TOKEN_IMPORT)) {
                if (is_static) {
                    log_on_line(&p->token.location, "import declarations cannot be static or const");
                    goto fail;
                }
                if (is_const) {
                    log_on_line(&p->token.location, "import declarations cannot be static or const");
                    goto fail;
                }
                if (!parse_import_declaration(p, module, is_public)) {
                    goto fail;
                }
                terminal = true;
            } else if (parser_accept(p, TOKEN_ALIAS)) {
                if (is_static) {
                    log_on_line(&p->token.location, "alias declarations cannot be static or const");
                    goto fail;
                }
                if (is_const) {
                    log_on_line(&p->token.location, "alias declarations cannot be static or const");
                    goto fail;
                }
                if (!parse_alias_declaration(p, module, is_public)) {
                    goto fail;
                }
                terminal = true;
            } else if (parser_accept(p, TOKEN_PUBLIC)) {
                is_public = true;
            } else if (parser_accept(p, TOKEN_STATIC)) {
                is_static = true;
            } else if (parser_accept(p, TOKEN_CONST)) {
                is_const = true;
            } else if (parser_accept_primitive(p)) {
                if (!parse_variable_declaration(p, module, is_public, is_static, is_const)) {
                    goto fail;
                }
                terminal = true;
            } else {
                log_on_line(&p->token.location, "unexpected token");
                goto fail;
            }
        }
    }

    free(p);
    return module;
fail:
    free(p);
    parser_free(module);
    return NULL;
}

void free_type_expression(TypeExpression* expr) {
    if (expr->tag == TYPE_EXPRESSION_ARRAY) {
        free_type_expression(expr->array.array);
        free(expr->array.array);
    }
}

void parser_free(Module* module) {
	if (module == NULL) {
        return;
    }

    if (module->imports != NULL) {
        for(size_t i = 0; i < module->import_count; i++) {
            free(module->imports[i].module_name);
        }
        free(module->imports);
    }
    
    if (module->aliases != NULL) {
        for(size_t i = 0; i < module->alias_count; i++) {
            free(module->aliases[i].name);
            free_type_expression(&module->aliases[i].value);
        }
        free(module->aliases);
    }

    if (module->variables != NULL) {
        for(size_t i = 0; i < module->variable_count; i++) {
            free(module->variables[i].name);
            free_type_expression(&module->variables[i].type);
        }
        free(module->variables);
    }

    free(module->name);
    free(module);
}