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

pub const CompilationError = error{TypeError};

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

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

    // opaque unsolved identifier
    OpaqueType,

    // solved custom type
    Struct,
    Enum,
};

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

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

// functions, for our purposes, other than symbols, have:
//  - a return type
//  - TODO 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: UnderlyingTypeMap,
    symbols: SymbolTable,

    /// Find a given identifier in the function. Can resolve to either a parameter
    pub fn findSymbol(self: *const @This(), identifier: []const u8) ?SymbolData {
        // try to find it on overall variable symbols

        // TODO

        //var var_sym = self.symbols.get(identifier);
        //if (var_sym != null) return var_sym.?.value;

        //var param_sym = self.parameters.get(identifier);
        //if (param_sym != null) return param_sym.?.value;
        return null;
    }
};

// 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,
};

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 builtin_type_identifiers = [_][]const u8{ "i32", "i64", "bool" };

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

/// 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,

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

    /// 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.toSlice()) |field, idx| {
            _ = try type_map.put(field.name.lexeme, field_types.at(idx));
        }

        _ = try self.symbol_table.put(struc.name.lexeme, SymbolData{ .Struct = type_map });
    }

    pub fn insertFn(
        self: *@This(),
        decl: ast.FnDecl,
        ret_type: SymbolUnderlyingType,
        param_types: TypeList,
        symbols: SymbolTable,
    ) !void {
        var type_map = UnderlyingTypeMap.init(self.allocator);

        for (decl.params.toSlice()) |param, idx| {
            _ = try type_map.put(param.name.lexeme, param_types.at(idx));
        }

        _ = try self.symbol_table.put(decl.func_name.lexeme, SymbolData{
            .Function = FunctionSymbol{
                .decl = decl,
                .return_type = ret_type,
                .parameters = type_map,
                .symbols = symbols,
            },
        });
    }

    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.toSlice()) |token, idx| {
            _ = try ident_map.put(token.lexeme, @intCast(u32, idx));
        }

        _ = try self.symbol_table.put(enu.name.lexeme, SymbolData{
            .Enum = ident_map,
        });
    }
};