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

/**
 * Parses an import declaration.
 * 
 * @param ts The token stream to parse from.
 * @param module The module being parsed.
 * @returns true on success, false on failure.
 */
static bool parse_import(TokenStream* ts, Module* module) {
    ImportDeclaration* new_imports = realloc(module->imports, (module->import_count + 1) * sizeof(ImportDeclaration));
    if (!new_imports) {
        fprintf(stderr, "Out of memory\n");
        exit(1);
    }
    module->imports = new_imports;

    Token t = tokenstream_next(ts);
    bool is_public = false;
    if (t.token == TOKEN_PUBLIC) {
        is_public = true;
        t = tokenstream_next(ts);
    }

    if (t.token != TOKEN_IDENTIFIER) {
        log_on_line(&t.location, t.location.column_end, "expected module name to import");
        return false;
    }
    
    char* name = (char*)malloc(t.text.length + 1);
    memcpy(name, t.text.data, t.text.length);
    name[t.text.length] = '\0';

    module->imports[module->import_count] = (ImportDeclaration){ .module_name = name, .is_public = is_public };
    module->import_count++;

    t = tokenstream_next(ts);
    if (t.token != TOKEN_SEMICOLON) {
        log_on_line(&t.location, t.location.column_end, "expected ';' after import");
        return false;
    }
    return true;
}

/**
 * Parses an alias declaration.
 * 
 * @param ts The token stream to parse from.
 * @param module The module being parsed.
 * @returns true on success, false on failure.
 */
static bool parse_alias(TokenStream* ts, Module* module) {
    AliasDeclaration* new_aliases = realloc(module->aliases, (module->alias_count + 1) * sizeof(AliasDeclaration));
    if (!new_aliases) {
        fprintf(stderr, "Out of memory\n");
        exit(1);
    }
    module->aliases = new_aliases;

    Token t = tokenstream_next(ts);
    if (t.token != TOKEN_IDENTIFIER) {
        log_on_line(&t.location, t.location.column_end, "expected alias name");
        return false;
    }
    char* name = (char*)malloc(t.text.length + 1);
    memcpy(name, t.text.data, t.text.length);
    name[t.text.length] = '\0';
    AliasDeclaration alias;
    alias.name = name;

    t = tokenstream_next(ts);
    if (t.token != TOKEN_ASSIGN) {
        log_on_line(&t.location, t.location.column_end, "expected '='");
        return false;
    }

    t = tokenstream_next(ts);
    
    TypeExpression type;
    if (t.token == TOKEN_IDENTIFIER && strncmp(t.text.data, "int32", t.text.length) == 0) {
        type = (TypeExpression){ .tag = TYPE_EXPRESSION_BUILTIN, .builtin = { .bitSize = 32, .isSigned = true } };
        t = tokenstream_next(ts);
        if (t.token == TOKEN_BRACKET_OPEN) {
            t = tokenstream_next(ts);
            if (t.token != TOKEN_BRACKET_CLOSE) {
                log_on_line(&t.location, t.location.column_end, "expected ']'");
                return false;
            }
            TypeExpression* inner = malloc(sizeof(TypeExpression));
            *inner = type;
            type = (TypeExpression){ .tag = TYPE_EXPRESSION_ARRAY, .array = { .array = inner } };
            t = tokenstream_next(ts);
        }
    } else {
        log_on_line(&t.location, t.location.column_end, "expected type");
        return false;
    }

    alias.value = type;
    
    module->aliases[module->alias_count] = alias;
    module->alias_count++;

    if (t.token != TOKEN_SEMICOLON) {
        log_on_line(&t.location, t.location.column_end, "expected ';'");
        return false;
    }
    return true;
}

Module* parser_parse(TokenStream* ts) {
    Token t = tokenstream_next(ts);
	if (t.token != TOKEN_MODULE) {
		log_on_line(&t.location, t.location.column_end, "expected 'module' keyword");
		return NULL;
	}

    t = tokenstream_next(ts);
    if (t.token != TOKEN_IDENTIFIER) {
        log_on_line(&t.location, t.location.column_end, "expected module name");
        return NULL;
    }

    Module* module = (Module*)malloc(sizeof(Module));
    if (module == NULL) {
        fprintf(stderr, "Out of memory\n");
        exit(1);
    }
    module->name = (const char*)malloc(t.text.length + 1);
    memcpy((void*)module->name, t.text.data, t.text.length);
    ((char*)module->name)[t.text.length] = '\0';
    module->imports = NULL;
    module->import_count = 0;
    module->aliases = NULL;
    module->alias_count = 0;

	t = tokenstream_next(ts);
	if (t.token != TOKEN_SEMICOLON) {
		log_on_line(&t.location, t.location.column_end, "expected ';' after module name");
		parser_free(module);
		return NULL;
	}

	while (1) {
		t = tokenstream_next(ts);
		if (t.token == TOKEN_IMPORT) {
            if (!parse_import(ts, module)) {
                parser_free(module);
                return NULL;
            }
        } else if (t.token == TOKEN_ALIAS) {
            if (!parse_alias(ts, module)) {
                parser_free(module);
                return NULL;
            }
        } else {
            break;
        }
	}

	return module;
}

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((void*)module->imports[i].module_name);
        }
        free(module->imports);
    }
    if (module->aliases != NULL) {
        for(size_t i = 0; i < module->alias_count; i++) {
            free((void*)module->aliases[i].name);
            free_type_expression(&module->aliases[i].value);
        }
        free(module->aliases);
    }
    free((void*)module->name);
	free(module);
}