const std = @import("std");
const Allocator = std.mem.Allocator;
const types = @import("types.zig");
const inst = @import("inst.zig");

const TypeId = types.TypeId;
const TypeInfo = types.TypeInfo;
const TypeTable = types.TypeTable;
const StringId = types.StringId;
const Ref = inst.Ref;
const BlockId = inst.BlockId;
const FuncId = inst.FuncId;
const GlobalId = inst.GlobalId;
const Inst = inst.Inst;
const Op = inst.Op;
const Block = inst.Block;
const Function = inst.Function;
const Global = inst.Global;
const Span = inst.Span;

// ── Module ──────────────────────────────────────────────────────────────

pub const Module = struct {
    types: TypeTable,
    functions: std.ArrayList(Function),
    globals: std.ArrayList(Global),
    /// Maps (protocol_ty, concrete_ty) → list of method FuncIds.
    impl_table: ImplTable,
    alloc: Allocator,

    pub fn init(alloc: Allocator) Module {
        return .{
            .types = TypeTable.init(alloc),
            .functions = std.ArrayList(Function).empty,
            .globals = std.ArrayList(Global).empty,
            .impl_table = ImplTable.init(alloc),
            .alloc = alloc,
        };
    }

    pub fn deinit(self: *Module) void {
        for (self.functions.items) |*func| {
            func.deinit(self.alloc);
        }
        self.functions.deinit(self.alloc);
        self.globals.deinit(self.alloc);
        self.impl_table.deinit();
        self.types.deinit();
    }

    pub fn addFunction(self: *Module, func: Function) FuncId {
        const id = FuncId.fromIndex(@intCast(self.functions.items.len));
        self.functions.append(self.alloc, func) catch unreachable;
        return id;
    }

    pub fn getFunction(self: *const Module, id: FuncId) *const Function {
        return &self.functions.items[id.index()];
    }

    pub fn getFunctionMut(self: *Module, id: FuncId) *Function {
        return &self.functions.items[id.index()];
    }

    pub fn addGlobal(self: *Module, global: Global) GlobalId {
        const id = GlobalId.fromIndex(@intCast(self.globals.items.len));
        self.globals.append(self.alloc, global) catch unreachable;
        return id;
    }
};

// ── ImplTable ───────────────────────────────────────────────────────────

pub const ImplKey = struct {
    protocol: TypeId,
    concrete: TypeId,
};

pub const ImplTable = struct {
    map: std.HashMap(ImplKey, []const FuncId, ImplKeyContext, 80),
    alloc: Allocator,

    pub fn init(alloc: Allocator) ImplTable {
        return .{
            .map = std.HashMap(ImplKey, []const FuncId, ImplKeyContext, 80).init(alloc),
            .alloc = alloc,
        };
    }

    pub fn deinit(self: *ImplTable) void {
        self.map.deinit();
    }

    pub fn put(self: *ImplTable, key: ImplKey, methods: []const FuncId) void {
        self.map.put(key, methods) catch unreachable;
    }

    pub fn get(self: *const ImplTable, key: ImplKey) ?[]const FuncId {
        return self.map.get(key);
    }

    const ImplKeyContext = struct {
        pub fn hash(_: ImplKeyContext, key: ImplKey) u64 {
            var h = std.hash.Wyhash.init(0);
            h.update(std.mem.asBytes(&key.protocol));
            h.update(std.mem.asBytes(&key.concrete));
            return h.final();
        }

        pub fn eql(_: ImplKeyContext, a: ImplKey, b: ImplKey) bool {
            return a.protocol == b.protocol and a.concrete == b.concrete;
        }
    };
};

// ── Builder ─────────────────────────────────────────────────────────────
// Fluent API for constructing one function at a time.

pub const Builder = struct {
    module: *Module,
    func: ?FuncId = null,
    current_block: ?BlockId = null,
    /// Running instruction counter within the current function (for Ref assignment).
    inst_counter: u32 = 0,

    pub fn init(module: *Module) Builder {
        return .{ .module = module };
    }

    // ── Function setup ──────────────────────────────────────────────

    pub fn beginFunction(self: *Builder, name: StringId, params: []const Function.Param, ret_ty: TypeId) FuncId {
        // Check if there's an existing extern stub with this name — upgrade it in-place
        for (self.module.functions.items, 0..) |*existing, i| {
            if (existing.name == name and existing.is_extern) {
                existing.is_extern = false;
                existing.linkage = .internal;
                existing.params = self.module.alloc.dupe(Function.Param, params) catch params;
                existing.ret = ret_ty;
                const id = FuncId.fromIndex(@intCast(i));
                self.func = id;
                self.inst_counter = @intCast(params.len);
                self.current_block = null;
                return id;
            }
        }
        const func = Function.init(name, params, ret_ty);
        const id = self.module.addFunction(func);
        self.func = id;
        // Reserve refs 0..N-1 for function parameters; instructions start at ref N.
        self.inst_counter = @intCast(params.len);
        self.current_block = null;
        return id;
    }

    /// Declare an extern function (no body, external linkage).
    pub fn declareExtern(self: *Builder, name: StringId, params: []const Function.Param, ret_ty: TypeId) FuncId {
        var func = Function.init(name, params, ret_ty);
        func.is_extern = true;
        func.linkage = .external;
        const id = self.module.addFunction(func);
        return id;
    }

    pub fn finalize(self: *Builder) void {
        self.func = null;
        self.current_block = null;
        self.inst_counter = 0;
    }

    // ── Blocks ──────────────────────────────────────────────────────

    pub fn appendBlock(self: *Builder, name: StringId, params: []const TypeId) BlockId {
        const f = self.currentFunc();
        const id = BlockId.fromIndex(@intCast(f.blocks.items.len));
        // Dupe params so the block owns the memory (callers may pass stack slices).
        const owned_params = if (params.len > 0)
            (self.module.alloc.dupe(TypeId, params) catch unreachable)
        else
            params;
        f.blocks.append(self.module.alloc, Block.init(name, owned_params)) catch unreachable;
        return id;
    }

    pub fn switchToBlock(self: *Builder, block: BlockId) void {
        self.current_block = block;
        // Record the starting ref index for this block
        const func = self.currentFunc();
        const blk = &func.blocks.items[block.index()];
        blk.first_ref = self.inst_counter;
    }

    /// Get the type of a previously emitted instruction Ref.
    pub fn getRefType(self: *Builder, ref: Ref) TypeId {
        if (self.func == null) return .s64;
        const func = self.currentFunc();
        const ref_idx = @intFromEnum(ref);
        // Check function parameters first (refs 0..N-1)
        if (ref_idx < func.params.len) {
            return func.params[ref_idx].ty;
        }
        for (func.blocks.items) |*block| {
            const first = block.first_ref;
            if (ref_idx >= first and ref_idx < first + @as(u32, @intCast(block.insts.items.len))) {
                return block.insts.items[ref_idx - first].ty;
            }
        }
        return .s64;
    }

    // ── Emit helpers ────────────────────────────────────────────────

    pub fn emit(self: *Builder, op: Op, ty: TypeId) Ref {
        return self.emitSpan(op, ty, .{});
    }

    fn emitSpan(self: *Builder, op: Op, ty: TypeId, span: Span) Ref {
        const block = self.currentBlock();
        const ref = Ref.fromIndex(self.inst_counter);
        self.inst_counter += 1;
        block.insts.append(self.module.alloc, .{ .op = op, .ty = ty, .span = span }) catch unreachable;
        return ref;
    }

    /// Emit an instruction with no meaningful result (terminators, stores).
    pub fn emitVoid(self: *Builder, op: Op, ty: TypeId) void {
        const block = self.currentBlock();
        self.inst_counter += 1;
        block.insts.append(self.module.alloc, .{ .op = op, .ty = ty }) catch unreachable;
    }

    // ── Constants ───────────────────────────────────────────────────

    pub fn constInt(self: *Builder, val: i64, ty: TypeId) Ref {
        return self.emit(.{ .const_int = val }, ty);
    }

    pub fn constFloat(self: *Builder, val: f64, ty: TypeId) Ref {
        return self.emit(.{ .const_float = val }, ty);
    }

    pub fn constBool(self: *Builder, val: bool) Ref {
        return self.emit(.{ .const_bool = val }, .bool);
    }

    pub fn constString(self: *Builder, val: StringId) Ref {
        return self.emit(.{ .const_string = val }, .string);
    }

    pub fn constNull(self: *Builder, ty: TypeId) Ref {
        return self.emit(.const_null, ty);
    }

    pub fn constUndef(self: *Builder, ty: TypeId) Ref {
        return self.emit(.const_undef, ty);
    }

    // ── Arithmetic ──────────────────────────────────────────────────

    pub fn add(self: *Builder, lhs: Ref, rhs: Ref, ty: TypeId) Ref {
        return self.emit(.{ .add = .{ .lhs = lhs, .rhs = rhs } }, ty);
    }

    pub fn sub(self: *Builder, lhs: Ref, rhs: Ref, ty: TypeId) Ref {
        return self.emit(.{ .sub = .{ .lhs = lhs, .rhs = rhs } }, ty);
    }

    pub fn mul(self: *Builder, lhs: Ref, rhs: Ref, ty: TypeId) Ref {
        return self.emit(.{ .mul = .{ .lhs = lhs, .rhs = rhs } }, ty);
    }

    pub fn div(self: *Builder, lhs: Ref, rhs: Ref, ty: TypeId) Ref {
        return self.emit(.{ .div = .{ .lhs = lhs, .rhs = rhs } }, ty);
    }

    // ── Comparison ──────────────────────────────────────────────────

    pub fn cmpEq(self: *Builder, lhs: Ref, rhs: Ref) Ref {
        return self.emit(.{ .cmp_eq = .{ .lhs = lhs, .rhs = rhs } }, .bool);
    }

    pub fn cmpLt(self: *Builder, lhs: Ref, rhs: Ref) Ref {
        return self.emit(.{ .cmp_lt = .{ .lhs = lhs, .rhs = rhs } }, .bool);
    }

    pub fn cmpGt(self: *Builder, lhs: Ref, rhs: Ref) Ref {
        return self.emit(.{ .cmp_gt = .{ .lhs = lhs, .rhs = rhs } }, .bool);
    }

    // ── Memory ──────────────────────────────────────────────────────

    pub fn alloca(self: *Builder, ty: TypeId) Ref {
        const ptr_ty = self.module.types.ptrTo(ty);
        return self.emit(.{ .alloca = ty }, ptr_ty);
    }

    pub fn load(self: *Builder, ptr: Ref, ty: TypeId) Ref {
        return self.emit(.{ .load = .{ .operand = ptr } }, ty);
    }

    pub fn store(self: *Builder, ptr: Ref, val: Ref) void {
        self.emitVoid(.{ .store = .{ .ptr = ptr, .val = val } }, .void);
    }

    // ── Struct ops ──────────────────────────────────────────────────

    pub fn structInit(self: *Builder, fields: []const Ref, ty: TypeId) Ref {
        const owned = self.module.alloc.dupe(Ref, fields) catch unreachable;
        return self.emit(.{ .struct_init = .{ .fields = owned } }, ty);
    }

    pub fn structGet(self: *Builder, base: Ref, field_index: u32, ty: TypeId) Ref {
        return self.emit(.{ .struct_get = .{ .base = base, .field_index = field_index } }, ty);
    }

    pub fn structGep(self: *Builder, base: Ref, field_index: u32, ty: TypeId) Ref {
        return self.emit(.{ .struct_gep = .{ .base = base, .field_index = field_index } }, ty);
    }

    pub fn structGepTyped(self: *Builder, base: Ref, field_index: u32, ty: TypeId, base_type: TypeId) Ref {
        return self.emit(.{ .struct_gep = .{ .base = base, .field_index = field_index, .base_type = base_type } }, ty);
    }

    // ── Enum ops ────────────────────────────────────────────────────

    pub fn enumInit(self: *Builder, tag: u32, payload: Ref, ty: TypeId) Ref {
        return self.emit(.{ .enum_init = .{ .tag = tag, .payload = payload } }, ty);
    }

    pub fn enumTag(self: *Builder, val: Ref, tag_ty: TypeId) Ref {
        return self.emit(.{ .enum_tag = .{ .operand = val } }, tag_ty);
    }

    // ── Optional ops ────────────────────────────────────────────────

    pub fn optionalWrap(self: *Builder, val: Ref, ty: TypeId) Ref {
        return self.emit(.{ .optional_wrap = .{ .operand = val } }, ty);
    }

    pub fn optionalUnwrap(self: *Builder, val: Ref, ty: TypeId) Ref {
        return self.emit(.{ .optional_unwrap = .{ .operand = val } }, ty);
    }

    pub fn optionalHasValue(self: *Builder, val: Ref) Ref {
        return self.emit(.{ .optional_has_value = .{ .operand = val } }, .bool);
    }

    // ── Calls ───────────────────────────────────────────────────────

    pub fn call(self: *Builder, callee: FuncId, args: []const Ref, ret_ty: TypeId) Ref {
        const owned = self.module.alloc.dupe(Ref, args) catch unreachable;
        return self.emit(.{ .call = .{ .callee = callee, .args = owned } }, ret_ty);
    }

    pub fn callClosure(self: *Builder, callee: Ref, args: []const Ref, ret_ty: TypeId) Ref {
        const owned = self.module.alloc.dupe(Ref, args) catch unreachable;
        return self.emit(.{ .call_closure = .{ .callee = callee, .args = owned } }, ret_ty);
    }

    pub fn callBuiltin(self: *Builder, builtin: inst.BuiltinId, args: []const Ref, ret_ty: TypeId) Ref {
        const owned = self.module.alloc.dupe(Ref, args) catch unreachable;
        return self.emit(.{ .call_builtin = .{ .builtin = builtin, .args = owned } }, ret_ty);
    }

    // ── Protocol ────────────────────────────────────────────────────

    pub fn protocolCallDynamic(self: *Builder, receiver: Ref, method_index: u32, args: []const Ref, ret_ty: TypeId) Ref {
        const owned = self.module.alloc.dupe(Ref, args) catch unreachable;
        return self.emit(.{ .protocol_call_dynamic = .{ .receiver = receiver, .method_index = method_index, .args = owned } }, ret_ty);
    }

    pub fn protocolErase(self: *Builder, concrete: Ref, protocol_type: TypeId) Ref {
        return self.emit(.{ .protocol_erase = .{ .concrete = concrete, .protocol_type = protocol_type } }, protocol_type);
    }

    // ── Closure ─────────────────────────────────────────────────────

    pub fn closureCreate(self: *Builder, func_id: FuncId, env: Ref, ty: TypeId) Ref {
        return self.emit(.{ .closure_create = .{ .func = func_id, .env = env } }, ty);
    }

    // ── Conversions ─────────────────────────────────────────────────

    pub fn widen(self: *Builder, operand: Ref, from: TypeId, to: TypeId) Ref {
        return self.emit(.{ .widen = .{ .operand = operand, .from = from, .to = to } }, to);
    }

    pub fn narrow(self: *Builder, operand: Ref, from: TypeId, to: TypeId) Ref {
        return self.emit(.{ .narrow = .{ .operand = operand, .from = from, .to = to } }, to);
    }

    // ── Any ─────────────────────────────────────────────────────────

    pub fn boxAny(self: *Builder, operand: Ref, source_type: TypeId) Ref {
        return self.emit(.{ .box_any = .{ .operand = operand, .source_type = source_type } }, .any);
    }

    // ── Context ─────────────────────────────────────────────────────

    pub fn contextLoad(self: *Builder, field: StringId, ty: TypeId) Ref {
        return self.emit(.{ .context_load = .{ .field = field, .value = .none } }, ty);
    }

    pub fn contextStore(self: *Builder, field: StringId, value: Ref) void {
        self.emitVoid(.{ .context_store = .{ .field = field, .value = value } }, .void);
    }

    // ── Terminators ─────────────────────────────────────────────────

    pub fn br(self: *Builder, target: BlockId, args: []const Ref) void {
        const owned = self.module.alloc.dupe(Ref, args) catch unreachable;
        self.emitVoid(.{ .br = .{ .target = target, .args = owned } }, .void);
    }

    pub fn condBr(self: *Builder, cond: Ref, then_target: BlockId, then_args: []const Ref, else_target: BlockId, else_args: []const Ref) void {
        const t_args = self.module.alloc.dupe(Ref, then_args) catch unreachable;
        const e_args = self.module.alloc.dupe(Ref, else_args) catch unreachable;
        self.emitVoid(.{ .cond_br = .{
            .cond = cond,
            .then_target = then_target,
            .then_args = t_args,
            .else_target = else_target,
            .else_args = e_args,
        } }, .void);
    }

    pub fn ret(self: *Builder, val: Ref, ty: TypeId) void {
        self.emitVoid(.{ .ret = .{ .operand = val } }, ty);
    }

    pub fn retVoid(self: *Builder) void {
        self.emitVoid(.ret_void, .void);
    }

    pub fn switchBr(self: *Builder, operand: Ref, cases: []const inst.SwitchBranch.Case, default: BlockId, default_args: []const Ref) void {
        const owned_cases = self.module.alloc.dupe(inst.SwitchBranch.Case, cases) catch unreachable;
        const owned_default_args = self.module.alloc.dupe(Ref, default_args) catch unreachable;
        self.emitVoid(.{ .switch_br = .{
            .operand = operand,
            .cases = owned_cases,
            .default = default,
            .default_args = owned_default_args,
        } }, .void);
    }

    pub fn emitUnreachable(self: *Builder) void {
        self.emitVoid(.@"unreachable", .void);
    }

    // ── Block params ───────────────────────────────────────────────

    pub fn blockParam(self: *Builder, block: BlockId, param_index: u32, ty: TypeId) Ref {
        return self.emit(.{ .block_param = .{ .block = block, .param_index = param_index } }, ty);
    }

    // ── Internal helpers ────────────────────────────────────────────

    pub fn currentFunc(self: *Builder) *Function {
        return self.module.getFunctionMut(self.func.?);
    }

    fn currentBlock(self: *Builder) *Block {
        const f = self.currentFunc();
        return &f.blocks.items[self.current_block.?.index()];
    }
};