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

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 = t;
    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) {
			break;
		}

        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;

		t = tokenstream_next(ts);
        bool is_public = false;
        if (t.token == TOKEN_IDENTIFIER && strncmp(t.text.data, "public", t.text.length) == 0) {
            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");
			parser_free(module);
			return NULL;
		}

		module->imports[module->import_count] = (ImportDeclaration){ .module_name = t, .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");
			parser_free(module);
			return NULL;
		}
	}

    // Now t holds the first non-import token. If it's not TOKEN_ALIAS, return.
    // If it is, process it.
    while (t.token == TOKEN_ALIAS) {
        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;

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

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

        t = tokenstream_next(ts);
        if (t.token != TOKEN_IDENTIFIER || strncmp(t.text.data, "int32", t.text.length) != 0) {
            log_on_line(&t.location, t.location.column_end, "expected type");
            parser_free(module);
            return NULL;
        }
        alias.value = (TypeExpression){ .tag = TYPE_EXPRESSION_BUILTIN, .builtin = { .bitSize = 32, .isSigned = true } };
        
        module->aliases[module->alias_count] = alias;
        module->alias_count++;

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

	return module;
}

void parser_free(Module* module) {
	if (module == NULL) return;
    if (module->imports != NULL) {
        free(module->imports);
    }
    if (module->aliases != NULL) {
        free(module->aliases);
    }
	free(module);
}