const std = @import("std");
const ast = @import("ast.zig");
const llvm = @import("llvm.zig");

pub const CompilationError = error{
    TypeError,
    UnknownName,
};

pub const SymbolTable = std.hash_map.StringHashMap(*SymbolData);
pub const TypeList = std.ArrayList(SymbolUnderlyingType);

pub const SymbolUnderlyingTypeEnum = enum {
    Integer32,
    Integer64,
    Double,
    Bool,

    // opaque unsolved identifier
    OpaqueType,

    // solved custom type
    Struct,
    Enum,
};

pub const SymbolUnderlyingType = union(SymbolUnderlyingTypeEnum) {
    Integer32: void,
    Integer64: void,
    Double: void,
    Bool: void,
    OpaqueType: []const u8,

    Struct: []const u8,
    Enum: []const u8,
};

pub const ScopeList = std.ArrayList(*Scope);
pub const Scope = struct {
    parent: ?*Scope,
    env: UnderlyingTypeMap,
    meta_map: VariableMetadataMap,

    /// List of children in the scope.
    children: ScopeList,
    cur_child_idx: usize = 0,

    allocator: *std.mem.Allocator,
    id: ?[]const u8 = null,

    pub fn create(allocator: *std.mem.Allocator, parent: ?*Scope, id: ?[]const u8) !*Scope {
        var scope = try allocator.create(Scope);

        scope.* = Scope{
            .parent = parent,
            .env = UnderlyingTypeMap.init(allocator),
            .children = ScopeList.init(allocator),
            .allocator = allocator,
            .id = id,
            .meta_map = VariableMetadataMap.init(allocator),
        };

        return scope;
    }

    pub fn createChild(self: *@This(), id: ?[]const u8) !*Scope {
        var child = try @This().create(self.allocator, self, id);
        try self.children.append(child);
        return child;
    }

    pub fn nextChild(self: *@This()) *Scope {
        var child = self.children.items[self.cur_child_idx];
        self.cur_child_idx += 1;
        return child;
    }

    pub fn deinit(self: *const @This()) void {
        // XXX: fix this someday
        // self.env.deinit();
    }
};

pub const Parameter = struct {
    idx: usize,
    name: []const u8,
    typ: SymbolUnderlyingType,

    llvm_alloca: llvm.LLVMValueRef = null,
};

pub const ParameterMap = std.StringHashMap(*Parameter);

// functions, for our purposes, other than symbols, have:
//  - a return type
//  - parameters
pub const FunctionSymbol = struct {
    decl: ast.FnDecl,
    return_type: SymbolUnderlyingType,

    /// Parameters for a function are also a table instead of an ArrayList
    /// because we want to resolve identifiers to them.
    parameters: ParameterMap,
    parameter_list: TypeList,

    scope: *Scope,

    // only contains metadata for parameters (i really need to simplify this)
    meta_map: VariableMetadataMap,
};

// structs are hashmaps pointing to SymbolUnderlyingType
pub const UnderlyingTypeMap = std.hash_map.StringHashMap(SymbolUnderlyingType);

// enums have maps from fields to u32s
pub const IdentifierMap = std.StringHashMap(u32);

// TODO const
pub const SymbolType = enum {
    Function,
    Struct,
    Enum,
    Variable,
    Const,
};

pub const SymbolData = union(SymbolType) {
    Function: FunctionSymbol,
    Struct: UnderlyingTypeMap,
    Enum: IdentifierMap,

    // variables (parameters and variables), for the type system
    // only have types
    Variable: SymbolUnderlyingType,
    Const: SymbolUnderlyingType,
};

const builtin_type_identifiers = [_][]const u8{ "i32", "i64", "bool" };

const builtin_types = [_]SymbolUnderlyingTypeEnum{ .Integer32, .Integer64, .Bool };

const Using = enum {
    Scope,
    Function,
};

pub const VariableMetadata = struct {
    typ: SymbolUnderlyingType,
    llvm_alloca: llvm.LLVMValueRef = null,

    using: Using,

    from_scope: ?*Scope = null,
    from_function: ?*FunctionSymbol = null,

    pub fn withScope(
        allocator: *std.mem.Allocator,
        scope: *Scope,
        typ: SymbolUnderlyingType,
    ) !*VariableMetadata {
        var meta = try allocator.create(VariableMetadata);
        std.debug.warn("VARMETA create from scope={}, meta={}\n", .{ @ptrToInt(scope), @ptrToInt(meta) });
        meta.* = VariableMetadata{ .typ = typ, .from_scope = scope, .using = .Scope };
        return meta;
    }

    pub fn withParam(
        allocator: *std.mem.Allocator,
        func: *FunctionSymbol,
        typ: SymbolUnderlyingType,
    ) !*VariableMetadata {
        var meta = try allocator.create(VariableMetadata);
        std.debug.warn("VARMETA create from fndecl={}, meta={}\n", .{ @ptrToInt(func), @ptrToInt(meta) });
        meta.* = VariableMetadata{ .typ = typ, .from_function = func, .using = .Function };
        return meta;
    }
};

// TODO rm const?
pub const VariableMetadataMap = std.StringHashMap(*VariableMetadata);

/// Represents the context for a full compiler run.
/// This is used to manage the symbol table for the compilation unit, etc.
pub const CompilationContext = struct {
    allocator: *std.mem.Allocator,
    symbol_table: SymbolTable,

    cur_function: ?*FunctionSymbol = null,
    current_scope: ?*Scope = null,

    pub fn init(allocator: *std.mem.Allocator) CompilationContext {
        return CompilationContext{
            .allocator = allocator,
            .symbol_table = SymbolTable.init(allocator),
        };
    }

    pub fn deinit(self: *@This()) void {
        self.symbol_table.deinit();
    }

    /// Create a new scope out of the current one and set it as the current.
    pub fn bumpScope(self: *@This(), scope_id: ?[]const u8) !void {
        if (self.current_scope == null) {
            @panic("can't bump scope from null");
        }

        std.debug.warn("==scope bump== '{}'\n", .{scope_id});

        var child = try self.current_scope.?.createChild(scope_id);
        self.current_scope = child;
    }

    /// Set a given scope as the current scope.
    pub fn setScope(self: *@This(), scope: *Scope) void {
        std.debug.warn("==set== set scope to {}\n", .{scope.id});
        self.current_scope = scope;
    }

    /// "dump" the current scope, making the new current scope be its parent.
    pub fn dumpScope(self: *@This()) void {
        if (self.current_scope == null) {
            @panic("can't dump scope from null");
        }

        const parent_id: ?[]const u8 = if (self.current_scope.?.parent == null) null else self.current_scope.?.parent.?.id;

        std.debug.warn("==scope dump== {} to {}\n", .{
            self.current_scope.?.id,
            parent_id,
        });

        self.current_scope = self.current_scope.?.parent;
    }

    pub fn setCurrentFunction(self: *@This(), func_ctx: ?FunctionAnalysisContext) void {
        self.cur_function = func_ctx;
    }

    /// Solve a given type as a string into a SymbolUnderlyingTypeEnum
    /// This does not help if you want a full SymbolUnderlyingType, use
    /// solveType() for that.
    pub fn solveTypeEnum(
        self: *@This(),
        typ_ident: []const u8,
    ) SymbolUnderlyingTypeEnum {
        inline for (builtin_type_identifiers) |typ, idx| {
            if (std.mem.eql(u8, typ, typ_ident)) return builtin_types[idx];
        }

        return .OpaqueType;
    }

    /// Solve a given type string into a full fleged SymbolUnderlyingType.
    /// This always works, since resolution into struct/enum custom types
    /// become the fallback.
    pub fn solveType(
        self: *@This(),
        typ_ident: []const u8,
    ) SymbolUnderlyingType {
        const typ_enum_val = self.solveTypeEnum(typ_ident);

        return switch (typ_enum_val) {
            .Integer32 => SymbolUnderlyingType{ .Integer32 = {} },
            .Integer64 => SymbolUnderlyingType{ .Integer64 = {} },
            .Bool => SymbolUnderlyingType{ .Bool = {} },
            .OpaqueType => SymbolUnderlyingType{ .OpaqueType = typ_ident },
            else => unreachable,
        };
    }

    pub fn insertStruct(self: *@This(), struc: ast.Struct, field_types: TypeList) !void {
        var type_map = UnderlyingTypeMap.init(self.allocator);

        for (struc.fields.items) |field, idx| {
            _ = try type_map.put(field.name.lexeme, field_types.items[idx]);
        }

        var symbol = try self.allocator.create(SymbolData);
        symbol.* = SymbolData{ .Struct = type_map };
        _ = try self.symbol_table.put(struc.name.lexeme, symbol);
    }

    pub fn insertFn(
        self: *@This(),
        decl: ast.FnDecl,
        ret_type: SymbolUnderlyingType,
        param_types: TypeList,
        scope: *Scope,
    ) !void {
        var param_map = ParameterMap.init(self.allocator);

        for (decl.params.items) |param, idx| {
            var param_sym = try self.allocator.create(Parameter);

            param_sym.* = Parameter{
                .name = param.name.lexeme,
                .idx = idx,
                .typ = param_types.items[idx],
            };
            _ = try param_map.put(param.name.lexeme, param_sym);
        }

        const lex = decl.func_name.lexeme;

        var symbol = try self.allocator.create(SymbolData);
        symbol.* = SymbolData{
            .Function = FunctionSymbol{
                .decl = decl,
                .return_type = ret_type,
                .parameters = param_map,
                .parameter_list = param_types,
                .scope = scope,

                .meta_map = VariableMetadataMap.init(self.allocator),
            },
        };

        _ = try self.symbol_table.put(lex, symbol);

        var kv = self.symbol_table.get(lex);
        self.cur_function = &kv.?.Function;
    }

    pub fn insertEnum(self: *@This(), enu: ast.Enum) !void {
        var ident_map = IdentifierMap.init(self.allocator);
        errdefer ident_map.deinit();

        // TODO change this when enums get support for custom values
        for (enu.fields.items) |token, idx| {
            _ = try ident_map.put(token.lexeme, @intCast(u32, idx));
        }

        var symbol = try self.allocator.create(SymbolData);
        symbol.* = SymbolData{ .Enum = ident_map };
        _ = try self.symbol_table.put(enu.name.lexeme, symbol);
    }

    pub fn insertConst(self: *@This(), constdecl: ast.SingleConst, typ: SymbolUnderlyingType) !void {
        var symbol = try self.allocator.create(SymbolData);
        symbol.* = SymbolData{ .Const = typ };
        _ = try self.symbol_table.put(constdecl.name.lexeme, symbol);
    }

    pub fn fetchGlobalSymbol(
        self: *@This(),
        identifier: []const u8,
        typ: SymbolType,
    ) !*SymbolData {
        var sym_kv = self.symbol_table.get(identifier);
        if (sym_kv == null) {
            std.debug.warn("Unknown {} '{}'\n", .{ typ, identifier });
            return CompilationError.TypeError;
        }

        var value = sym_kv.?;

        var sym_typ = @as(SymbolType, value.*);
        if (sym_typ != typ) {
            std.debug.warn("Expected {}, got {}\n", .{ sym_typ, typ });
            return CompilationError.TypeError;
        }

        return value;
    }

    fn resolveVarTypeInScope(
        self: *@This(),
        scope_opt: ?*Scope,
        name: []const u8,
        create_meta: bool,
    ) error{OutOfMemory}!?*VariableMetadata {
        if (scope_opt == null) return null;
        var scope = scope_opt.?;

        var kv_opt = scope.env.get(name);
        if (kv_opt) |kv| {
            if (create_meta) {
                return try VariableMetadata.withScope(
                    self.allocator,
                    scope,
                    kv,
                );
            } else {
                return null;
            }
        } else {
            return self.resolveVarTypeInScope(scope.parent, name, create_meta);
        }
    }

    /// Resolve a given name's type, assuming it is a variable.
    pub fn resolveVarType(
        self: *@This(),
        name: []const u8,
        create_meta: bool,
    ) !?*VariableMetadata {
        var var_type: ?*VariableMetadata = null;

        if (self.current_scope) |scope| {
            var_type = try self.resolveVarTypeInScope(scope, name, create_meta);
            if (var_type) |typ| return typ;
        }

        if (self.cur_function) |cur_function| {
            var kv_opt = cur_function.parameters.get(name);
            if (kv_opt) |kv| {
                if (create_meta) {
                    return try VariableMetadata.withParam(
                        self.allocator,
                        cur_function,
                        kv.typ,
                    );
                } else {
                    return null;
                }
            }
        }

        std.debug.warn("Unknown name {}\n", .{name});
        return CompilationError.UnknownName;
    }

    pub fn insertMetadata(self: *@This(), name: []const u8, metadata: *VariableMetadata) !void {
        if (self.current_scope) |scope| {
            _ = try scope.meta_map.put(name, metadata);
            return;
        }

        if (self.cur_function) |func| {
            _ = try func.meta_map.put(name, metadata);
            return;
        }
    }
};