const std = @import("std");

const EnumFlags = @import("EnumFlags.zig").EnumFlags;

pub const FieldInfo = @import("Class/FieldInfo.zig");
pub const MethodInfo = @import("Class/MethodInfo.zig");
pub const AttributeInfo = @import("Class/AttributeInfo.zig");

allocator: std.mem.Allocator,

magic: u32,
minor_version: u16,
major_version: u16,
constant_pool: []?ConstantPoolInfo,
access_flags: ClassAccessFlags,
this_class: u16,
super_class: u16,
interfaces: []u16,
fields: []FieldInfo,
methods: []MethodInfo,
attributes: []AttributeInfo,

const Self = @This();

pub const ClassFlags = enum(u16) {
    public = 0x0001,
    final = 0x0010,
    super = 0x0020,
    interface = 0x0200,
    abstract = 0x0400,
    synthetic = 0x1000,
    annotation = 0x2000,
    @"enum" = 0x4000,
};

pub const ClassAccessFlags = EnumFlags(ClassFlags);

pub const ParseError = std.Io.Reader.Error || std.mem.Allocator.Error || error {
    InvalidMagicNumber,
    InvalidTag,
};

pub const ResolveError = error {
    InvalidConstantPoolIndex,
    InvalidConstantType,
};

pub fn parse(input: *std.Io.Reader, allocator: std.mem.Allocator) ParseError!Self {
    const magic = try input.takeInt(u32, .big);
    if (magic != 0xCAFEBABE) {
        return ParseError.InvalidMagicNumber;
    }

    const minor_version = try input.takeInt(u16, .big);
    const major_version = try input.takeInt(u16, .big);

    const constant_pool = try parseConstantPool(input, allocator);
    errdefer {
        for (constant_pool) |*cp_info| {
            if (cp_info.*) |*c| c.deinit(allocator);
        }
        allocator.free(constant_pool);
    }

    const access_flags: ClassAccessFlags = .{ .mask = try input.takeInt(u16, .big) };

    const this_class = try input.takeInt(u16, .big);
    const super_class = try input.takeInt(u16, .big);

    const interfaces = try parseInterfaces(input, allocator);
    errdefer allocator.free(interfaces);

    const fields = try parseFields(input, allocator);
    errdefer {
        for (fields) |*field| {
            field.deinit(allocator);
        }
        allocator.free(fields);
    }

    const methods = try parseMethods(input, allocator);
    errdefer {
        for (methods) |*method| {
            method.deinit(allocator);
        }
        allocator.free(methods);
    }

    const attributes = try parseAttributes(input, allocator);
    errdefer {
        for (attributes) |*attribute| {
            attribute.deinit(allocator);
        }
        allocator.free(attributes);
    }

    return .{
        .allocator = allocator,
        .magic = magic,
        .minor_version = minor_version,
        .major_version = major_version,
        .constant_pool = constant_pool,
        .access_flags = access_flags,
        .this_class = this_class,
        .super_class = super_class,
        .interfaces = interfaces,
        .fields = fields,
        .methods = methods,
        .attributes = attributes,
    };
}

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

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

        const cp_info = try ConstantPoolInfo.parse(input, allocator);
        constant_pool[i] = cp_info;

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

    return constant_pool;
}

fn parseInterfaces(input: *std.Io.Reader, allocator: std.mem.Allocator) ParseError![]u16 {
    const interfaces_count = try input.takeInt(u16, .big);
    const interfaces = try allocator.alloc(u16, interfaces_count);
    errdefer allocator.free(interfaces);

    for (0..interfaces_count) |i| {
        interfaces[i] = try input.takeInt(u16, .big);
    }

    return interfaces;
}

fn parseFields(input: *std.Io.Reader, allocator: std.mem.Allocator) ParseError![]FieldInfo {
    const fields_count = try input.takeInt(u16, .big);
    const fields = try allocator.alloc(FieldInfo, fields_count);
    errdefer allocator.free(fields);

    for (0..fields_count) |i| {
        fields[i] = try FieldInfo.parse(input, allocator);
    }

    return fields;
}

fn parseMethods(input: *std.Io.Reader, allocator: std.mem.Allocator) ParseError![]MethodInfo {
    const methods_count = try input.takeInt(u16, .big);
    const methods = try allocator.alloc(MethodInfo, methods_count);
    errdefer allocator.free(methods);

    for (0..methods_count) |i| {
        methods[i] = try MethodInfo.parse(input, allocator);
    }

    return methods;
}

pub fn parseAttributes(input: *std.Io.Reader, allocator: std.mem.Allocator) ParseError![]AttributeInfo {
    const attributes_count = try input.takeInt(u16, .big);
    const attributes = try allocator.alloc(AttributeInfo, attributes_count);
    errdefer allocator.free(attributes);

    for (0..attributes_count) |i| {
        attributes[i] = try AttributeInfo.parse(input, allocator);
    }

    return attributes;
}

pub fn deinit(self: *Self) void {
    self.freeConstantPool();
    self.freeInterfaces();
    self.freeFields();
    self.freeMethods();
    self.freeAttributes();
}

fn freeConstantPool(self: *Self) void {
    for (self.constant_pool) |*cp_info| {
        if (cp_info.*) |*c| c.deinit(self.allocator);
    }
    self.allocator.free(self.constant_pool);
}

fn freeInterfaces(self: *Self) void {
    self.allocator.free(self.interfaces);
}

fn freeFields(self: *Self) void {
    for (self.fields) |*field_info| {
        field_info.deinit(self.allocator);
    }
    self.allocator.free(self.fields);
}

fn freeMethods(self: *Self) void {
    for (self.methods) |*method_info| {
        method_info.deinit(self.allocator);
    }
    self.allocator.free(self.methods);
}

fn freeAttributes(self: *Self) void {
    for (self.attributes) |*attr| {
        attr.deinit(self.allocator);
    }
    self.allocator.free(self.attributes);
}

pub fn jsonStringify(self: *const Self, jws: *std.json.Stringify) !void {
    try jws.beginObject();

    inline for (std.meta.fields(Self)) |field| {
        if (comptime std.mem.eql(u8, field.name, "allocator")) continue;
        try jws.objectField(field.name);
        try jws.write(@field(self, field.name));
    }

    try jws.endObject();
}

pub fn format(self: *const Self, w: *std.Io.Writer) std.Io.Writer.Error!void {
    const depth = 2;
    var indent: u8 = 0;

    try w.print("magic: 0x{X}\n", .{ self.magic });

    try w.print("version: {d}.{d}\n", .{ self.major_version, self.minor_version });

    try w.print("flags: {f}\n", .{ self.access_flags });

    const this_class = self.thisClass() catch return error.WriteFailed;
    _ = try w.write("this class:\n");
    try this_class.format(w, depth, indent + 1, self.constant_pool);

    _ = try w.write("super class:\n");
    if (self.superClass() catch return error.WriteFailed) |super_class| {
        try super_class.format(w, depth, indent + 1, self.constant_pool);
    }

    _ = try w.write("interfaces:\n");
    for (self.interfaces, 0..) |interface, i| {
        indent += 1;
        defer indent -= 1;

        if (interface >= self.constant_pool.len) return error.WriteFailed;
        const cp_info = self.constant_pool[interface - 1];
        if (cp_info == null or @as(ConstantPoolInfo.Tag, cp_info.?) != .class) return error.WriteFailed;

        try w.splatByteAll(' ', depth * indent);
        try w.print("{d}:\n", .{ i });
        try cp_info.?.format(w, depth, indent + 1, self.constant_pool);
    }

    _ = try w.write("fields:\n");
    for (self.fields, 0..) |field, i| {
        indent += 1;
        defer indent -= 1;

        try w.splatByteAll(' ', depth * indent);
        try w.print("{d}:\n", .{ i });
        try field.format(w, depth, indent + 1, self.allocator, self.constant_pool);
    }

    _ = try w.write("methods:\n");
    for (self.methods, 0..) |method, i| {
        indent += 1;
        defer indent -= 1;

        try w.splatByteAll(' ', depth * indent);
        try w.print("{d}:\n", .{ i });
        try method.format(w, depth, indent + 1, self.allocator, self.constant_pool);
    }

    _ = try w.write("attributes:\n");
    for (self.attributes, 0..) |*attribute, i| {
        indent += 1;
        defer indent -= 1;

        try w.splatByteAll(' ', depth * indent);
        try w.print("{d}:\n", .{ i });
        try attribute.format(w, depth, indent + 1, self.allocator, self.constant_pool);
    }
}

pub fn thisClass(self: *const Self) ResolveError!ConstantPoolInfo {
    const this_class = self.this_class - 1;
    if (this_class >= self.constant_pool.len) return ResolveError.InvalidConstantPoolIndex;
    const cp_info = self.constant_pool[this_class];
    if (cp_info == null or @as(ConstantPoolInfo.Tag, cp_info.?) != .class) return ResolveError.InvalidConstantType;
    return cp_info.?;
}

pub fn superClass(self: *const Self) ResolveError!?ConstantPoolInfo {
    if (self.super_class == 0) {
        const this_class = (try self.thisClass()).class;
        const this_class_name = (try this_class.name(self.constant_pool)).utf8;
        if (!std.mem.eql(u8, this_class_name.bytes, "java/lang/Object")) return ResolveError.InvalidConstantType;
        return null;
    } else {
        const super_class = self.super_class - 1;
        if (super_class >= self.constant_pool.len) return ResolveError.InvalidConstantPoolIndex;
        const cp_info = self.constant_pool[super_class];
        if (cp_info == null or @as(ConstantPoolInfo.Tag, cp_info.?) != .class) return ResolveError.InvalidConstantType;

        if (self.access_flags.contains(.interface)) {
            const super_class_name = (try cp_info.?.class.name(self.constant_pool)).utf8;
            if (!std.mem.eql(u8, super_class_name.bytes, "java/lang/Object")) return ResolveError.InvalidConstantType;
        } else {
            // TODO: Assert that superclass does not have ACC_FINAL flag set.
        }
        return cp_info.?;
    }
}

pub const ConstantPoolInfo = union(ConstantPoolInfo.Tag) {
    class: 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,
    invoke_dynamic: InvokeDynamic,

    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,
        invoke_dynamic = 18,
    };

    pub const Class = @import("Class/ConstantPoolInfo/Class.zig");
    pub const FieldRef = @import("Class/ConstantPoolInfo/FieldRef.zig");
    pub const MethodRef = @import("Class/ConstantPoolInfo/MethodRef.zig");
    pub const InterfaceMethodRef = @import("Class/ConstantPoolInfo/InterfaceMethodRef.zig");
    pub const String = @import("Class/ConstantPoolInfo/String.zig");

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

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

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

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

    pub const NameAndType = @import("Class/ConstantPoolInfo/NameAndType.zig");

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

    pub const MethodHandle = @import("Class/ConstantPoolInfo/MethodHandle.zig");
    pub const MethodType = @import("Class/ConstantPoolInfo/MethodType.zig");
    pub const InvokeDynamic = @import("Class/ConstantPoolInfo/InvokeDynamic.zig");

    pub fn parse(input: *std.Io.Reader, allocator: std.mem.Allocator) ParseError!ConstantPoolInfo {
        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),
                },
            },
            .invoke_dynamic => .{
                .invoke_dynamic = .{
                    .bootstrap_method_attr_index = try input.takeInt(u16, .big),
                    .name_and_type_index = try input.takeInt(u16, .big),
                },
            },
        };
    }

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

    pub fn format(self: *const ConstantPoolInfo, w: *std.Io.Writer, depth: usize, indent: u8, constant_pool: []const ?ConstantPoolInfo) std.Io.Writer.Error!void {
        try w.splatByteAll(' ', depth * indent);
        try w.print("constant_value_type: {s}\n", .{ @tagName(self.*) });

        switch (self.*) {
            .class => |class| {
                try w.splatByteAll(' ', depth * indent);
                _ = try w.write("name:\n");
                const class_name = class.name(constant_pool) catch return error.WriteFailed;
                try class_name.format(w, depth, indent + 1, constant_pool);
            },
            .field_ref => |field_ref| {
                try w.splatByteAll(' ', depth * indent);
                _ = try w.write("class:\n");
                const class = field_ref.class(constant_pool) catch return error.WriteFailed;
                try class.format(w, depth, indent + 1, constant_pool);

                try w.splatByteAll(' ', depth * indent);
                _ = try w.write("name_and_type:\n");
                const name_and_type = field_ref.nameAndType(constant_pool) catch return error.WriteFailed;
                try name_and_type.format(w, depth, indent + 1, constant_pool);
            },
            .method_ref => |method_ref| {
                try w.splatByteAll(' ', depth * indent);
                _ = try w.write("class:\n");
                const class = method_ref.class(constant_pool) catch return error.WriteFailed;
                try class.format(w, depth, indent + 1, constant_pool);

                try w.splatByteAll(' ', depth * indent);
                _ = try w.write("name_and_type:\n");
                const name_and_type = method_ref.nameAndType(constant_pool) catch return error.WriteFailed;
                try name_and_type.format(w, depth, indent + 1, constant_pool);
            },
            .interface_method_ref => |interface_method_ref| {
                try w.splatByteAll(' ', depth * indent);
                _ = try w.write("class:\n");
                const class = interface_method_ref.class(constant_pool) catch return error.WriteFailed;
                try class.format(w, depth, indent + 1, constant_pool);

                try w.splatByteAll(' ', depth * indent);
                _ = try w.write("name_and_type:\n");
                const name_and_type = interface_method_ref.nameAndType(constant_pool) catch return error.WriteFailed;
                try name_and_type.format(w, depth, indent + 1, constant_pool);
            },
            .string => |string| {
                const s = string.string(constant_pool) catch return error.WriteFailed;
                try s.format(w, depth, indent + 1, constant_pool);
            },
            .integer => |integer| {
                try w.splatByteAll(' ', depth * indent);
                try w.print("integer: {d}\n", .{ integer.bytes });
            },
            .float => |float| {
                const f: f32 = @bitCast(float.bytes);

                try w.splatByteAll(' ', depth * indent);
                try w.print("float: {d}\n", .{ f });
            },
            .long => |long| {
                const l: u64 = (@as(u64, long.high_bytes) << 32) | long.low_bytes;

                try w.splatByteAll(' ', depth * indent);
                try w.print("long: {d}\n", .{ l });
            },
            .double => |double| {
                const l: u64 = (@as(u64, double.high_bytes) << 32) | double.low_bytes;
                const d: f64 = @bitCast(l);

                try w.splatByteAll(' ', depth * indent);
                try w.print("double: {d}\n", .{ d });
            },
            .name_and_type => |name_and_type| {
                try w.splatByteAll(' ', depth * indent);
                _ = try w.write("name:\n");
                const name = name_and_type.name(constant_pool) catch return error.WriteFailed;
                try name.format(w, depth, indent + 1, constant_pool);

                try w.splatByteAll(' ', depth * indent);
                _ = try w.write("descriptor:\n");
                const descriptor = name_and_type.descriptor(constant_pool) catch return error.WriteFailed;
                try descriptor.format(w, depth, indent + 1, 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 w.splatByteAll(' ', depth * indent);
                _ = try w.write("reference:\n");
                const reference = method_handle.reference(constant_pool) catch return error.WriteFailed;
                try reference.format(w, depth, indent + 1, constant_pool);
            },
            .method_type => |method_type| {
                try w.splatByteAll(' ', depth * indent);
                _ = try w.write("descriptor:\n");
                const descriptor = method_type.descriptor(constant_pool) catch return error.WriteFailed;
                try descriptor.format(w, depth, indent + 1, constant_pool);
            },
            .invoke_dynamic => |invoke_dynamic| {
                try w.splatByteAll(' ', depth * indent);
                _ = try w.write("name_and_type:\n");
                const name_and_type = invoke_dynamic.nameAndType(constant_pool) catch return error.WriteFailed;
                try name_and_type.format(w, depth, indent + 1, constant_pool);
            },
        }
    }
};