const std = @import("std");
const Allocator = std.mem.Allocator;
const ast = @import("../ast.zig");
const Node = ast.Node;
const sx_types = @import("../types.zig");
const ir_types = @import("types.zig");
const TypeId = ir_types.TypeId;
const TypeInfo = ir_types.TypeInfo;
const TypeTable = ir_types.TypeTable;
const StringId = ir_types.StringId;

// ── AST Node → TypeId ───────────────────────────────────────────────────
// Resolve an AST type node into an IR TypeId. Used during lowering when
// we only have the parsed AST (no codegen type registry).

pub fn resolveAstType(node: ?*const Node, table: *TypeTable) TypeId {
    const n = node orelse return .s64; // no annotation → default
    return switch (n.data) {
        .type_expr => |te| resolveTypeName(te.name, table),
        .identifier => |id| resolveTypeName(id.name, table),
        .array_type_expr => |at| resolveArrayType(&at, table),
        .slice_type_expr => |st| resolveSliceType(&st, table),
        .pointer_type_expr => |pt| resolvePointerType(&pt, table),
        .many_pointer_type_expr => |mpt| resolveManyPointerType(&mpt, table),
        .optional_type_expr => |ot| resolveOptionalType(&ot, table),
        .function_type_expr => |ft| resolveFunctionType(&ft, table),
        .closure_type_expr => |ct| resolveClosureType(&ct, table),
        .tuple_type_expr => |tt| resolveTupleType(&tt, table),
        .tuple_literal => |tl| resolveTupleLiteralAsType(&tl, table),
        .parameterized_type_expr => |pt| resolveParameterizedType(&pt, table),
        .inferred_type => .s64, // inferred — default until we have type inference
        // Inline type declarations (used as field types)
        .enum_decl => |ed| resolveInlineEnum(&ed, table),
        .struct_decl => |sd| resolveInlineStruct(&sd, table),
        .union_decl => |ud| resolveInlineUnion(&ud, table),
        else => {
            std.debug.print("type_bridge: unhandled node type {s}\n", .{@tagName(n.data)});
            return .s64;
        },
    };
}

// ── types.Type → TypeId ─────────────────────────────────────────────────
// Translate an existing codegen Type value into an IR TypeId. Used when
// we have access to the codegen's resolved type info (Phase 3+).

pub fn bridgeType(ty: sx_types.Type, table: *TypeTable) TypeId {
    return switch (ty) {
        .signed => |w| switch (w) {
            8 => .s8,
            16 => .s16,
            32 => .s32,
            64 => .s64,
            else => .s64,
        },
        .unsigned => |w| switch (w) {
            8 => .u8,
            16 => .u16,
            32 => .u32,
            64 => .u64,
            else => .u64,
        },
        .f32 => .f32,
        .f64 => .f64,
        .void_type => .void,
        .boolean => .bool,
        .string_type => .string,
        .any_type => .any,
        .usize_type => .usize,
        .isize_type => .isize,
        .enum_type => |name| resolveNamedType(name, .@"enum", table),
        .struct_type => |name| resolveNamedType(name, .@"struct", table),
        .union_type => |name| resolveNamedType(name, .@"union", table),
        .array_type => |info| blk: {
            const elem = resolveTypeName(info.element_name, table);
            break :blk table.arrayOf(elem, info.length);
        },
        .slice_type => |info| blk: {
            const elem = resolveTypeName(info.element_name, table);
            break :blk table.sliceOf(elem);
        },
        .pointer_type => |info| blk: {
            const pointee = resolveTypeName(info.pointee_name, table);
            break :blk table.ptrTo(pointee);
        },
        .many_pointer_type => |info| blk: {
            const elem = resolveTypeName(info.element_name, table);
            break :blk table.manyPtrTo(elem);
        },
        .optional_type => |info| blk: {
            const child = resolveTypeName(info.child_name, table);
            break :blk table.optionalOf(child);
        },
        .vector_type => |info| blk: {
            const elem = resolveTypeName(info.element_name, table);
            break :blk table.vectorOf(elem, info.length);
        },
        .function_type => |info| blk: {
            const alloc = table.alloc;
            var param_ids = std.ArrayList(TypeId).empty;
            for (info.param_types) |pt| {
                param_ids.append(alloc, bridgeType(pt, table)) catch unreachable;
            }
            const ret_id = bridgeType(info.return_type.*, table);
            break :blk table.functionType(param_ids.items, ret_id);
        },
        .closure_type => |info| blk: {
            const alloc = table.alloc;
            var param_ids = std.ArrayList(TypeId).empty;
            for (info.param_types) |pt| {
                param_ids.append(alloc, bridgeType(pt, table)) catch unreachable;
            }
            const ret_id = bridgeType(info.return_type.*, table);
            break :blk table.closureType(param_ids.items, ret_id);
        },
        .tuple_type => |info| blk: {
            const alloc = table.alloc;
            var field_ids = std.ArrayList(TypeId).empty;
            for (info.field_types) |ft| {
                field_ids.append(alloc, bridgeType(ft, table)) catch unreachable;
            }
            var name_ids: ?[]const StringId = null;
            if (info.field_names) |names| {
                var ids = std.ArrayList(StringId).empty;
                for (names) |n| {
                    ids.append(alloc, table.internString(n)) catch unreachable;
                }
                name_ids = ids.items;
            }
            break :blk table.intern(.{ .tuple = .{
                .fields = field_ids.items,
                .names = name_ids,
            } });
        },
        .meta_type => .any, // meta types map to Any for now
    };
}

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

const NamedKind = enum { @"struct", @"enum", @"union" };

fn resolveNamedType(name: []const u8, kind: NamedKind, table: *TypeTable) TypeId {
    // Check if primitive first
    if (resolveTypePrimitive(name)) |id| return id;

    // Register as a named type
    const name_id = table.internString(name);
    return switch (kind) {
        .@"struct" => table.intern(.{ .@"struct" = .{ .name = name_id, .fields = &.{} } }),
        .@"enum" => table.intern(.{ .@"enum" = .{ .name = name_id, .variants = &.{} } }),
        .@"union" => table.intern(.{ .@"union" = .{ .name = name_id, .fields = &.{} } }),
    };
}

fn resolveTypeName(name: []const u8, table: *TypeTable) TypeId {
    // Try primitive first
    if (resolveTypePrimitive(name)) |id| return id;

    // Arbitrary bit-width integers: s1-s64, u1-u64
    if (name.len >= 2 and (name[0] == 's' or name[0] == 'u')) {
        if (std.fmt.parseInt(u8, name[1..], 10)) |width| {
            if (width >= 1 and width <= 64) {
                if (name[0] == 's') {
                    return table.intern(.{ .signed = width });
                } else {
                    return table.intern(.{ .unsigned = width });
                }
            }
        } else |_| {}
    }

    // Sentinel-terminated slice: [:0]u8 → string
    if (name.len >= 5 and name[0] == '[' and name[1] == ':') {
        if (std.mem.indexOfScalar(u8, name, ']')) |close| {
            const sentinel = name[2..close];
            const elem = name[close + 1 ..];
            if (std.mem.eql(u8, sentinel, "0") and std.mem.eql(u8, elem, "u8")) {
                return .string;
            }
        }
    }

    // Many-pointer: [*]T
    if (name.len >= 4 and name[0] == '[' and name[1] == '*' and name[2] == ']') {
        const elem = resolveTypeName(name[3..], table);
        return table.manyPtrTo(elem);
    }

    // Pointer: *T
    if (name.len >= 2 and name[0] == '*') {
        const pointee = resolveTypeName(name[1..], table);
        return table.ptrTo(pointee);
    }

    // Optional: ?T
    if (name.len >= 2 and name[0] == '?') {
        const child = resolveTypeName(name[1..], table);
        return table.optionalOf(child);
    }

    // Assume it's a named struct/enum/union type
    const name_id = table.internString(name);
    // Check if already registered (e.g., as a union from enum_decl)
    if (table.findByName(name_id)) |existing| return existing;
    return table.intern(.{ .@"struct" = .{ .name = name_id, .fields = &.{} } });
}

pub fn resolveTypePrimitive(name: []const u8) ?TypeId {
    if (name.len == 0) return null;
    // Fast path for common types
    if (std.mem.eql(u8, name, "s64")) return .s64;
    if (std.mem.eql(u8, name, "s32")) return .s32;
    if (std.mem.eql(u8, name, "s16")) return .s16;
    if (std.mem.eql(u8, name, "s8")) return .s8;
    if (std.mem.eql(u8, name, "u64")) return .u64;
    if (std.mem.eql(u8, name, "u32")) return .u32;
    if (std.mem.eql(u8, name, "u16")) return .u16;
    if (std.mem.eql(u8, name, "u8")) return .u8;
    if (std.mem.eql(u8, name, "f32")) return .f32;
    if (std.mem.eql(u8, name, "f64")) return .f64;
    if (std.mem.eql(u8, name, "bool")) return .bool;
    if (std.mem.eql(u8, name, "string")) return .string;
    if (std.mem.eql(u8, name, "void")) return .void;
    if (std.mem.eql(u8, name, "Any")) return .any;
    if (std.mem.eql(u8, name, "Type")) return .any; // Type-as-value: boxed string
    if (std.mem.eql(u8, name, "noreturn")) return .noreturn;
    if (std.mem.eql(u8, name, "usize")) return .usize;
    if (std.mem.eql(u8, name, "isize")) return .isize;
    return null;
}

fn resolveArrayType(at: *const ast.ArrayTypeExpr, table: *TypeTable) TypeId {
    const elem = resolveAstType(at.element_type, table);
    const length: u32 = switch (at.length.data) {
        .int_literal => |lit| @intCast(@as(u64, @bitCast(lit.value))),
        else => 0,
    };
    return table.arrayOf(elem, length);
}

fn resolveSliceType(st: *const ast.SliceTypeExpr, table: *TypeTable) TypeId {
    const elem = resolveAstType(st.element_type, table);
    return table.sliceOf(elem);
}

fn resolvePointerType(pt: *const ast.PointerTypeExpr, table: *TypeTable) TypeId {
    const pointee = resolveAstType(pt.pointee_type, table);
    return table.ptrTo(pointee);
}

fn resolveManyPointerType(mpt: *const ast.ManyPointerTypeExpr, table: *TypeTable) TypeId {
    const elem = resolveAstType(mpt.element_type, table);
    return table.manyPtrTo(elem);
}

fn resolveOptionalType(ot: *const ast.OptionalTypeExpr, table: *TypeTable) TypeId {
    const child = resolveAstType(ot.inner_type, table);
    return table.optionalOf(child);
}

fn resolveFunctionType(ft: *const ast.FunctionTypeExpr, table: *TypeTable) TypeId {
    const alloc = table.alloc;
    var param_ids = std.ArrayList(TypeId).empty;
    for (ft.param_types) |pt| {
        param_ids.append(alloc, resolveAstType(pt, table)) catch unreachable;
    }
    const ret_id = if (ft.return_type) |rt| resolveAstType(rt, table) else TypeId.void;
    const cc: ir_types.TypeInfo.CallConv = if (ft.call_conv == .c) .c else .default;
    return table.functionTypeCC(param_ids.items, ret_id, cc);
}

fn resolveClosureType(ct: *const ast.ClosureTypeExpr, table: *TypeTable) TypeId {
    const alloc = table.alloc;
    var param_ids = std.ArrayList(TypeId).empty;
    for (ct.param_types) |pt| {
        param_ids.append(alloc, resolveAstType(pt, table)) catch unreachable;
    }
    const ret_id = if (ct.return_type) |rt| resolveAstType(rt, table) else TypeId.void;
    return table.closureType(param_ids.items, ret_id);
}

fn resolveTupleType(tt: *const ast.TupleTypeExpr, table: *TypeTable) TypeId {
    const alloc = table.alloc;
    var field_ids = std.ArrayList(TypeId).empty;
    for (tt.field_types) |ft| {
        field_ids.append(alloc, resolveAstType(ft, table)) catch unreachable;
    }
    var name_ids: ?[]const StringId = null;
    if (tt.field_names) |names| {
        var ids = std.ArrayList(StringId).empty;
        for (names) |n| {
            ids.append(alloc, table.internString(n)) catch unreachable;
        }
        name_ids = ids.items;
    }
    return table.intern(.{ .tuple = .{
        .fields = field_ids.items,
        .names = name_ids,
    } });
}

// Treat a tuple value literal as the corresponding tuple TYPE — valid only when
// every element is itself a type expression. Non-type elements report a clear
// diagnostic and degrade to .s64 for that slot (which the snapshot will catch).
fn resolveTupleLiteralAsType(tl: *const ast.TupleLiteral, table: *TypeTable) TypeId {
    const alloc = table.alloc;
    var field_ids = std.ArrayList(TypeId).empty;
    var name_ids_list = std.ArrayList(StringId).empty;
    var any_named = false;
    for (tl.elements) |el| {
        if (!isTypeShapedAstNode(el.value, table)) {
            std.debug.print("type_bridge: tuple literal element is not a type (tag={s}) — cannot use as tuple type\n", .{@tagName(el.value.data)});
            field_ids.append(alloc, .s64) catch unreachable;
        } else {
            field_ids.append(alloc, resolveAstType(el.value, table)) catch unreachable;
        }
        if (el.name) |n| {
            any_named = true;
            name_ids_list.append(alloc, table.internString(n)) catch unreachable;
        } else {
            name_ids_list.append(alloc, table.internString("")) catch unreachable;
        }
    }
    const names: ?[]const StringId = if (any_named) name_ids_list.items else null;
    return table.intern(.{ .tuple = .{
        .fields = field_ids.items,
        .names = names,
    } });
}

// Returns true when this AST node, on its own, denotes a type rather than a
// value. Used to guard tuple-literal-as-type reinterpretation: a tuple literal
// becomes a tuple type only when every element is a type.
pub fn isTypeShapedAstNode(node: *const Node, table: *TypeTable) bool {
    return switch (node.data) {
        .type_expr,
        .pointer_type_expr,
        .many_pointer_type_expr,
        .array_type_expr,
        .slice_type_expr,
        .optional_type_expr,
        .function_type_expr,
        .closure_type_expr,
        .tuple_type_expr,
        .parameterized_type_expr,
        => true,
        .identifier => |id| table.findByName(table.internString(id.name)) != null,
        .tuple_literal => |tl| blk: {
            for (tl.elements) |el| {
                if (!isTypeShapedAstNode(el.value, table)) break :blk false;
            }
            break :blk true;
        },
        else => false,
    };
}

fn resolveParameterizedType(pt: *const ast.ParameterizedTypeExpr, table: *TypeTable) TypeId {
    // Strip module prefix (e.g. "std.Vector" → "Vector")
    const base_name = if (std.mem.lastIndexOfScalar(u8, pt.name, '.')) |dot| pt.name[dot + 1 ..] else pt.name;
    // Vector(N, T) is a built-in parameterized type
    if (std.mem.eql(u8, base_name, "Vector")) {
        if (pt.args.len == 2) {
            const length: u32 = switch (pt.args[0].data) {
                .int_literal => |lit| @intCast(@as(u64, @bitCast(lit.value))),
                else => 0,
            };
            const elem = resolveAstType(pt.args[1], table);
            return table.vectorOf(elem, length);
        }
    }
    // Generic struct instantiation — register as named type
    const name_id = table.internString(pt.name);
    return table.intern(.{ .@"struct" = .{ .name = name_id, .fields = &.{} } });
}

// ── Inline type declarations ─────────────────────────────────────────

fn resolveInlineEnum(ed: *const ast.EnumDecl, table: *TypeTable) TypeId {
    const alloc = table.alloc;
    const name_id = table.internString(ed.name);

    // Check if already registered
    if (table.findByName(name_id)) |existing| return existing;

    // Enum with payloads → tagged union
    const has_payloads = ed.variant_types.len > 0;
    if (has_payloads) {
        var fields = std.ArrayList(TypeInfo.StructInfo.Field).empty;
        for (ed.variant_names, 0..) |vn, i| {
            var field_ty: TypeId = .void;
            if (i < ed.variant_types.len) {
                if (ed.variant_types[i]) |vt| {
                    // For inline structs (__anon), rename to EnumName.variant_name
                    if (vt.data == .struct_decl) {
                        const sd = &vt.data.struct_decl;
                        if (std.mem.eql(u8, sd.name, "__anon")) {
                            const qualified = std.fmt.allocPrint(alloc, "{s}.{s}", .{ ed.name, vn }) catch "__anon";
                            const qname_id = table.internString(qualified);
                            if (table.findByName(qname_id)) |existing| {
                                field_ty = existing;
                            } else {
                                var sfields = std.ArrayList(TypeInfo.StructInfo.Field).empty;
                                for (sd.field_names, sd.field_types) |fname, ftype_node| {
                                    const fty = resolveAstType(ftype_node, table);
                                    sfields.append(alloc, .{
                                        .name = table.internString(fname),
                                        .ty = fty,
                                    }) catch unreachable;
                                }
                                const sinfo: TypeInfo = .{ .@"struct" = .{
                                    .name = qname_id,
                                    .fields = sfields.items,
                                } };
                                field_ty = table.intern(sinfo);
                                table.update(field_ty, sinfo);
                            }
                        } else {
                            field_ty = resolveAstType(vt, table);
                        }
                    } else {
                        field_ty = resolveAstType(vt, table);
                    }
                }
            }
            fields.append(alloc, .{
                .name = table.internString(vn),
                .ty = field_ty,
            }) catch unreachable;
        }
        // Resolve backing type and tag type from enum struct
        // e.g. enum struct { tag: u32; _: u32; payload: [30]u32; } { ... }
        var backing_type: ?TypeId = null;
        var tag_type: ?TypeId = null;
        if (ed.backing_type) |bt| {
            const backing_ty = resolveAstType(bt, table);
            backing_type = backing_ty;
            // Extract tag type from first field of backing struct
            const backing_info = table.get(backing_ty);
            if (backing_info == .@"struct") {
                if (backing_info.@"struct".fields.len > 0) {
                    tag_type = backing_info.@"struct".fields[0].ty;
                }
            }
        }

        // Build explicit tag values from variant_values (e.g., quit :: 0x100)
        var explicit_tag_vals: ?[]const i64 = null;
        if (ed.variant_values.len > 0) {
            var vals = std.ArrayList(i64).empty;
            for (0..ed.variant_names.len) |i| {
                if (i < ed.variant_values.len) {
                    if (ed.variant_values[i]) |vv| {
                        if (vv.data == .int_literal) {
                            vals.append(alloc, vv.data.int_literal.value) catch unreachable;
                            continue;
                        }
                    }
                }
                vals.append(alloc, @intCast(i)) catch unreachable;
            }
            explicit_tag_vals = vals.items;
        }

        const info: TypeInfo = .{ .tagged_union = .{
            .name = name_id,
            .fields = fields.items,
            .tag_type = tag_type orelse .s64, // enum unions are always tagged (default i64)
            .backing_type = backing_type,
            .explicit_tag_values = explicit_tag_vals,
        } };
        const id = table.intern(info);
        table.update(id, info);
        return id;
    }

    // Plain enum (no payloads)
    var variants = std.ArrayList(StringId).empty;
    for (ed.variant_names) |vn| {
        variants.append(alloc, table.internString(vn)) catch unreachable;
    }
    // Build explicit values for flags (power-of-2) or custom values
    var explicit_vals: ?[]const i64 = null;
    if (ed.is_flags) {
        var vals = std.ArrayList(i64).empty;
        for (ed.variant_names, 0..) |_, i| {
            if (i < ed.variant_values.len) {
                if (ed.variant_values[i]) |vv| {
                    if (vv.data == .int_literal) {
                        vals.append(alloc, vv.data.int_literal.value) catch unreachable;
                        continue;
                    }
                }
            }
            // Auto power-of-2: 1, 2, 4, 8, ...
            vals.append(alloc, @as(i64, 1) << @intCast(i)) catch unreachable;
        }
        explicit_vals = vals.items;
    } else if (ed.variant_values.len > 0) {
        var vals = std.ArrayList(i64).empty;
        for (0..ed.variant_names.len) |i| {
            if (i < ed.variant_values.len) {
                if (ed.variant_values[i]) |vv| {
                    if (vv.data == .int_literal) {
                        vals.append(alloc, vv.data.int_literal.value) catch unreachable;
                        continue;
                    }
                }
            }
            vals.append(alloc, @intCast(i)) catch unreachable;
        }
        explicit_vals = vals.items;
    }
    // Resolve backing type for sized enums (e.g. enum u32 { ... })
    var enum_backing: ?TypeId = null;
    if (ed.backing_type) |bt| {
        // Only use simple backing types (u8, u16, u32, etc.), not struct backing (enum struct)
        if (bt.data != .struct_decl) {
            enum_backing = resolveAstType(bt, table);
        }
    }

    const info: TypeInfo = .{ .@"enum" = .{
        .name = name_id,
        .variants = variants.items,
        .is_flags = ed.is_flags,
        .explicit_values = explicit_vals,
        .backing_type = enum_backing,
    } };
    const id = table.intern(info);
    table.update(id, info);
    return id;
}

fn resolveInlineStruct(sd: *const ast.StructDecl, table: *TypeTable) TypeId {
    const alloc = table.alloc;
    const name_id = table.internString(sd.name);

    if (table.findByName(name_id)) |existing| return existing;

    var fields = std.ArrayList(TypeInfo.StructInfo.Field).empty;
    for (sd.field_names, sd.field_types) |fname, ftype_node| {
        const field_ty = resolveAstType(ftype_node, table);
        fields.append(alloc, .{
            .name = table.internString(fname),
            .ty = field_ty,
        }) catch unreachable;
    }
    const info: TypeInfo = .{ .@"struct" = .{
        .name = name_id,
        .fields = fields.items,
    } };
    const id = table.intern(info);
    table.update(id, info);
    return id;
}

fn resolveInlineUnion(ud: *const ast.UnionDecl, table: *TypeTable) TypeId {
    const alloc = table.alloc;
    const name_id = table.internString(ud.name);

    if (table.findByName(name_id)) |existing| return existing;

    var fields = std.ArrayList(TypeInfo.StructInfo.Field).empty;
    for (ud.field_names, ud.field_types) |fname, ftype_node| {
        const field_ty = resolveAstType(ftype_node, table);
        fields.append(alloc, .{
            .name = table.internString(fname),
            .ty = field_ty,
        }) catch unreachable;
    }
    const info: TypeInfo = .{ .@"union" = .{
        .name = name_id,
        .fields = fields.items,
    } };
    const id = table.intern(info);
    table.update(id, info);
    return id;
}
