const std = @import("std");
const builtin = @import("builtin");

const c = @import("c");

pub const ffi = @import("ffi.zig");

const log = std.log.scoped(.pjrt);

test {
    std.testing.refAllDecls(@This());
}

pub const meta = struct {
    // We could calculate it like PJRT does, but it turns out that some of those
    // were wrong in PJRT itself [1], which gets propagated to binary plugins. In
    // order to mirror that, we just the value as computed by PJRT itself, through
    // comptime reflection. We could make the argument to remove that one day since
    // [1] has been fixed. The problem is that this problem could happen again in
    // as the way PJRT does it is not very robust.
    //
    // 1. https://github.com/openxla/xla/issues/10032
    pub fn structSize(comptime T: type) usize {
        // unsafe on purpose, we want this to fail if that ever changes
        const typedef_name = comptime blk: {
            const needle = ".struct_";
            const idx = std.mem.indexOf(u8, @typeName(T), needle).?;
            break :blk @typeName(T)[idx + needle.len ..];
        };
        return @field(c, typedef_name ++ "_STRUCT_SIZE");
    }

    pub fn Struct(comptime T: type) type {
        const fields = std.meta.fields(T);
        var names: [fields.len][]const u8 = undefined;
        var types: [fields.len]type = undefined;
        var attributes: [fields.len]std.builtin.Type.StructField.Attributes = undefined;
        for (fields, &names, &types, &attributes) |field, *name, *type_, *attr| {
            name.* = field.name;
            type_.* = field.type;
            attr.* = .{
                .default_value_ptr = @ptrCast(if (std.mem.eql(u8, field.name, "struct_size"))
                    &structSize(T)
                else
                    &std.mem.zeroes(field.type)),
            };
        }
        return @Struct(
            .@"extern",
            null,
            &names,
            &types,
            &attributes,
        );
    }
};

pub const ApiError = error{
    Cancelled,
    Unknown,
    InvalidArgument,
    DeadlineExceeded,
    NotFound,
    AlreadyExists,
    PermissionDenied,
    ResourceExhausted,
    FailedPrecondition,
    Aborted,
    OutOfRange,
    Unimplemented,
    Internal,
    Unavailable,
    DataLoss,
    Unauthenticated,
};

fn InnerMixin(comptime innerT: type) type {
    return struct {
        fn inner(self: anytype) *innerT {
            return @ptrCast(@alignCast(@constCast(self)));
        }
    };
}

pub const Api = struct {
    pub const Version = struct {
        major: i64,
        minor: i64,

        pub fn format(self: Version, writer: *std.Io.Writer) std.Io.Writer.Error!void {
            try writer.print("{d}.{d}", .{ self.major, self.minor });
        }
    };

    const Funcs = std.meta.FieldEnum(c.PJRT_Api);

    inner: c.PJRT_Api,

    pub fn loadFrom(library: [:0]const u8) !*const Api {
        const basename = std.Io.Dir.path.basename(library);
        log.info("Loading: {s}...", .{basename});

        var lib: std.DynLib = switch (builtin.os.tag) {
            .linux, .macos => blk: {
                const rtld: std.c.RTLD = switch (builtin.os.tag) {
                    // We use RTLD_GLOBAL so that symbols from NEEDED libraries are available in the global namespace.
                    .linux => .{ .LAZY = true, .GLOBAL = true, .NODELETE = true },
                    .macos => .{ .LAZY = true, .LOCAL = true },
                    else => unreachable,
                };
                break :blk .{
                    .inner = .{
                        .handle = std.c.dlopen(library, rtld) orelse {
                            log.err("Unable to dlopen plugin: {s}", .{library});
                            return error.FileNotFound;
                        },
                    },
                };
            },
            else => std.DynLib.open(library) catch |err| {
                log.err("Unable to dlopen plugin: {s}", .{library});
                return err;
            },
        };

        const api = fromDynLib(&lib) catch |err| {
            log.err("Unable to load PJRT API from plugin: {s}: {}", .{ library, err });
            return err;
        };
        log.info("Loaded: {s}", .{basename});
        return api;
    }

    pub fn fromDynLib(lib: *std.DynLib) !*const Api {
        const DynGetPjrtApi = lib.lookup(*const fn () callconv(.c) *const Api, "GetPjrtApi") orelse {
            return error.MissingGetPjrtApi;
        };

        const api = DynGetPjrtApi();
        _ = try api.call(.PJRT_Plugin_Initialize, .{});

        return api;
    }

    fn PJRTFnArg(comptime func: Funcs) type {
        const fti = @typeInfo(@FieldType(c.PJRT_Api, @tagName(func)));
        const fn_ptr = @typeInfo(fti.optional.child);
        const fn_type_info = @typeInfo(fn_ptr.pointer.child);
        const arg_array_type_info = @typeInfo(fn_type_info.@"fn".params[0].type.?);
        return arg_array_type_info.pointer.child;
    }

    fn PJRTFnArgWithDefault(comptime func: Funcs) type {
        const argT = PJRTFnArg(func);
        return switch (@typeInfo(argT)) {
            .@"struct" => meta.Struct(argT),
            else => argT,
        };
    }

    inline fn innerCall(self: *const Api, comptime method: Funcs, arg: *PJRTFnArg(method)) ApiError!void {
        if (@offsetOf(c.PJRT_Api, @tagName(method)) > self.inner.struct_size) {
            std.debug.panic("PJRT Api method {s} not available in this plugin", .{@tagName(method)});
        }
        const fn_ptr = @field(&self.inner, @tagName(method)).?;
        const result = fn_ptr(arg);
        if (@TypeOf(result) == void) {
            return;
        }
        if (result) |pjrt_c_error| {
            const pjrt_error: *Error = @ptrCast(pjrt_c_error);
            log.err("[{s}] {s}", .{ @tagName(method), pjrt_error.getMessage(self) });
            return pjrt_error.getCode(self).toApiError();
        }
    }

    inline fn call(self: *const Api, comptime method: Funcs, arg: PJRTFnArgWithDefault(method)) ApiError!PJRTFnArgWithDefault(method) {
        var ret = arg;
        try innerCall(self, method, @ptrCast(&ret));
        return ret;
    }

    pub const Extensions = struct {
        pub const Type = enum(c.PJRT_Extension_Type) {
            custom_call = c.PJRT_Extension_Type_Gpu_Custom_Call,
            profiler = c.PJRT_Extension_Type_Profiler,
            stream = c.PJRT_Extension_Type_Stream,
            layouts = c.PJRT_Extension_Type_Layouts,
            ffi = c.PJRT_Extension_Type_FFI,
            memory_descriptions = c.PJRT_Extension_Type_MemoryDescriptions,
            triton = c.PJRT_Extension_Type_Triton,
            raw_buffer = c.PJRT_Extension_Type_RawBuffer,
            phase_compile = c.PJRT_Extension_Type_PhaseCompile,
            unknown = c.PJRT_Extension_Type_Unknown,
        };

        pub const Extension = union(Type) {
            custom_call: *const c.PJRT_Gpu_Custom_Call,
            profiler: *const c.PJRT_Profiler_Extension,
            stream: *const c.PJRT_Stream_Extension,
            layouts: *const c.PJRT_Layouts_Extension,
            ffi: *const c.PJRT_FFI_Extension,
            memory_descriptions: *const c.PJRT_MemoryDescriptions_Extension,
            triton: *const c.PJRT_Triton_Extension,
            raw_buffer: *const c.PJRT_RawBuffer_Extension,
            phase_compile: *const c.PJRT_PhaseCompile_Extension,
            unknown: *const c.PJRT_Extension_Base,
        };

        pub const Iterator = struct {
            current: ?*const c.PJRT_Extension_Base,

            pub fn next(self: *Iterator) ?Extension {
                defer if (self.current) |cur| {
                    self.current = cur.next;
                };
                if (self.current) |cur| {
                    if (std.enums.fromInt(Extensions.Type, cur.type)) |e| {
                        return switch (e) {
                            inline else => |t| @unionInit(Extension, @tagName(t), @ptrCast(self.current)),
                        };
                    }
                    return .{ .unknown = cur };
                }
                return null;
            }
        };
    };

    pub fn extensions(self: *const Api) Extensions.Iterator {
        return .{ .current = @ptrCast(@alignCast(self.inner.extension_start)) };
    }

    pub fn extension(self: *const Api, ext_type: Extensions.Type) ?Extensions.Extension {
        var it = self.extensions();
        while (it.next()) |ext| {
            if (std.meta.activeTag(ext) == ext_type) {
                return ext;
            }
        }
        return null;
    }

    pub inline fn version(self: *const Api) Version {
        return .{
            .major = @intCast(self.inner.pjrt_api_version.major_version),
            .minor = @intCast(self.inner.pjrt_api_version.minor_version),
        };
    }

    pub fn stablehloCurrentVersion(self: *const Api) ?[]const u8 {
        const state = struct {
            var buf: [32]u8 = undefined;
            var str: ?[:0]const u8 = null;
        };
        if (state.str) |str| {
            return str;
        }
        if (self.pluginAttribute("stablehlo_current_version")) |nv| {
            switch (nv.value()) {
                .int64list => |v| {
                    state.str = std.fmt.bufPrintZ(&state.buf, "{d}.{d}.{d}", .{ v[0], v[1], v[2] }) catch unreachable;
                },
                else => unreachable,
            }
        }
        return state.str;
    }

    pub fn createExecuteContext(api: *const Api) ApiError!*ExecuteContext {
        const ret = try api.call(.PJRT_ExecuteContext_Create, .{});
        return @ptrCast(ret.context.?);
    }

    pub fn ffi(api: *const Api) ?Ffi {
        if (api.extension(.ffi)) |ext| {
            return .{ .inner = ext.ffi };
        }
        return null;
    }

    pub fn profiler(self: *const Api, options_pb: []const u8) ApiError!?Profiler {
        if (self.extension(.profiler)) |ext| {
            return try Profiler.init(self, ext.profiler, options_pb);
        }

        return null;
    }

    pub fn pluginAttribute(api: *const Api, key: []const u8) ?NamedValue {
        const attributes = api.pluginAttributes();
        for (attributes) |attr| {
            if (std.mem.eql(u8, attr.name(), key)) {
                return attr;
            }
        }
        return null;
    }

    pub fn pluginAttributes(api: *const Api) []const NamedValue {
        const ret = api.call(.PJRT_Plugin_Attributes, .{
            .extension_start = null,
        }) catch unreachable;

        if (ret.attributes == null) return &.{};

        return @ptrCast(ret.attributes[0..ret.num_attributes]);
    }
};

pub const ErrorCode = enum(c.PJRT_Error_Code) {
    cancelled = c.PJRT_Error_Code_CANCELLED,
    unknown = c.PJRT_Error_Code_UNKNOWN,
    invalid_argument = c.PJRT_Error_Code_INVALID_ARGUMENT,
    deadline_exceeded = c.PJRT_Error_Code_DEADLINE_EXCEEDED,
    not_found = c.PJRT_Error_Code_NOT_FOUND,
    already_exists = c.PJRT_Error_Code_ALREADY_EXISTS,
    permission_denied = c.PJRT_Error_Code_PERMISSION_DENIED,
    resource_exhausted = c.PJRT_Error_Code_RESOURCE_EXHAUSTED,
    failed_precondition = c.PJRT_Error_Code_FAILED_PRECONDITION,
    aborted = c.PJRT_Error_Code_ABORTED,
    out_of_range = c.PJRT_Error_Code_OUT_OF_RANGE,
    unimplemented = c.PJRT_Error_Code_UNIMPLEMENTED,
    internal = c.PJRT_Error_Code_INTERNAL,
    unavailable = c.PJRT_Error_Code_UNAVAILABLE,
    data_loss = c.PJRT_Error_Code_DATA_LOSS,
    unauthenticated = c.PJRT_Error_Code_UNAUTHENTICATED,

    pub fn toApiError(code: ErrorCode) ApiError {
        return switch (code) {
            .cancelled => ApiError.Cancelled,
            .unknown => ApiError.Unknown,
            .invalid_argument => ApiError.InvalidArgument,
            .deadline_exceeded => ApiError.DeadlineExceeded,
            .not_found => ApiError.NotFound,
            .already_exists => ApiError.AlreadyExists,
            .permission_denied => ApiError.PermissionDenied,
            .resource_exhausted => ApiError.ResourceExhausted,
            .failed_precondition => ApiError.FailedPrecondition,
            .aborted => ApiError.Aborted,
            .out_of_range => ApiError.OutOfRange,
            .unimplemented => ApiError.Unimplemented,
            .internal => ApiError.Internal,
            .unavailable => ApiError.Unavailable,
            .data_loss => ApiError.DataLoss,
            .unauthenticated => ApiError.Unauthenticated,
        };
    }
};

pub const Error = opaque {
    const inner = InnerMixin(c.PJRT_Error).inner;

    pub fn deinit(self: *Error, api: *const Api) void {
        _ = api.call(.PJRT_Error_Destroy, .{
            .@"error" = self.inner(),
        }) catch unreachable;
    }

    pub fn getCode(self: *Error, api: *const Api) ErrorCode {
        const ret = api.call(.PJRT_Error_GetCode, .{
            .@"error" = self.inner(),
        }) catch unreachable;
        return @enumFromInt(ret.code);
    }

    pub fn getMessage(self: *Error, api: *const Api) []const u8 {
        const ret = api.call(.PJRT_Error_Message, .{
            .@"error" = self.inner(),
        }) catch unreachable;
        return ret.message[0..ret.message_size];
    }
};

pub const ClientInitError = error{LoadingFailed} || ApiError;

pub const ShapeSpec = extern struct {
    comptime {
        std.debug.assert(@sizeOf(ShapeSpec) == @sizeOf(c.PJRT_ShapeSpec));
    }

    inner: meta.Struct(c.PJRT_ShapeSpec),

    pub fn init(dims_: []const i64, bt: BufferType) ShapeSpec {
        return .{
            .inner = .{
                .dims = @ptrCast(@constCast(dims_)),
                .num_dims = dims_.len,
                .element_type = @intFromEnum(bt),
            },
        };
    }

    pub fn dims(self: ShapeSpec) []const i64 {
        return self.inner.dims[0..self.inner.num_dims];
    }

    pub fn bufferType(self: ShapeSpec) BufferType {
        return @enumFromInt(self.inner.element_type);
    }
};

pub const Stream = opaque {};

pub const Client = opaque {
    const inner = InnerMixin(c.PJRT_Client).inner;

    pub const ProgramFormat = enum {
        hlo,
        mlir,
    };

    pub fn init(api: *const Api, create_options: []const NamedValue) ClientInitError!*Client {
        // log.info("Loaded PJRT runtime plugin: {s}", .{api.Platform});
        const ret = try api.call(.PJRT_Client_Create, .{
            .create_options = @ptrCast(create_options),
            .num_options = create_options.len,
        });
        return @ptrCast(ret.client.?);
    }

    pub fn deinit(self: *Client, api: *const Api) void {
        _ = api.call(.PJRT_Client_Destroy, .{
            .client = self.inner(),
        }) catch {};
    }

    pub fn platformName(self: *const Client, api: *const Api) []const u8 {
        const ret = api.call(.PJRT_Client_PlatformName, .{
            .client = self.inner(),
        }) catch unreachable;
        return ret.platform_name[0..ret.platform_name_size];
    }

    pub fn devices(self: *const Client, api: *const Api) []const *Device {
        const ret = api.call(.PJRT_Client_Devices, .{
            .client = self.inner(),
        }) catch unreachable;
        return @ptrCast(ret.devices[0..ret.num_devices]);
    }

    pub fn addressableDevices(self: *const Client, api: *const Api) []const *Device {
        const ret = api.call(.PJRT_Client_AddressableDevices, .{
            .client = self.inner(),
        }) catch unreachable;
        return @ptrCast(ret.addressable_devices[0..ret.num_addressable_devices]);
    }

    pub const CompileArgs = struct {
        bytecode: []const u8,
        bytecode_format: ProgramFormat,
        compile_options_pb: []const u8,
    };

    pub fn compileRaw(self: *const Client, api: *const Api, args: CompileArgs) ApiError!*LoadedExecutable {
        const bytecode_format_ = @tagName(args.bytecode_format);
        const ret = try api.call(.PJRT_Client_Compile, .{
            .program = @ptrCast(&meta.Struct(c.PJRT_Program){
                .code = @ptrCast(@constCast(args.bytecode)),
                .code_size = args.bytecode.len,
                .format = @ptrCast(@constCast(bytecode_format_)),
                .format_size = bytecode_format_.len,
            }),
            .compile_options = @ptrCast(@constCast(args.compile_options_pb)),
            .compile_options_size = args.compile_options_pb.len,
            .client = self.inner(),
        });
        return @ptrCast(ret.executable.?);
    }

    pub fn compile(self: *const Client, api: *const Api, io: std.Io, args: CompileArgs) ApiError!*LoadedExecutable {
        var future = io.async(compileRaw, .{ self, api, args });
        return future.await(io);
    }

    pub const BufferFromHostBufferArgs = struct {
        data: [*]const u8,
        buffer_type: BufferType,
        dims: []const i64,
        byte_strides: ?[]const i64,
        layout: MemoryLayout,
        host_buffer_semantics: HostBufferSemantics,
        dst: union(enum) {
            device: *const Device,
            memory: *const Memory,
        },
    };

    pub fn bufferFromHostBuffer(self: *const Client, api: *const Api, args: BufferFromHostBufferArgs) ApiError!struct { *Buffer, ?*Event } {
        const ret = try api.call(.PJRT_Client_BufferFromHostBuffer, .{
            .client = self.inner(),
            .data = @constCast(args.data),
            .type = @intFromEnum(args.buffer_type),
            .dims = @ptrCast(@constCast(args.dims)),
            .num_dims = args.dims.len,
            .byte_strides = if (args.byte_strides) |bs| @ptrCast(@constCast(bs)) else null,
            .num_byte_strides = if (args.byte_strides) |bs| bs.len else 0,
            .device_layout = @ptrCast(@constCast(&args.layout.toCStruct())),
            .host_buffer_semantics = @intFromEnum(args.host_buffer_semantics),
            .device = if (args.dst == .device) @ptrCast(@constCast(args.dst.device)) else null,
            .memory = if (args.dst == .memory) args.dst.memory.inner() else null,
        });

        return .{
            @ptrCast(ret.buffer.?),
            @ptrCast(ret.done_with_host_buffer),
        };
    }

    pub fn deserializeAndLoad(self: *const Client, api: *const Api, bytes: []const u8) ApiError!*LoadedExecutable {
        const ret = try api.call(.PJRT_Executable_DeserializeAndLoad, .{
            .client = self.inner(),
            .serialized_executable = bytes.ptr,
            .serialized_executable_size = bytes.len,
        });
        return @ptrCast(ret.loaded_executable.?);
    }

    pub const CreateViewOfDeviceBufferArgs = struct {
        data: *anyopaque,
        dims: []const i64,
        element_type: BufferType,
        layout: MemoryLayout,
        device: ?*const Device = null,
        on_delete_callback: *const fn (device_buffer_ptr: ?*anyopaque, ctx: ?*anyopaque) callconv(.c) void = &struct {
            fn call(_: ?*anyopaque, _: ?*anyopaque) callconv(.c) void {}
        }.call,
        on_delete_callback_arg: ?*anyopaque = null,
        stream: ?*const Stream = null,
    };

    pub fn createViewOfDeviceBuffer(self: *const Client, api: *const Api, args: CreateViewOfDeviceBufferArgs) ApiError!*Buffer {
        const layout = args.layout.toCStruct();
        const ret = try api.call(.PJRT_Client_CreateViewOfDeviceBuffer, .{
            .client = self.inner(),
            .device_buffer_ptr = @constCast(args.data),
            .dims = args.dims.ptr,
            .num_dims = args.dims.len,
            .element_type = @intFromEnum(args.element_type),
            .layout = @ptrCast(@constCast(&layout)),
            .device = @ptrCast(@constCast(args.device)),
            .on_delete_callback = args.on_delete_callback,
            .on_delete_callback_arg = args.on_delete_callback_arg,
            .stream = @bitCast(@intFromPtr(args.stream)),
        });
        return @ptrCast(ret.buffer.?);
    }

    pub fn addressableMemories(self: *const Client, api: *const Api) []*const Memory {
        const ret = api.call(.PJRT_Client_AddressableMemories, .{
            .client = self.inner(),
        }) catch return &.{};
        if (ret.addressable_memories) |memories| {
            return @ptrCast(@constCast(memories[0..ret.num_addressable_memories]));
        }
        return &.{};
    }

    pub fn dmaMap(self: *const Client, api: *const Api, data: []const u8) ApiError!void {
        _ = try api.call(.PJRT_Client_DmaMap, .{
            .client = self.inner(),
            .data = @ptrCast(@constCast(data)),
            .size = @intCast(data.len),
        });
    }

    pub fn dmaUnmap(self: *const Client, api: *const Api, data: []const u8) ApiError!void {
        _ = try api.call(.PJRT_Client_DmaUnmap, .{
            .client = self.inner(),
            .data = @ptrCast(@constCast(data)),
        });
    }

    pub const CreateBuffersForAsyncHostToDeviceArgs = struct {
        shape_specs: []const ShapeSpec,
        device_layouts: ?[]*const MemoryLayout = null,
        memory: *const Memory,
    };

    pub fn createBuffersForAsyncHostToDevice(self: *const Client, api: *const Api, args: CreateBuffersForAsyncHostToDeviceArgs) ApiError!*AsyncHostToDeviceTransferManager {
        const ret = try api.call(.PJRT_Client_CreateBuffersForAsyncHostToDevice, .{
            .client = self.inner(),
            .shape_specs = @ptrCast(@constCast(args.shape_specs)),
            .num_shape_specs = args.shape_specs.len,
            .device_layouts = if (args.device_layouts) |layouts| @ptrCast(@constCast(layouts)) else null,
            .num_device_layouts = if (args.device_layouts) |layouts| layouts.len else 0,
            .memory = args.memory.inner(),
        });
        return @ptrCast(ret.transfer_manager.?);
    }

    pub fn defaultMemoryLayout(
        self: *const Client,
        api: *const Api,
        element_type: BufferType,
        dims: []const i64,
    ) !DefaultMemoryLayout {
        const ext = if (api.extension(.layouts)) |e| e.layouts else return error.LayoutsExtensionUnavailable;

        var get_args: c.PJRT_Layouts_PJRT_Client_GetDefaultLayout_Args = .{
            .struct_size = c.PJRT_Layouts_PJRT_Client_GetDefaultLayout_Args_STRUCT_SIZE,
            .extension_start = null,
            .client = self.inner(),
            .type = @intFromEnum(element_type),
            .dims = dims.ptr,
            .num_dims = dims.len,
            .layout = null,
        };
        if (ext.PJRT_Layouts_PJRT_Client_GetDefaultLayout.?(&get_args)) |pjrt_c_error| {
            const pjrt_error: *Error = @ptrCast(pjrt_c_error);
            defer pjrt_error.deinit(api);
            return pjrt_error.getCode(api).toApiError();
        }
        defer {
            var destroy_args: c.PJRT_Layouts_MemoryLayout_Destroy_Args = .{
                .struct_size = c.PJRT_Layouts_MemoryLayout_Destroy_Args_STRUCT_SIZE,
                .extension_start = null,
                .layout = get_args.layout,
            };
            if (ext.PJRT_Layouts_MemoryLayout_Destroy.?(&destroy_args)) |pjrt_c_error| {
                const pjrt_error: *Error = @ptrCast(pjrt_c_error);
                defer pjrt_error.deinit(api);
                log.err("[PJRT_Layouts_MemoryLayout_Destroy] {s}", .{pjrt_error.getMessage(api)});
            }
        }

        var serialize_args: c.PJRT_Layouts_MemoryLayout_Serialize_Args = .{
            .struct_size = c.PJRT_Layouts_MemoryLayout_Serialize_Args_STRUCT_SIZE,
            .extension_start = null,
            .layout = get_args.layout,
            .serialized_bytes = null,
            .serialized_bytes_size = 0,
            .serialized_layout = null,
            .serialized_layout_deleter = null,
        };
        if (ext.PJRT_Layouts_MemoryLayout_Serialize.?(&serialize_args)) |pjrt_c_error| {
            const pjrt_error: *Error = @ptrCast(pjrt_c_error);
            defer pjrt_error.deinit(api);
            return pjrt_error.getCode(api).toApiError();
        }
        defer if (serialize_args.serialized_layout_deleter) |deleter| {
            deleter(serialize_args.serialized_layout);
        };

        const serialized = serialize_args.serialized_bytes[0..serialize_args.serialized_bytes_size];
        return try DefaultMemoryLayout.parseSerialized(serialized);
    }

    pub const CreateUninitializedBufferArgs = struct {
        dims: []const i64,
        element_type: BufferType,
        layout: MemoryLayout,
        dst: union(enum) {
            device: *const Device,
            memory: *const Memory,
        },
    };

    pub fn createUninitializedBuffer(self: *const Client, api: *const Api, args: CreateUninitializedBufferArgs) ApiError!*Buffer {
        var layout = args.layout.toCStruct();
        const ret = try api.call(.PJRT_Client_CreateUninitializedBuffer, .{
            .client = self.inner(),
            .shape_dims = args.dims.ptr,
            .shape_num_dims = args.dims.len,
            .shape_element_type = @intFromEnum(args.element_type),
            .shape_layout = @ptrCast(&layout),
            .device = if (args.dst == .device) @ptrCast(@constCast(args.dst.device)) else null,
            .memory = if (args.dst == .memory) args.dst.memory.inner() else null,
        });
        return @ptrCast(ret.buffer.?);
    }
};

pub const Device = opaque {
    const inner = InnerMixin(c.PJRT_Device).inner;

    pub fn getDescription(self: *const Device, api: *const Api) *const DeviceDescription {
        const ret = api.call(.PJRT_Device_GetDescription, .{
            .device = self.inner(),
        }) catch unreachable;
        return @ptrCast(ret.device_description.?);
    }

    pub fn isAddressable(self: *const Device, api: *const Api) bool {
        const ret = api.call(.PJRT_Device_IsAddressable, .{
            .device = self.inner(),
        }) catch unreachable;
        return ret.is_addressable;
    }

    pub fn localHardwareId(self: *const Device, api: *const Api) usize {
        const ret = api.call(.PJRT_Device_LocalHardwareId, .{
            .device = self.inner(),
        }) catch unreachable;
        return @intCast(ret.local_hardware_id);
    }

    pub fn addressableMemories(self: *const Device, api: *const Api) []const *Memory {
        const ret = api.call(
            .PJRT_Device_AddressableMemories,
            .{ .device = self.inner() },
        ) catch return &.{};
        return @ptrCast(ret.memories[0..ret.num_memories]);
    }

    pub fn addressableMemory(client: *const Client, api: *const Api, kind: Memory.Kind) ?*const Memory {
        for (client.addressableMemories(api)) |mem| {
            if (mem.kind(api) == kind) {
                return mem;
            }
        }
        return null;
    }

    pub fn defaultMemory(self: *const Device, api: *const Api) *const Memory {
        const ret = api.call(.PJRT_Device_DefaultMemory, .{
            .device = self.inner(),
        }) catch unreachable;
        return @ptrCast(ret.memory);
    }

    pub const MemoryStats = struct {
        /// Number of bytes in use.
        bytes_in_use: u64,
        /// The peak bytes in use.
        peak_bytes_in_use: ?u64,
        /// Number of allocations.
        num_allocs: ?u64,
        /// The largest single allocation seen.
        largest_alloc_size: ?u64,
        /// The upper limit of user-allocatable device memory in bytes.
        bytes_limit: ?u64,
        /// Number of bytes reserved.
        bytes_reserved: ?u64,
        /// The peak number of bytes reserved.
        peak_bytes_reserved: ?u64,
        /// The upper limit on the number bytes of reservable memory.
        bytes_reservable_limit: ?u64,
        /// Largest free block size in bytes.
        largest_free_block_bytes: ?u64,
        /// Number of bytes of memory held by the allocator. This may be higher than
        /// bytes_in_use if the allocator holds a pool of memory (e.g. BFCAllocator).
        pool_bytes: ?u64,
        peak_pool_bytes: ?u64,

        pub const zeroes = std.mem.zeroes(MemoryStats);

        fn fromCStruct(v: meta.Struct(c.PJRT_Device_MemoryStats_Args)) MemoryStats {
            return .{
                .bytes_in_use = @intCast(v.bytes_in_use),
                .peak_bytes_in_use = if (v.peak_bytes_in_use_is_set) @intCast(v.peak_bytes_in_use) else null,
                .num_allocs = if (v.num_allocs_is_set) @intCast(v.num_allocs) else null,
                .largest_alloc_size = if (v.largest_alloc_size_is_set) @intCast(v.largest_alloc_size) else null,
                .bytes_limit = if (v.bytes_limit_is_set) @intCast(v.bytes_limit) else null,
                .bytes_reserved = if (v.bytes_reserved_is_set) @intCast(v.bytes_reserved) else null,
                .peak_bytes_reserved = if (v.peak_bytes_reserved_is_set) @intCast(v.peak_bytes_reserved) else null,
                .bytes_reservable_limit = if (v.bytes_reservable_limit_is_set) @intCast(v.bytes_reservable_limit) else null,
                .largest_free_block_bytes = if (v.largest_free_block_bytes_is_set) @intCast(v.largest_free_block_bytes) else null,
                .pool_bytes = if (v.pool_bytes_is_set) @intCast(v.pool_bytes) else null,
                .peak_pool_bytes = if (v.peak_pool_bytes_is_set) @intCast(v.peak_pool_bytes) else null,
            };
        }
    };

    pub fn memoryStats(self: *const Device, api: *const Api) ApiError!MemoryStats {
        const ret = try api.call(.PJRT_Device_MemoryStats, .{
            .device = self.inner(),
        });
        return .fromCStruct(@bitCast(ret));
    }
};

pub const DeviceDescription = opaque {
    const inner = InnerMixin(c.PJRT_DeviceDescription).inner;

    pub fn id(self: *const DeviceDescription, api: *const Api) usize {
        const ret = api.call(.PJRT_DeviceDescription_Id, .{
            .device_description = self.inner(),
        }) catch unreachable;
        return @intCast(ret.id);
    }

    pub fn processIndex(self: *const DeviceDescription, api: *const Api) usize {
        const ret = api.call(.PJRT_DeviceDescription_ProcessIndex, .{
            .device_description = self.inner(),
        }) catch unreachable;
        return @intCast(ret.process_index);
    }

    pub fn kind(self: *const DeviceDescription, api: *const Api) []const u8 {
        const ret = api.call(.PJRT_DeviceDescription_Kind, .{
            .device_description = self.inner(),
        }) catch unreachable;
        return ret.device_kind[0..ret.device_kind_size];
    }

    pub fn debugString(self: *const DeviceDescription, api: *const Api) []const u8 {
        const ret = api.call(.PJRT_DeviceDescription_DebugString, .{
            .device_description = self.inner(),
        }) catch unreachable;
        return ret.debug_string[0..ret.debug_string_size];
    }

    pub fn toString(self: *const DeviceDescription, api: *const Api) []const u8 {
        const ret = api.call(.PJRT_DeviceDescription_ToString, .{
            .device_description = self.inner(),
        }) catch unreachable;
        return ret.to_string[0..ret.to_string_size];
    }

    pub fn attributes(self: *const DeviceDescription, api: *const Api) []const NamedValue {
        const ret = api.call(.PJRT_DeviceDescription_Attributes, .{
            .device_description = self.inner(),
        }) catch unreachable;

        if (ret.attributes == null) return &.{};

        return @ptrCast(ret.attributes[0..ret.num_attributes]);
    }

    pub fn attribute(self: *const DeviceDescription, api: *const Api, name: []const u8) ?NamedValue.Value {
        for (self.attributes(api)) |attr| {
            if (std.mem.eql(u8, attr.name(), name)) {
                return attr.value();
            }
        }
        return null;
    }
};

pub const GetCostAnalysisError = std.mem.Allocator.Error || ApiError;

pub const SerializeResult = struct {
    bytes: []const u8,
    handle: *anyopaque,
    deleter: *const fn (?*anyopaque) callconv(.c) void,

    pub fn deinit(self: *SerializeResult) void {
        self.deleter(self.handle);
        self.bytes = &.{};
        self.* = undefined;
    }
};

pub const ExecuteContext = opaque {
    pub fn deinit(self: *ExecuteContext, api: *const Api) void {
        _ = api.call(.PJRT_ExecuteContext_Destroy, .{
            .context = @ptrCast(self),
        }) catch {};
    }
};

pub const Executable = opaque {
    const inner = InnerMixin(c.PJRT_Executable).inner;

    pub fn deinit(self: *Executable, api: *const Api) void {
        _ = api.call(.PJRT_Executable_Destroy, .{
            .executable = self.inner(),
        }) catch unreachable;
    }

    pub fn costAnalysis(self: *const Executable, api: *const Api) GetCostAnalysisError![]const NamedValue {
        const ret = try api.call(.PJRT_Executable_GetCostAnalysis, .{
            .executable = self.inner(),
        });
        const values: [*]const NamedValue = @ptrCast(ret.properties);
        return values[0..ret.num_properties];
    }

    pub fn serialize(self: *const Executable, api: *const Api) ApiError!SerializeResult {
        const ret = try api.call(.PJRT_Executable_Serialize, .{
            .executable = self.inner(),
        });

        return .{
            .bytes = ret.serialized_bytes[0..ret.serialized_bytes_size],
            .handle = ret.serialized_executable.?,
            .deleter = @ptrCast(ret.serialized_executable_deleter.?),
        };
    }

    pub const CompiledMemoryStats = struct {
        /// Mirrors xla::CompiledMemoryStats.
        /// Device default memory (e.g., HBM for GPU/TPU) usage stats.
        generated_code_size_in_bytes: u64,
        argument_size_in_bytes: u64,
        output_size_in_bytes: u64,
        /// much: How argument is reused for output.
        alias_size_in_bytes: u64,
        temp_size_in_bytes: u64,

        /// memory: Host usage stats.
        host_generated_code_size_in_bytes: u64,
        host_argument_size_in_bytes: u64,
        host_output_size_in_bytes: u64,
        host_alias_size_in_bytes: u64,
        host_temp_size_in_bytes: u64,
    };

    pub fn getCompiledMemoryStats(self: *const Executable, api: *const Api) ApiError!CompiledMemoryStats {
        const ret = try api.call(.PJRT_Executable_GetCompiledMemoryStats, .{
            .executable = self.inner(),
        });
        return .{
            .generated_code_size_in_bytes = @intCast(ret.generated_code_size_in_bytes),
            .argument_size_in_bytes = @intCast(ret.argument_size_in_bytes),
            .output_size_in_bytes = @intCast(ret.output_size_in_bytes),
            .alias_size_in_bytes = @intCast(ret.alias_size_in_bytes),
            .temp_size_in_bytes = @intCast(ret.temp_size_in_bytes),
            .host_generated_code_size_in_bytes = @intCast(ret.host_generated_code_size_in_bytes),
            .host_argument_size_in_bytes = @intCast(ret.host_argument_size_in_bytes),
            .host_output_size_in_bytes = @intCast(ret.host_output_size_in_bytes),
            .host_alias_size_in_bytes = @intCast(ret.host_alias_size_in_bytes),
            .host_temp_size_in_bytes = @intCast(ret.host_temp_size_in_bytes),
        };
    }
};

pub const LoadedExecutable = opaque {
    const inner = InnerMixin(c.PJRT_LoadedExecutable).inner;

    pub fn deinit(self: *LoadedExecutable, api: *const Api) void {
        _ = api.call(.PJRT_LoadedExecutable_Destroy, .{
            .executable = self.inner(),
        }) catch {};
    }

    pub fn delete(self: *LoadedExecutable, api: *const Api) void {
        _ = api.call(.PJRT_LoadedExecutable_Delete, .{
            .executable = self.inner(),
        }) catch unreachable;
    }

    pub fn isDeleted(self: *const LoadedExecutable, api: *const Api) bool {
        const ret = api.call(.PJRT_LoadedExecutable_IsDeleted, .{
            .executable = self.inner(),
        }) catch unreachable;
        return ret.is_deleted;
    }

    pub fn addressableDevices(self: *const LoadedExecutable, api: *const Api) []const *Device {
        const ret = api.call(.PJRT_LoadedExecutable_AddressableDevices, .{
            .executable = self.inner(),
        }) catch unreachable;
        return @ptrCast(ret.addressable_devices[0..ret.num_addressable_devices]);
    }

    pub const ExecuteArgs = struct {
        num_args: usize,
        arguments: []const [*]const *const Buffer,
        results: []const [*]*Buffer,
        events: ?[]?*Event,
        non_donatable_input_indices: []const i64 = &.{},
        context: ?*ExecuteContext,
    };

    pub fn execute(self: *const LoadedExecutable, api: *const Api, args: ExecuteArgs) ApiError!void {
        var options: meta.Struct(c.PJRT_ExecuteOptions) = .{
            .non_donatable_input_indices = @ptrCast(args.non_donatable_input_indices),
            .num_non_donatable_input_indices = args.non_donatable_input_indices.len,
            .context = @ptrCast(args.context),
        };
        _ = try api.call(.PJRT_LoadedExecutable_Execute, .{
            .executable = self.inner(),
            .options = @ptrCast(&options),
            .argument_lists = @ptrCast(args.arguments),
            .num_devices = args.arguments.len,
            .num_args = args.num_args,
            .output_lists = @ptrCast(args.results),
            .device_complete_events = if (args.events) |ev| @ptrCast(ev) else null,
        });
    }

    pub fn executable(self: *const LoadedExecutable, api: *const Api) ApiError!*Executable {
        const ret = try api.call(.PJRT_LoadedExecutable_GetExecutable, .{
            .loaded_executable = self.inner(),
        });
        return @ptrCast(ret.executable.?);
    }
};

pub const BufferType = enum(c.PJRT_Buffer_Type) {
    invalid = c.PJRT_Buffer_Type_INVALID,
    bool = c.PJRT_Buffer_Type_PRED,
    i2 = c.PJRT_Buffer_Type_S2,
    i4 = c.PJRT_Buffer_Type_S4,
    i8 = c.PJRT_Buffer_Type_S8,
    i16 = c.PJRT_Buffer_Type_S16,
    i32 = c.PJRT_Buffer_Type_S32,
    i64 = c.PJRT_Buffer_Type_S64,
    u2 = c.PJRT_Buffer_Type_U2,
    u4 = c.PJRT_Buffer_Type_U4,
    u8 = c.PJRT_Buffer_Type_U8,
    u16 = c.PJRT_Buffer_Type_U16,
    u32 = c.PJRT_Buffer_Type_U32,
    u64 = c.PJRT_Buffer_Type_U64,
    f16 = c.PJRT_Buffer_Type_F16,
    f32 = c.PJRT_Buffer_Type_F32,
    f64 = c.PJRT_Buffer_Type_F64,
    bf16 = c.PJRT_Buffer_Type_BF16,
    c64 = c.PJRT_Buffer_Type_C64,
    c128 = c.PJRT_Buffer_Type_C128,
    f8e5m2 = c.PJRT_Buffer_Type_F8E5M2,
    f8e4m3fn = c.PJRT_Buffer_Type_F8E4M3FN,
    f8e4m3b11fnuz = c.PJRT_Buffer_Type_F8E4M3B11FNUZ,
    f8e5m2fnuz = c.PJRT_Buffer_Type_F8E5M2FNUZ,
    f8e4m3fnuz = c.PJRT_Buffer_Type_F8E4M3FNUZ,
    f8e4m3 = c.PJRT_Buffer_Type_F8E4M3,
    f8e3m4 = c.PJRT_Buffer_Type_F8E3M4,
    f8e8m0 = c.PJRT_Buffer_Type_F8E8M0FNU,
    f4e2m1 = c.PJRT_Buffer_Type_F4E2M1FN,
};

pub const MemoryLayoutType = enum(c.PJRT_Buffer_MemoryLayout_Type) {
    tiled = c.PJRT_Buffer_MemoryLayout_Type_Tiled,
    strides = c.PJRT_Buffer_MemoryLayout_Type_Strides,
};

pub const MemoryLayout = union(MemoryLayoutType) {
    pub const Type = MemoryLayoutType;

    pub const Tiled = struct {
        minor_to_major: []const i64,
        tile_dims: []const i64,
        tile_dims_sizes: []const usize,
    };

    pub const Strides = struct {
        byte_strides: []const i64,
    };

    tiled: Tiled,
    strides: Strides,

    fn toCStruct(self: MemoryLayout) c.PJRT_Buffer_MemoryLayout {
        return switch (self) {
            .tiled => |v| c.PJRT_Buffer_MemoryLayout{
                .struct_size = c.PJRT_Buffer_MemoryLayout_STRUCT_SIZE,
                .extension_start = null,
                .type = c.PJRT_Buffer_MemoryLayout_Type_Tiled,
                .unnamed_0 = .{
                    .tiled = c.PJRT_Buffer_MemoryLayout_Tiled{
                        .struct_size = c.PJRT_Buffer_MemoryLayout_Tiled_STRUCT_SIZE,
                        .extension_start = null,
                        .minor_to_major = v.minor_to_major.ptr,
                        .minor_to_major_size = v.minor_to_major.len,
                        .tile_dims = v.tile_dims.ptr,
                        .tile_dim_sizes = v.tile_dims_sizes.ptr,
                        .num_tiles = v.tile_dims_sizes.len,
                    },
                },
            },
            .strides => |v| c.PJRT_Buffer_MemoryLayout{
                .struct_size = c.PJRT_Buffer_MemoryLayout_STRUCT_SIZE,
                .extension_start = null,
                .type = c.PJRT_Buffer_MemoryLayout_Type_Strides,
                .unnamed_0 = .{
                    .strides = c.PJRT_Buffer_MemoryLayout_Strides{
                        .struct_size = c.PJRT_Buffer_MemoryLayout_Strides_STRUCT_SIZE,
                        .extension_start = null,
                        .byte_strides = v.byte_strides.ptr,
                        .num_byte_strides = v.byte_strides.len,
                    },
                },
            },
        };
    }
};

pub const DefaultMemoryLayout = struct {
    pub const MAX_MINOR_TO_MAJOR: usize = 16;
    pub const MAX_TILE_DIMS: usize = 128;
    pub const MAX_NUM_TILES: usize = 16;

    minor_to_major_storage: [MAX_MINOR_TO_MAJOR]i64 = undefined,
    minor_to_major_len: usize = 0,

    tile_dims_storage: [MAX_TILE_DIMS]i64 = undefined,
    tile_dims_len: usize = 0,

    tile_dims_sizes_storage: [MAX_NUM_TILES]usize = undefined,
    tile_dims_sizes_len: usize = 0,

    pub fn toMemoryLayout(self: *const DefaultMemoryLayout) MemoryLayout {
        return .{
            .tiled = .{
                .minor_to_major = self.minor_to_major_storage[0..self.minor_to_major_len],
                .tile_dims = self.tile_dims_storage[0..self.tile_dims_len],
                .tile_dims_sizes = self.tile_dims_sizes_storage[0..self.tile_dims_sizes_len],
            },
        };
    }

    pub fn parseSerialized(serialized: []const u8) !DefaultMemoryLayout {
        var parsed: DefaultMemoryLayout = .{};
        try parseSerializedInto(serialized, &parsed);
        return parsed;
    }

    fn parseSerializedInto(serialized: []const u8, out: *DefaultMemoryLayout) !void {
        out.* = .{};
        const raw = std.mem.trim(u8, serialized, " \t\n\r\x00");

        if (raw.len < 2 or raw[0] != '{' or raw[raw.len - 1] != '}') {
            return error.InvalidLayoutString;
        }

        const body = raw[1 .. raw.len - 1];
        const first_colon = std.mem.indexOfScalar(u8, body, ':');

        const minor_to_major_str = std.mem.trim(u8, if (first_colon) |index| body[0..index] else body, " \t\n\r");

        if (minor_to_major_str.len > 0) {
            var mtm_iter = std.mem.splitScalar(u8, minor_to_major_str, ',');
            while (mtm_iter.next()) |part_raw| {
                if (out.minor_to_major_len >= MAX_MINOR_TO_MAJOR) return error.LayoutTooLarge;
                out.minor_to_major_storage[out.minor_to_major_len] = try parseI64Token(part_raw);
                out.minor_to_major_len += 1;
            }
        }

        if (first_colon) |index| {
            const attrs = body[index + 1 ..];
            var cursor: usize = 0;
            while (cursor < attrs.len) {
                const t_pos = std.mem.indexOfScalarPos(u8, attrs, cursor, 'T') orelse break;
                cursor = t_pos + 1;
                while (cursor < attrs.len and std.ascii.isWhitespace(attrs[cursor])) : (cursor += 1) {}
                if (cursor >= attrs.len or attrs[cursor] != '(') continue;

                while (cursor < attrs.len and attrs[cursor] == '(') {
                    const tuple_start = cursor + 1;
                    const tuple_end = std.mem.indexOfScalarPos(u8, attrs, tuple_start, ')') orelse return error.InvalidLayoutString;
                    const tuple = attrs[tuple_start..tuple_end];

                    var dims_in_tile: usize = 0;
                    var tuple_iter = std.mem.splitScalar(u8, tuple, ',');
                    while (tuple_iter.next()) |dim_raw| {
                        if (out.tile_dims_len >= MAX_TILE_DIMS) return error.LayoutTooLarge;
                        out.tile_dims_storage[out.tile_dims_len] = try parseI64Token(dim_raw);
                        out.tile_dims_len += 1;
                        dims_in_tile += 1;
                    }
                    if (dims_in_tile == 0) return error.InvalidLayoutString;

                    if (out.tile_dims_sizes_len >= MAX_NUM_TILES) return error.LayoutTooLarge;
                    out.tile_dims_sizes_storage[out.tile_dims_sizes_len] = dims_in_tile;
                    out.tile_dims_sizes_len += 1;

                    cursor = tuple_end + 1;
                }
            }
        }
    }

    fn parseI64Token(token_raw: []const u8) !i64 {
        const token = std.mem.trim(u8, token_raw, " \t\n\r");
        if (token.len == 0) return error.InvalidLayoutString;
        return std.fmt.parseInt(i64, token, 10);
    }
};

pub const HostBufferSemantics = enum(c.PJRT_HostBufferSemantics) {
    ImmutableOnlyDuringCall = c.PJRT_HostBufferSemantics_kImmutableOnlyDuringCall,
    ImmutableUntilTransferCompletes = c.PJRT_HostBufferSemantics_kImmutableUntilTransferCompletes,
    ImmutableZeroCopy = c.PJRT_HostBufferSemantics_kImmutableZeroCopy,
    MutableZeroCopy = c.PJRT_HostBufferSemantics_kMutableZeroCopy,
};

pub const Buffer = opaque {
    const inner = InnerMixin(c.PJRT_Buffer).inner;

    pub fn deinit(self: *Buffer, api: *const Api) void {
        _ = api.call(.PJRT_Buffer_Destroy, .{
            .buffer = self.inner(),
        }) catch unreachable;
    }

    pub fn device(self: *const Buffer, api: *const Api) ApiError!*Device {
        const ret = try api.call(.PJRT_Buffer_Device, .{
            .buffer = self.inner(),
        });
        return @ptrCast(ret.device.?);
    }

    pub fn delete(self: *Buffer, api: *const Api) void {
        _ = api.call(.PJRT_Buffer_Delete, .{
            .buffer = self.inner(),
        }) catch unreachable;
    }

    pub fn isDeleted(self: *const Buffer, api: *const Api) bool {
        const ret = api.call(.PJRT_Buffer_IsDeleted, .{
            .buffer = self.inner(),
        }) catch unreachable;
        return ret.is_deleted;
    }

    pub fn isOnCpu(self: *const Buffer, api: *const Api) bool {
        const ret = api.call(.PJRT_Buffer_IsOnCpu, .{
            .buffer = self.inner(),
        }) catch unreachable;
        return ret.is_on_cpu;
    }

    pub fn toHostBuffer(self: *const Buffer, api: *const Api, dst: []u8) ApiError!?*Event {
        const ret = try api.call(.PJRT_Buffer_ToHostBuffer, .{
            .src = self.inner(),
            .dst = dst.ptr,
            .dst_size = dst.len,
        });
        return @ptrCast(ret.event);
    }

    pub fn elementType(self: *const Buffer, api: *const Api) BufferType {
        const ret = api.call(.PJRT_Buffer_ElementType, .{
            .buffer = self.inner(),
        }) catch unreachable;
        return @enumFromInt(ret.type);
    }

    pub fn dimensions(self: *const Buffer, api: *const Api) []const i64 {
        const ret = api.call(.PJRT_Buffer_Dimensions, .{
            .buffer = self.inner(),
        }) catch unreachable;
        if (ret.num_dims == 0) {
            return &.{};
        }
        return ret.dims[0..ret.num_dims];
    }

    pub fn unpaddedDimensions(self: *const Buffer, api: *const Api) ApiError![]const i64 {
        const ret = try api.call(.PJRT_Buffer_UnpaddedDimensions, .{
            .buffer = self.inner(),
        });
        return ret.unpadded_dims[0..ret.num_dims];
    }

    pub fn onDeviceSizeInBytes(self: *const Buffer, api: *const Api) ApiError!usize {
        const ret = try api.call(.PJRT_Buffer_OnDeviceSizeInBytes, .{
            .buffer = self.inner(),
        });
        return @intCast(ret.on_device_size_in_bytes);
    }

    pub fn copyToDevice(self: *const Buffer, api: *const Api, dst_device: *Device) ApiError!*Buffer {
        const ret = try api.call(.PJRT_Buffer_CopyToDevice, .{
            .buffer = self.inner(),
            .dst_device = dst_device.inner(),
        });
        return @ptrCast(ret.dst_buffer.?);
    }

    pub fn readyEvent(self: *const Buffer, api: *const Api) *Event {
        const ret = api.call(.PJRT_Buffer_ReadyEvent, .{
            .buffer = self.inner(),
        }) catch unreachable;
        return @ptrCast(ret.event.?);
    }

    pub fn opaqueDeviceMemoryDataPointer(self: *const Buffer, api: *const Api) ApiError!*anyopaque {
        const ret = try api.call(.PJRT_Buffer_OpaqueDeviceMemoryDataPointer, .{
            .buffer = self.inner(),
        });
        return ret.device_memory_ptr.?;
    }

    pub fn copyRawToHost(self: *const Buffer, api: *const Api, dst: []u8, offset: i64) ApiError!?*Event {
        const ret = try api.call(.PJRT_Buffer_CopyRawToHost, .{
            .buffer = self.inner(),
            .dst = @ptrCast(dst),
            .offset = offset,
            .transfer_size = @intCast(dst.len),
        });
        return @ptrCast(ret.event);
    }

    pub fn copyToMemory(self: *const Buffer, api: *const Api, dst_memory: *const Memory) ApiError!*Buffer {
        const ret = try api.call(.PJRT_Buffer_CopyToMemory, .{
            .buffer = self.inner(),
            .dst_memory = dst_memory.inner(),
        });
        return @ptrCast(ret.dst_buffer);
    }

    pub fn memory(self: *const Buffer, api: *const Api) *const Memory {
        const ret = api.call(.PJRT_Buffer_Memory, .{
            .buffer = self.inner(),
        }) catch unreachable;
        return @ptrCast(ret.memory);
    }

    pub fn increaseExternalReferenceCount(self: *const Buffer, api: *const Api) ApiError!void {
        _ = try api.call(.PJRT_Buffer_IncreaseExternalReferenceCount, .{
            .buffer = self.inner(),
        });
    }

    pub fn decreaseExternalReferenceCount(self: *const Buffer, api: *const Api) ApiError!void {
        _ = try api.call(.PJRT_Buffer_DecreaseExternalReferenceCount, .{
            .buffer = self.inner(),
        });
    }
};

pub const Event = opaque {
    const inner = InnerMixin(c.PJRT_Event).inner;

    pub fn deinit(self: *Event, api: *const Api) void {
        _ = api.call(.PJRT_Event_Destroy, .{
            .event = self.inner(),
        }) catch unreachable;
    }

    pub fn isReady(self: *const Event, api: *const Api) bool {
        const ret = api.call(.PJRT_Event_IsReady, .{
            .event = self.inner(),
        }) catch unreachable;
        return ret.is_ready;
    }

    pub fn getEventError(self: *const Event, api: *const Api) ?*Error {
        var args: meta.Struct(c.PJRT_Event_Error_Args) = .{
            .event = self.inner(),
        };
        const result: ?*c.PJRT_Error = api.inner.PJRT_Event_Error.?(@ptrCast(&args));
        return @ptrCast(result);
    }

    pub fn await(self: *Event, api: *const Api, io: std.Io) ApiError!void {
        if (self.isReady(api)) {
            return;
        }

        const Ctx = struct {
            err: ?*Error = null,
            event: std.Io.Event = .unset,
            io: std.Io,
        };
        var ctx: Ctx = .{ .io = io };
        try self.onReady(api, Ctx, struct {
            fn call(err: ?*Error, ctx_: *Ctx) void {
                ctx_.err = err;
                ctx_.event.set(ctx_.io);
            }
        }.call, &ctx);
        ctx.event.waitUncancelable(io);

        if (ctx.err) |e| {
            defer e.deinit(api);
            const err_code = e.getCode(api).toApiError();
            log.err("{t} {s}", .{ err_code, e.getMessage(api) });
            return err_code;
        }
    }

    pub fn awaitRaw(self: *const Event, api: *const Api) ApiError!void {
        _ = try api.call(.PJRT_Event_Await, .{
            .event = self.inner(),
        });
    }

    pub fn onReady(self: *Event, api: *const Api, comptime T: type, func: fn (err: ?*Error, user_arg: *T) void, user_arg: *T) ApiError!void {
        _ = try api.call(.PJRT_Event_OnReady, .{
            .event = self.inner(),
            .callback = @ptrCast(&struct {
                fn call(err: ?*Error, user_arg_: ?*anyopaque) callconv(.c) void {
                    func(err, @ptrCast(@alignCast(user_arg_.?)));
                }
            }.call),
            .user_arg = user_arg,
        });
    }
};

pub const Memory = opaque {
    pub const Kind = enum {
        device,
        host_pinned,
        host_unpinned,

        pub fn pjrtName(k: Kind) []const u8 {
            return switch (k) {
                .device => "device",
                .host_pinned => "pinned_host",
                .host_unpinned => "unpinned_host",
            };
        }
    };

    const inner = InnerMixin(c.PJRT_Memory).inner;

    pub fn id(self: *const Memory, api: *const Api) usize {
        const ret = api.call(.PJRT_Memory_Id, .{ .memory = self.inner() }) catch unreachable;
        return @intCast(ret.id);
    }

    pub fn kind(self: *const Memory, api: *const Api) Kind {
        const ret = api.call(.PJRT_Memory_Kind, .{ .memory = self.inner() }) catch unreachable;
        return switch (ret.kind_size) {
            "device".len => .device,
            "pinned_host".len => .host_pinned,
            "unpinned_host".len => .host_unpinned,
            else => @panic("Memory kind not supported"),
        };
    }

    pub fn kind_(self: *const Memory, api: *const Api) []const u8 {
        const ret = api.call(.PJRT_Memory_Kind, .{ .memory = self.inner() }) catch unreachable;
        return ret.kind[0..ret.kind_size];
    }

    pub fn kindId(self: *const Memory, api: *const Api) u32 {
        const ret = api.call(.PJRT_Memory_Kind_Id, .{
            .memory = self.inner(),
        }) catch unreachable;
        return @bitCast(ret.kind_id);
    }

    pub fn debugString(self: *const Memory, api: *const Api) []const u8 {
        const ret = api.call(.PJRT_Memory_DebugString, .{
            .memory = self.inner(),
        }) catch unreachable;
        if (ret.debug_string) |debug_string| {
            return debug_string[0..ret.debug_string_size];
        }
        return &.{};
    }

    pub fn toString(self: *const Memory, api: *const Api) []const u8 {
        const ret = api.call(.PJRT_Memory_ToString, .{
            .memory = self.inner(),
        }) catch unreachable;
        if (ret.to_string) |to_string| {
            return to_string[0..ret.to_string_size];
        }
        return &.{};
    }

    pub fn addressableByDevices(self: *const Memory, api: *const Api) []const *Device {
        const ret = api.call(.PJRT_Memory_AddressableByDevices, .{
            .memory = self.inner(),
        }) catch unreachable;
        if (ret.devices) |devices| {
            return @ptrCast(devices[0..ret.num_devices]);
        }
        return &.{};
    }
};

pub const AsyncHostToDeviceTransferManager = opaque {
    const inner = InnerMixin(c.PJRT_AsyncHostToDeviceTransferManager).inner;

    pub fn deinit(self: *AsyncHostToDeviceTransferManager, api: *const Api) void {
        _ = api.call(.PJRT_AsyncHostToDeviceTransferManager_Destroy, .{
            .transfer_manager = self.inner(),
        }) catch unreachable;
    }

    pub fn transferData(self: *AsyncHostToDeviceTransferManager, api: *const Api, buffer_index: usize, data: []const u8, offset: i64, is_last_transfer: bool) ApiError!*Event {
        const ret = try api.call(.PJRT_AsyncHostToDeviceTransferManager_TransferData, .{
            .transfer_manager = self.inner(),
            .buffer_index = @intCast(buffer_index),
            .data = data.ptr,
            .offset = offset,
            .transfer_size = @intCast(data.len),
            .is_last_transfer = is_last_transfer,
        });
        return @ptrCast(ret.done_with_h2d_transfer.?);
    }

    pub fn retrieveBuffer(self: *AsyncHostToDeviceTransferManager, api: *const Api, buffer_index: usize) ApiError!*Buffer {
        const ret = try api.call(.PJRT_AsyncHostToDeviceTransferManager_RetrieveBuffer, .{
            .transfer_manager = self.inner(),
            .buffer_index = @intCast(buffer_index),
        });
        return @ptrCast(ret.buffer_out.?);
    }

    pub fn device(self: *AsyncHostToDeviceTransferManager, api: *const Api) ApiError!*Device {
        const ret = try api.call(.PJRT_AsyncHostToDeviceTransferManager_Device, .{
            .transfer_manager = self.inner(),
        });
        return @ptrCast(ret.device_out.?);
    }

    pub fn bufferCount(self: *AsyncHostToDeviceTransferManager, api: *const Api) ApiError!usize {
        const ret = try api.call(.PJRT_AsyncHostToDeviceTransferManager_BufferCount, .{
            .transfer_manager = self.inner(),
        });
        return ret.buffer_count;
    }

    pub fn bufferSize(self: *AsyncHostToDeviceTransferManager, api: *const Api, buffer_index: usize) ApiError!usize {
        const ret = try api.call(.PJRT_AsyncHostToDeviceTransferManager_BufferSize, .{
            .transfer_manager = self.inner(),
            .buffer_index = @intCast(buffer_index),
        });
        return ret.buffer_size;
    }

    pub fn setBufferError(self: *AsyncHostToDeviceTransferManager, api: *const Api, buffer_index: usize, error_code: c.PJRT_Error_Code, error_message: []const u8) ApiError!void {
        _ = try api.call(.PJRT_AsyncHostToDeviceTransferManager_SetBufferError, .{
            .transfer_manager = self.inner(),
            .buffer_index = @intCast(buffer_index),
            .error_code = error_code,
            .error_message = error_message.ptr,
            .error_message_size = error_message.len,
        });
    }

    pub fn addMetadata(self: *AsyncHostToDeviceTransferManager, api: *const Api, transfer_metadata: []const NamedValue) ApiError!void {
        _ = try api.call(.PJRT_AsyncHostToDeviceTransferManager_AddMetadata, .{
            .transfer_manager = self.inner(),
            .transfer_metadata = @ptrCast(transfer_metadata),
            .num_metadata = transfer_metadata.len,
        });
    }
};

pub const NamedValue = extern struct {
    comptime {
        std.debug.assert(@sizeOf(NamedValue) == @sizeOf(c.PJRT_NamedValue));
    }

    inner: c.PJRT_NamedValue,

    pub const Kind = enum(c.PJRT_NamedValue_Type) {
        string = c.PJRT_NamedValue_kString,
        int64 = c.PJRT_NamedValue_kInt64,
        int64list = c.PJRT_NamedValue_kInt64List,
        float = c.PJRT_NamedValue_kFloat,
        bool = c.PJRT_NamedValue_kBool,
    };

    pub const Value = union(Kind) {
        string: []const u8,
        int64: i64,
        int64list: []const i64,
        float: f32,
        bool: bool,
    };

    pub fn kind(self: NamedValue) Kind {
        return @enumFromInt(self.inner.type);
    }

    pub fn name(self: NamedValue) []const u8 {
        return self.inner.name[0..self.inner.name_size];
    }

    pub fn value(self: NamedValue) Value {
        return switch (self.kind()) {
            .string => .{ .string = self.inner.unnamed_0.string_value[0..self.inner.value_size] },
            .int64 => .{ .int64 = self.inner.unnamed_0.int64_value },
            .int64list => .{ .int64list = self.inner.unnamed_0.int64_array_value[0..self.inner.value_size] },
            .float => .{ .float = self.inner.unnamed_0.float_value },
            .bool => .{ .bool = self.inner.unnamed_0.bool_value },
        };
    }

    pub fn init(comptime kind_: Kind, name_: []const u8, value_: std.meta.fieldInfo(Value, kind_).type) NamedValue {
        return .{
            .inner = .{
                .struct_size = c.PJRT_NamedValue_STRUCT_SIZE,
                .extension_start = null,
                .name = @ptrCast(name_),
                .name_size = name_.len,
                .type = @intFromEnum(kind_),
                .unnamed_0 = switch (kind_) {
                    .string => .{ .string_value = @ptrCast(@constCast(value_)) },
                    .int64 => .{ .int64_value = value_ },
                    .int64list => .{ .int64_array_value = @ptrCast(@constCast(value_)) },
                    .float => .{ .float_value = value_ },
                    .bool => .{ .bool_value = value_ },
                },
                .value_size = switch (kind_) {
                    .string, .int64list => value_.len,
                    inline else => 1,
                },
            },
        };
    }

    pub fn format(self: NamedValue, writer: *std.Io.Writer) !void {
        try writer.print("{s}{{ .name = {s},", .{ @typeName(NamedValue), self.name() });
        switch (self.value()) {
            .string => |v| try writer.print(" .string = \"{s}\" ", .{v}),
            .int64 => |v| try writer.print(" .int64 = {d} ", .{v}),
            .int64list => |v| try writer.print(" .int64list = {any} ", .{v}),
            .float => |v| try writer.print(" .float = {d} ", .{v}),
            .bool => |v| try writer.print(" .bool = {} ", .{v}),
        }
        try writer.writeAll("}");
    }
};

pub const Ffi = extern struct {
    inner: *const c.PJRT_FFI,

    pub const UserData = extern struct {
        type_id: i64,
        user_data: *anyopaque,

        fn toCStruct(self: UserData) c.PJRT_FFI_UserData {
            return .{
                .type_id = self.type_id,
                .data = self.user_data,
            };
        }
    };

    pub const TypeInfo = struct {
        deleter: ?*const fn (*anyopaque) callconv(.c) void = null,
        serialize: ?*const fn () callconv(.c) void = null,
        deserialize: ?*const fn () callconv(.c) void = null,

        pub fn toCStruct(self: TypeInfo) c.PJRT_FFI_Type_Info {
            return .{
                .deleter = @ptrCast(self.deleter),
                .serialize = @ptrCast(self.serialize),
                .deserialize = @ptrCast(self.deserialize),
            };
        }
    };

    // todo : support all missing handlers available in GPU plugin extension: handler_instantiate, handler_prepare, handler_initialize
    // introduced by https://github.com/openxla/xla/commit/ef85a7bcc308313492ebc50295a8a08b4e51b8f5
    pub fn register(
        self: *const Ffi,
        api: *const Api,
        target_name: []const u8,
        platform_name: []const u8,
        func: *const ffi.Handler,
        traits: ffi.HandlerTraits,
    ) ApiError!void {
        var ret: meta.Struct(c.PJRT_FFI_Register_Handler_Args) = .{
            .target_name = target_name.ptr,
            .target_name_size = target_name.len,
            .handler = @as(?*anyopaque, @ptrCast(@constCast(func))),
            .platform_name = platform_name.ptr,
            .platform_name_size = platform_name.len,
            .traits = @bitCast(traits),
        };
        const result = self.inner.register_handler.?(@ptrCast(&ret));
        if (result) |pjrt_c_error| {
            const pjrt_error: *Error = @ptrCast(pjrt_c_error);
            log.err("registerFfi error: {s}", .{pjrt_error.getMessage(api)});
            return pjrt_error.getCode(api).toApiError();
        }
    }

    pub fn registerTypeId(self: *const Ffi, api: *const Api, type_name: []const u8, type_info: ?*const c.PJRT_FFI_Type_Info) ApiError!ffi.TypeId {
        var ret: meta.Struct(c.PJRT_FFI_Type_Register_Args) = .{
            .type_name = type_name.ptr,
            .type_name_size = type_name.len,
            .type_id = 0, // let the plugin assign a unique type ID
            .type_info = @ptrCast(@constCast(type_info)),
        };
        const result = self.inner.type_register.?(&ret);
        if (result) |pjrt_c_error| {
            const pjrt_error: *Error = @ptrCast(pjrt_c_error);
            return pjrt_error.getCode(api).toApiError();
        }

        return .{ .type_id = ret.type_id };
    }

    pub fn addUserData(self: *const Ffi, api: *const Api, context: *ExecuteContext, user_data: UserData) ApiError!void {
        var ret: meta.Struct(c.PJRT_FFI_UserData_Add_Args) = .{
            .context = @ptrCast(context),
            .user_data = user_data.toCStruct(),
        };
        const result = self.inner.user_data_add.?(&ret);
        if (result) |pjrt_c_error| {
            const pjrt_error: *Error = @ptrCast(pjrt_c_error);
            log.err("addUserData error: {s}", .{pjrt_error.getMessage(api)});
            return pjrt_error.getCode(api).toApiError();
        }
    }
};

pub const Profiler = struct {
    pjrt_api: *const Api,
    api: *const c.PLUGIN_Profiler_Api,
    inner: *c.PLUGIN_Profiler,

    pub fn init(api_: *const Api, prof_ext: *const c.PJRT_Profiler_Extension, options_pb: []const u8) ApiError!Profiler {
        var args: c.PLUGIN_Profiler_Create_Args = .{
            .struct_size = meta.structSize(c.PLUGIN_Profiler_Create_Args),
            .options = options_pb.ptr,
            .options_size = options_pb.len,
            .profiler = null,
        };

        const profiler_api: *const c.PLUGIN_Profiler_Api = @ptrCast(prof_ext.profiler_api);
        if (profiler_api.create.?(&args)) |err| {
            const pjrt_err: *Error = @ptrCast(err);
            return pjrt_err.getCode(api_).toApiError();
        }

        return .{
            .pjrt_api = api_,
            .api = profiler_api,
            .inner = args.profiler.?,
        };
    }

    pub fn start(self: *Profiler) ApiError!void {
        var args: c.PLUGIN_Profiler_Start_Args = .{
            .struct_size = meta.structSize(c.PLUGIN_Profiler_Start_Args),
            .profiler = self.inner,
        };

        if (self.api.start.?(&args)) |err| {
            const pjrt_err: *Error = @ptrCast(err);
            return pjrt_err.getCode(self.pjrt_api).toApiError();
        }
    }

    pub fn stop(self: *Profiler) ApiError!void {
        var args: c.PLUGIN_Profiler_Stop_Args = .{
            .struct_size = meta.structSize(c.PLUGIN_Profiler_Stop_Args),
            .profiler = self.inner,
        };

        if (self.api.stop.?(&args)) |err| {
            const pjrt_err: *Error = @ptrCast(err);
            return pjrt_err.getCode(self.pjrt_api).toApiError();
        }
    }

    pub fn collectData(self: *Profiler, allocator: std.mem.Allocator) ![]u8 {
        var args: c.PLUGIN_Profiler_CollectData_Args = .{
            .struct_size = meta.structSize(c.PLUGIN_Profiler_CollectData_Args),
            .profiler = self.inner,
            .buffer = null,
            .buffer_size_in_bytes = 0,
        };

        if (self.api.collect_data.?(&args)) |err| {
            const pjrt_err: *Error = @ptrCast(err);
            return pjrt_err.getCode(self.pjrt_api).toApiError();
        }

        if (args.buffer != null) {
            var data = args.buffer[0..args.buffer_size_in_bytes];
            if (data.len > 0 and data[data.len - 1] == 0) {
                data = data[0 .. data.len - 1];
            }
            return allocator.dupe(u8, data);
        }

        if (args.buffer_size_in_bytes == 0) return &[_]u8{};

        const buffer = try allocator.alloc(u8, args.buffer_size_in_bytes);
        errdefer allocator.free(buffer);

        args.buffer = buffer.ptr;
        if (self.api.collect_data.?(&args)) |err| {
            const pjrt_err: *Error = @ptrCast(err);
            return pjrt_err.getCode(self.pjrt_api).toApiError();
        }

        // Trim trailing null byte logic
        var data = buffer[0..args.buffer_size_in_bytes];
        if (data.len > 0 and data[data.len - 1] == 0) {
            data = data[0 .. data.len - 1];
        }

        const pb = try allocator.dupe(u8, data);
        allocator.free(buffer);

        return pb;
    }
};