const std = @import("std");

const Class = @import("../Class.zig");
const AttributeInfo = Class.AttributeInfo;
const ParseError = Class.ParseError;
const ResolveError = Class.ResolveError;

constant_pool: []?Info,

const Self = @This();

pub const Info = union(Info.Tag) {
    class: Info.Class,
    field_ref: FieldRef,
    method_ref: MethodRef,
    interface_method_ref: InterfaceMethodRef,
    string: String,
    integer: Integer,
    float: Float,
    long: Long,
    double: Double,
    name_and_type: NameAndType,
    utf8: Utf8,
    method_handle: MethodHandle,
    method_type: MethodType,
    dynamic: Dynamic,
    invoke_dynamic: InvokeDynamic,
    module: Module,
    package: Package,

    pub const Tag = enum(u8) {
        class = 7,
        field_ref = 9,
        method_ref = 10,
        interface_method_ref = 11,
        string = 8,
        integer = 3,
        float = 4,
        long = 5,
        double = 6,
        name_and_type = 12,
        utf8 = 1,
        method_handle = 15,
        method_type = 16,
        dynamic = 17,
        invoke_dynamic = 18,
        module = 19,
        package = 20,

        pub fn loadable(self: Tag) bool {
            return switch (self) {
                .integer, .float, .long, .double, .class, .string, .method_handle, .method_type, .dynamic => true,
                else => false,
            };
        }
    };

    pub const Class = struct {
        name_index: u16,

        pub fn name(self: *const Info.Class, constant_pool: Self) ResolveError!Info {
            return constant_pool.getTag(self.name_index, .utf8);
        }
    };

    pub const FieldRef = struct {
        class_index: u16,
        name_and_type_index: u16,

        pub fn class(self: *const FieldRef, constant_pool: Self) ResolveError!Info {
            return constant_pool.getTag(self.class_index, .class);
        }

        pub fn nameAndType(self: *const FieldRef, constant_pool: Self) ResolveError!Info {
            return constant_pool.getTag(self.name_and_type_index, .name_and_type);
        }
    };

    pub const MethodRef = struct {
        class_index: u16,
        name_and_type_index: u16,

        pub fn class(self: *const MethodRef, constant_pool: Self) ResolveError!Info {
            return constant_pool.getTag(self.class_index, .class);
        }

        pub fn nameAndType(self: *const MethodRef, constant_pool: Self) ResolveError!Info {
            return constant_pool.getTag(self.name_and_type_index, .name_and_type);
        }
    };

    pub const InterfaceMethodRef = struct {
        class_index: u16,
        name_and_type_index: u16,

        pub fn class(self: *const InterfaceMethodRef, constant_pool: Self) ResolveError!Info {
            return constant_pool.getTag(self.class_index, .class);
        }

        pub fn nameAndType(self: *const InterfaceMethodRef, constant_pool: Self) ResolveError!Info {
            return constant_pool.getTag(self.name_and_type_index, .name_and_type);
        }
    };

    pub const String = struct {
        string_index: u16,

        pub fn string(self: *const String, constant_pool: Self) ResolveError!Info {
            return constant_pool.getTag(self.string_index, .utf8);
        }
    };

    pub const Integer = struct {
        bytes: u32,
    };

    pub const Float = struct {
        bytes: u32,

        pub fn float(self: *const Float) f32 {
            return @bitCast(self.bytes);
        }
    };

    pub const Long = struct {
        high_bytes: u32,
        low_bytes: u32,

        pub fn long(self: *const Long) u64 {
            return (@as(u64, self.high_bytes) << 32) | self.low_bytes;
        }
    };

    pub const Double = struct {
        high_bytes: u32,
        low_bytes: u32,

        pub fn double(self: *const Double) f64 {
            return @bitCast((@as(u64, self.high_bytes) << 32) | self.low_bytes);
        }
    };

    pub const NameAndType = struct {
        name_index: u16,
        descriptor_index: u16,

        pub fn name(self: *const NameAndType, constant_pool: Self) ResolveError!Info {
            return constant_pool.getTag(self.name_index, .utf8);
        }

        pub fn descriptor(self: *const NameAndType, constant_pool: Self) ResolveError!Info {
            return constant_pool.getTag(self.descriptor_index, .utf8);
        }
    };

    pub const Utf8 = struct {
        bytes: []const u8,
    };

    pub const MethodHandle = struct {
        reference_kind: Kind,
        reference_index: u16,

        pub const Kind = enum(u8) {
            get_field = 1,
            get_static = 2,
            put_field = 3,
            put_static = 4,
            invoke_virtual = 5,
            invoke_static = 6,
            invoke_special = 7,
            new_invoke_special = 8,
            invoke_interface = 9,
        };

        pub fn reference(self: *const MethodHandle, constant_pool: Self) ResolveError!Info {
            const info = try constant_pool.get(self.reference_index);
            return switch (self.reference_kind) {
                .get_field, .get_static, .put_field, .put_static => blk: {
                    if (!info.is(.field_ref)) break :blk ResolveError.InvalidConstantType;
                    break :blk info;
                },
                .invoke_virtual, .invoke_static, .invoke_special, .new_invoke_special => blk: {
                    if (!info.is(.method_ref)) break :blk ResolveError.InvalidConstantType;
                    break :blk info;
                },
                .invoke_interface => blk: {
                    if (!info.is(.interface_method_ref)) break :blk ResolveError.InvalidConstantType;
                    break :blk info;
                },
            };
        }
    };

    pub const MethodType = struct {
        descriptor_index: u16,

        pub fn descriptor(self: *const MethodType, constant_pool: Self) ResolveError!Info {
            return constant_pool.getTag(self.descriptor_index, .utf8);
        }
    };

    pub const Dynamic = struct {
        bootstrap_method_attr_index: u16,
        name_and_type_index: u16,

        pub fn nameAndType(self: *const Dynamic, constant_pool: Self) ResolveError!Info {
            return constant_pool.getTag(self.name_and_type_index, .name_and_type);
        }
    };

    pub const InvokeDynamic = struct {
        bootstrap_method_attr_index: u16,
        name_and_type_index: u16,

        pub fn nameAndType(self: *const InvokeDynamic, constant_pool: Self) ResolveError!Info {
            return constant_pool.getTag(self.name_and_type_index, .name_and_type);
        }
    };

    pub const Module = struct {
        name_index: u16,

        pub fn name(self: *const Module, constant_pool: Self) ResolveError!Info {
            return constant_pool.getTag(self.name_index, .utf8);
        }
    };

    pub const Package = struct {
        name_index: u16,

        pub fn name(self: *const Package, constant_pool: Self) ResolveError!Info {
            return constant_pool.getTag(self.name_index, .utf8);
        }
    };

    pub fn parse(input: *std.Io.Reader, allocator: std.mem.Allocator) ParseError!Info {
        const tag_byte = try input.takeByte();
        const tag = std.enums.fromInt(Tag, tag_byte) orelse return ParseError.InvalidTag;

        return switch (tag) {
            .class => .{
                .class = .{
                    .name_index = try input.takeInt(u16, .big),
                },
            },
            .field_ref => .{
                .field_ref = .{
                    .class_index = try input.takeInt(u16, .big),
                    .name_and_type_index = try input.takeInt(u16, .big),
                },
            },
            .method_ref => .{
                .method_ref = .{
                    .class_index = try input.takeInt(u16, .big),
                    .name_and_type_index = try input.takeInt(u16, .big),
                },
            },
            .interface_method_ref => .{
                .interface_method_ref = .{
                    .class_index = try input.takeInt(u16, .big),
                    .name_and_type_index = try input.takeInt(u16, .big),
                },
            },
            .string => .{
                .string = .{
                    .string_index = try input.takeInt(u16, .big),
                },
            },
            .integer => .{
                .integer = .{
                    .bytes = try input.takeInt(u32, .big),
                },
            },
            .float => .{
                .integer = .{
                    .bytes = try input.takeInt(u32, .big),
                },
            },
            .long => .{
                .long = .{
                    .high_bytes = try input.takeInt(u32, .big),
                    .low_bytes = try input.takeInt(u32, .big),
                },
            },
            .double => .{
                .double = .{
                    .high_bytes = try input.takeInt(u32, .big),
                    .low_bytes = try input.takeInt(u32, .big),
                },
            },
            .name_and_type => .{
                .name_and_type = .{
                    .name_index = try input.takeInt(u16, .big),
                    .descriptor_index = try input.takeInt(u16, .big),
                },
            },
            .utf8 => blk: {
                const length = try input.takeInt(u16, .big);
                const bytes = try input.readAlloc(allocator, length);

                break :blk .{
                    .utf8 = .{
                        .bytes = bytes,
                    },
                };
            },
            .method_handle => blk: {
                const kind_byte = try input.takeByte();
                const kind = std.enums.fromInt(MethodHandle.Kind, kind_byte) orelse return ParseError.InvalidTag;

                break :blk .{
                    .method_handle = .{
                        .reference_kind = kind,
                        .reference_index = try input.takeInt(u16, .big),
                    },
                };
            },
            .method_type => .{
                .method_type = .{
                    .descriptor_index = try input.takeInt(u16, .big),
                },
            },
            .dynamic => .{
                .dynamic = .{
                    .bootstrap_method_attr_index = try input.takeInt(u16, .big),
                    .name_and_type_index = try input.takeInt(u16, .big),
                },
            },
            .invoke_dynamic => .{
                .invoke_dynamic = .{
                    .bootstrap_method_attr_index = try input.takeInt(u16, .big),
                    .name_and_type_index = try input.takeInt(u16, .big),
                },
            },
            .module => .{
                .module = .{
                    .name_index = try input.takeInt(u16, .big),
                },
            },
            .package => .{
                .package = .{
                    .name_index = try input.takeInt(u16, .big),
                },
            },
        };
    }

    pub fn deinit(self: *Info, allocator: std.mem.Allocator) void {
        switch (self.*) {
            .utf8 => |info| allocator.free(info.bytes),
            else => {},
        }
    }

    pub fn format(self: *const Info, w: *std.Io.Writer) std.Io.Writer.Error!void {
        switch (self.*) {
            .class => |class| try w.print("#{d}", .{ class.name_index }),
            .field_ref => |field_ref| try w.print("#{d}.#{d}", .{ field_ref.class_index, field_ref.name_and_type_index }),
            .method_ref => |method_ref| try w.print("#{d}.#{d}", .{ method_ref.class_index, method_ref.name_and_type_index }),
            .interface_method_ref => |interface_method_ref| try w.print("#{d}.#{d}", .{ interface_method_ref.class_index, interface_method_ref.name_and_type_index }),
            .string => |string| try w.print("#{d}", .{ string.string_index }),
            .integer => |integer| try w.print("{d}", .{ integer.bytes }),
            .float => |float| try w.print("{d}", .{ float.float() }),
            .long => |long| try w.print("{d}", .{ long.long() }),
            .double => |double| try w.print("{d}", .{ double.double() }),
            .name_and_type => |name_and_type| try w.print("#{d}:#{d}", .{ name_and_type.name_index, name_and_type.descriptor_index }),
            .utf8 => |utf8|  try w.print("\"{s}\"", .{ utf8.bytes }),
            .method_handle => |method_handle| try w.print("{t} #{d}", .{ method_handle.reference_kind, method_handle.reference_index }),
            .method_type => |method_type| try w.print("#{d}", .{ method_type.descriptor_index }),
            .dynamic => |dynamic| try w.print("#{d} #{d}", .{ dynamic.bootstrap_method_attr_index, dynamic.name_and_type_index }),
            .invoke_dynamic => |invoke_dynamic| try w.print("#{d} #{d}", .{ invoke_dynamic.bootstrap_method_attr_index, invoke_dynamic.name_and_type_index }),
            .module => |module| try w.print("#{d}", .{ module.name_index }),
            .package => |package| try w.print("#{d}", .{ package.name_index }),
        }
    }

    pub fn indentedFormat(
        self: *const Info,
        w: *std.Io.Writer,
        depth: usize,
        indent: u8,
        constant_pool: Self,
    ) std.Io.Writer.Error!void {
        try w.splatByteAll(' ', depth * indent);
        try w.print("constant_value_type: {t}\n", .{ self.* });

        switch (self.*) {
            .class => |*class| {
                try formatInfo(w, depth, indent, "name", class, .name, constant_pool);
            },
            .field_ref => |*field_ref| {
                try formatInfo(w, depth, indent, "class", field_ref, .class, constant_pool);

                try formatInfo(w, depth, indent, "name_and_type", field_ref, .nameAndType, constant_pool);
            },
            .method_ref => |*method_ref| {
                try formatInfo(w, depth, indent, "class", method_ref, .class, constant_pool);

                try formatInfo(w, depth, indent, "name_and_type", method_ref, .nameAndType, constant_pool);
            },
            .interface_method_ref => |*interface_method_ref| {
                try formatInfo(w, depth, indent, "class", interface_method_ref, .class, constant_pool);

                try formatInfo(w, depth, indent, "name_and_type", interface_method_ref, .nameAndType, constant_pool);
            },
            .string => |*string| {
                try formatInfo(w, depth, indent, "string", string, .string, constant_pool);
            },
            .integer => |*integer| {
                try w.splatByteAll(' ', depth * indent);
                try w.print("integer: {d}\n", .{ integer.bytes });
            },
            .float => |*float| {
                try w.splatByteAll(' ', depth * indent);
                try w.print("float: {d}\n", .{ float.float() });
            },
            .long => |*long| {
                try w.splatByteAll(' ', depth * indent);
                try w.print("long: {d}\n", .{ long.long() });
            },
            .double => |*double| {
                try w.splatByteAll(' ', depth * indent);
                try w.print("double: {d}\n", .{ double.double() });
            },
            .name_and_type => |*name_and_type| {
                try formatInfo(w, depth, indent, "name", name_and_type, .name, constant_pool);

                try formatInfo(w, depth, indent, "descriptor", name_and_type, .descriptor, constant_pool);
            },
            .utf8 => |*utf8| {
                try w.splatByteAll(' ', depth * indent);
                try w.print("bytes: \"{s}\"\n", .{ utf8.bytes });
            },
            .method_handle => |*method_handle| {
                try w.splatByteAll(' ', depth * indent);
                try w.print("kind: {t}\n", .{ method_handle.reference_kind });

                try formatInfo(w, depth, indent, "descriptor", method_handle, .reference, constant_pool);
            },
            .method_type => |*method_type| {
                try formatInfo(w, depth, indent, "descriptor", method_type, .descriptor, constant_pool);
            },
            .dynamic => |*dynamic| {
                try formatInfo(w, depth, indent, "name_and_type", dynamic, .nameAndType, constant_pool);
            },
            .invoke_dynamic => |*invoke_dynamic| {
                try formatInfo(w, depth, indent, "name_and_type", invoke_dynamic, .nameAndType, constant_pool);
            },
            .module => |*module| {
                try formatInfo(w, depth, indent, "name", module, .name, constant_pool);
            },
            .package => |*package| {
                try formatInfo(w, depth, indent, "name", package, .name, constant_pool);
            },
        }
    }

    pub fn is(self: *const Info, tag: Tag) bool {
        return @as(Tag, self.*) == tag;
    }
};

pub fn parse(input: *std.Io.Reader, allocator: std.mem.Allocator) ParseError!Self {
    const constant_pool_size = try input.takeInt(u16, .big);
    const constant_pool = try allocator.alloc(?Info, constant_pool_size - 1);
    errdefer allocator.free(constant_pool);

    var i: usize = 0;
    while (i < constant_pool_size - 1) {
        defer i += 1;

        const info: Info = try .parse(input, allocator);
        constant_pool[i] = info;

        switch (info) {
            .long, .double => {
                i += 1;
                constant_pool[i] = null;
            },
            else => {},
        }
    }

    return .{
        .constant_pool = constant_pool,
    };
}

pub fn deinit(self: *Self, allocator: std.mem.Allocator) void {
    for (self.constant_pool) |*info| {
        if (info.*) |*c| c.deinit(allocator);
    }
    allocator.free(self.constant_pool);
}

pub fn get(self: *const Self, index: u16) ResolveError!Info {
    const actual_index = index - 1;
    if (actual_index >= self.constant_pool.len) return ResolveError.InvalidConstantPoolIndex;
    const info = self.constant_pool[actual_index];
    if (info == null) return ResolveError.InvalidConstantType;
    return info.?;
}

pub fn getOptional(self: *const Self, index: u16) ResolveError!?Info {
    if (index == 0) return null;
    return self.get(index);
}

pub fn getTag(self: *const Self, index: u16, expected_tag: Info.Tag) ResolveError!Info {
    const info = try self.get(index);
    if (!info.is(expected_tag)) return ResolveError.InvalidConstantType;
    return info;
}

pub fn getTagOptional(self: *const Self, index: u16, expected_tag: Info.Tag) ResolveError!?Info {
    if (index == 0) return null;
    return try self.getTag(index, expected_tag);
}

pub fn getLoadable(self: *const Self, index: u16) ResolveError!Info {
    const info = try self.get(index);
    if (!@as(Info.Tag, info).loadable()) return ResolveError.InvalidConstantType;
    return info;
}

pub fn getLoadableOptional(self: *const Self, index: u16) ResolveError!?Info {
    if (index == 0) return null;
    return self.getLoadable(index);
}

pub fn formatInfo(
    w: *std.Io.Writer,
    depth: usize,
    indent: u8,
    comptime name: []const u8,
    context: anytype,
    comptime getter: std.meta.DeclEnum(@TypeOf(context.*)),
    constant_pool: Self,
) std.Io.Writer.Error!void {
    const func = @field(@TypeOf(context.*), @tagName(getter));

    try w.splatByteAll(' ', depth * indent);
    _ = try w.write(name ++ ":\n");
    const info = func(context, constant_pool) catch return error.WriteFailed;
    try info.indentedFormat(w, depth, indent + 1, constant_pool);
}