@preconcurrency import COnnxRuntime
import Foundation

public enum ORTLoggingLevel: Int32 {
    case verbose
    case info
    case warning
    case error
    case fatal

    internal var native: OrtLoggingLevel {
        OrtLoggingLevel(UInt32(rawValue))
    }
}

public enum ORTValueType: Int32 {
    case unknown
    case tensor

    internal init(native: ONNXType) throws {
        switch native {
        case ONNX_TYPE_UNKNOWN:
            self = .unknown
        case ONNX_TYPE_TENSOR:
            self = .tensor
        default:
            throw ORTError.unsupported("Only tensor ONNX values are currently supported")
        }
    }
}

public enum ORTTensorElementDataType: Int32 {
    case undefined
    case float
    case int8
    case uInt8
    case int32
    case uInt32
    case int64
    case uInt64
    case string

    internal var native: ONNXTensorElementDataType {
        switch self {
        case .undefined:
            return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED
        case .float:
            return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT
        case .int8:
            return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8
        case .uInt8:
            return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8
        case .int32:
            return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32
        case .uInt32:
            return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32
        case .int64:
            return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64
        case .uInt64:
            return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64
        case .string:
            return ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING
        }
    }

    internal init(native: ONNXTensorElementDataType) throws {
        switch native {
        case ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED:
            self = .undefined
        case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT:
            self = .float
        case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8:
            self = .int8
        case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8:
            self = .uInt8
        case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32:
            self = .int32
        case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32:
            self = .uInt32
        case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64:
            self = .int64
        case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64:
            self = .uInt64
        case ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING:
            self = .string
        default:
            throw ORTError.unsupported("Unsupported tensor element type \(native.rawValue)")
        }
    }
}

public enum ORTGraphOptimizationLevel: Int32 {
    case none
    case basic
    case extended
    case layout
    case all

    internal var native: GraphOptimizationLevel {
        switch self {
        case .none:
            return ORT_DISABLE_ALL
        case .basic:
            return ORT_ENABLE_BASIC
        case .extended:
            return ORT_ENABLE_EXTENDED
        case .layout, .all:
            return ORT_ENABLE_ALL
        }
    }
}

public typealias ORTCAPIRegisterCustomOpsFnPtr = UnsafeMutableRawPointer

private final class ORTValueStorage {
    let handle: OpaquePointer
    let retainedBuffer: NSMutableData?

    init(handle: OpaquePointer, retainedBuffer: NSMutableData? = nil) {
        self.handle = handle
        self.retainedBuffer = retainedBuffer
    }

    deinit {
        ORTNative.api.pointee.ReleaseValue?(handle)
    }
}

public func ORTVersion() -> String? {
    ORTNative.apiBase?.pointee.GetVersionString().map { String(cString: $0) }
}

public func ORTIsCoreMLExecutionProviderAvailable() -> Bool {
    false
}

public func ORTSetSeed(_ seed: Int64) {
    guard let trainingApi = ORTNative.trainingApi else {
        return
    }

    _ = trainingApi.pointee.SetSeed?(seed)
}

public struct ORTValueTypeInfo {
    public let type: ORTValueType
    public let tensorTypeAndShapeInfo: ORTTensorTypeAndShapeInfo?

    public init(type: ORTValueType, tensorTypeAndShapeInfo: ORTTensorTypeAndShapeInfo?) {
        self.type = type
        self.tensorTypeAndShapeInfo = tensorTypeAndShapeInfo
    }
}

public struct ORTTensorTypeAndShapeInfo {
    public let elementType: ORTTensorElementDataType
    public let shape: [Int]

    public init(elementType: ORTTensorElementDataType, shape: [Int]) {
        self.elementType = elementType
        self.shape = shape
    }
}

public struct ORTCoreMLExecutionProviderOptions {
    public var useCPUOnly: Bool
    public var useCPUAndGPU: Bool
    public var enableOnSubgraphs: Bool
    public var onlyEnableForDevicesWithANE: Bool
    public var onlyAllowStaticInputShapes: Bool
    public var createMLProgram: Bool

    public init(
        useCPUOnly: Bool = false,
        useCPUAndGPU: Bool = false,
        enableOnSubgraphs: Bool = false,
        onlyEnableForDevicesWithANE: Bool = false,
        onlyAllowStaticInputShapes: Bool = false,
        createMLProgram: Bool = false
    ) {
        self.useCPUOnly = useCPUOnly
        self.useCPUAndGPU = useCPUAndGPU
        self.enableOnSubgraphs = enableOnSubgraphs
        self.onlyEnableForDevicesWithANE = onlyEnableForDevicesWithANE
        self.onlyAllowStaticInputShapes = onlyAllowStaticInputShapes
        self.createMLProgram = createMLProgram
    }

    internal var providerOptions: [String: String] {
        [
            "use_cpu_only": ORTNative.boolString(useCPUOnly),
            "use_cpu_and_gpu": ORTNative.boolString(useCPUAndGPU),
            "enable_on_subgraphs": ORTNative.boolString(enableOnSubgraphs),
            "only_enable_for_devices_with_ane": ORTNative.boolString(onlyEnableForDevicesWithANE),
            "only_allow_static_input_shapes": ORTNative.boolString(onlyAllowStaticInputShapes),
            "create_mlprogram": ORTNative.boolString(createMLProgram),
        ]
    }
}

public struct ORTXnnpackExecutionProviderOptions {
    public var intraOpNumThreads: Int32

    public init(intraOpNumThreads: Int32 = 0) {
        self.intraOpNumThreads = intraOpNumThreads
    }
}

public final class ORTEnv {
    public let loggingLevel: ORTLoggingLevel
    internal let handle: OpaquePointer

    public init(loggingLevel: ORTLoggingLevel) throws {
        self.loggingLevel = loggingLevel

        var env: OpaquePointer?
        try ORTNative.check(
            ORTNative.api.pointee.CreateEnv?(loggingLevel.native, "onnxruntime-swift-interface", &env),
            action: "CreateEnv"
        )
        guard let env else {
            throw ORTError.runtimeFailure(action: "CreateEnv", message: "Environment was not returned")
        }
        handle = env
    }

    deinit {
        ORTNative.api.pointee.ReleaseEnv?(handle)
    }
}

public final class ORTSession {
    private let handle: OpaquePointer

    public init(
        env: ORTEnv,
        modelPath: String,
        sessionOptions: ORTSessionOptions? = nil
    ) throws {
        var session: OpaquePointer?
        let createdOptions = sessionOptions == nil ? try ORTSessionOptions() : nil
        let optionsHandle = sessionOptions?.handle ?? createdOptions!.handle

        try modelPath.withCString { modelPathPointer in
            try ORTNative.check(
                ORTNative.api.pointee.CreateSession?(env.handle, modelPathPointer, optionsHandle, &session),
                action: "CreateSession"
            )
        }

        guard let session else {
            throw ORTError.runtimeFailure(action: "CreateSession", message: "Session was not returned")
        }
        handle = session
    }

    deinit {
        ORTNative.api.pointee.ReleaseSession?(handle)
    }

    public func run(
        inputs: [String: ORTValue],
        outputs: [String: ORTValue],
        runOptions: ORTRunOptions? = nil
    ) throws {
        let sortedInputs = inputs.keys.sorted()
        let sortedOutputs = outputs.keys.sorted()
        let inputValues = sortedInputs.map { Optional(inputs[$0]!.handle) }
        var outputValues = sortedOutputs.map { Optional(outputs[$0]!.handle) }
        let outputCount = outputValues.count

        try ORTNative.withCStringArray(sortedInputs) { inputNames in
            try ORTNative.withCStringArray(sortedOutputs) { outputNames in
                try inputNames.withUnsafeBufferPointer { inputNameBuffer in
                    try outputNames.withUnsafeBufferPointer { outputNameBuffer in
                        try inputValues.withUnsafeBufferPointer { inputValueBuffer in
                            try outputValues.withUnsafeMutableBufferPointer { outputValueBuffer in
                                try ORTNative.check(
                                    ORTNative.api.pointee.Run?(
                                        handle,
                                        runOptions?.handle,
                                        inputNameBuffer.baseAddress,
                                        inputValueBuffer.baseAddress,
                                        inputValues.count,
                                        outputNameBuffer.baseAddress,
                                        outputCount,
                                        outputValueBuffer.baseAddress
                                    ),
                                    action: "Run"
                                )
                            }
                        }
                    }
                }
            }
        }
    }

    public func run(
        inputs: [String: ORTValue],
        outputNames: Set<String>,
        runOptions: ORTRunOptions? = nil
    ) throws -> [String: ORTValue] {
        let sortedInputs = inputs.keys.sorted()
        let sortedOutputs = outputNames.sorted()

        let inputValues = sortedInputs.map { Optional(inputs[$0]!.handle) }
        var outputValues = Array<OpaquePointer?>(repeating: nil, count: sortedOutputs.count)
        let outputCount = sortedOutputs.count

        try ORTNative.withCStringArray(sortedInputs) { inputNames in
            try ORTNative.withCStringArray(sortedOutputs) { outputNamesArray in
                try inputNames.withUnsafeBufferPointer { inputNameBuffer in
                    try outputNamesArray.withUnsafeBufferPointer { outputNameBuffer in
                        try inputValues.withUnsafeBufferPointer { inputValueBuffer in
                            try outputValues.withUnsafeMutableBufferPointer { outputValueBuffer in
                                try ORTNative.check(
                                    ORTNative.api.pointee.Run?(
                                        handle,
                                        runOptions?.handle,
                                        inputNameBuffer.baseAddress,
                                        inputValueBuffer.baseAddress,
                                        inputValues.count,
                                        outputNameBuffer.baseAddress,
                                        outputCount,
                                        outputValueBuffer.baseAddress
                                    ),
                                    action: "Run"
                                )
                            }
                        }
                    }
                }
            }
        }

        var results: [String: ORTValue] = [:]
        for (index, name) in sortedOutputs.enumerated() {
            guard let valueHandle = outputValues[index] else {
                throw ORTError.runtimeFailure(action: "Run", message: "Output \(name) was not returned")
            }
            results[name] = ORTValue(adopting: valueHandle)
        }
        return results
    }

    public func inputNames() throws -> [String] {
        var count = 0
        try ORTNative.check(ORTNative.api.pointee.SessionGetInputCount?(handle, &count), action: "SessionGetInputCount")
        return try ORTNative.withAllocator { allocator in
            var results: [String] = []
            results.reserveCapacity(count)
            for index in 0..<count {
                var pointer: UnsafeMutablePointer<CChar>?
                try ORTNative.check(
                    ORTNative.api.pointee.SessionGetInputName?(handle, index, allocator, &pointer),
                    action: "SessionGetInputName"
                )
                results.append(try ORTNative.copyAllocatedString(pointer, allocator: allocator))
            }
            return results
        }
    }

    public func overridableInitializerNames() throws -> [String] {
        var count = 0
        try ORTNative.check(
            ORTNative.api.pointee.SessionGetOverridableInitializerCount?(handle, &count),
            action: "SessionGetOverridableInitializerCount"
        )
        return try ORTNative.withAllocator { allocator in
            var results: [String] = []
            results.reserveCapacity(count)
            for index in 0..<count {
                var pointer: UnsafeMutablePointer<CChar>?
                try ORTNative.check(
                    ORTNative.api.pointee.SessionGetOverridableInitializerName?(handle, index, allocator, &pointer),
                    action: "SessionGetOverridableInitializerName"
                )
                results.append(try ORTNative.copyAllocatedString(pointer, allocator: allocator))
            }
            return results
        }
    }

    public func outputNames() throws -> [String] {
        var count = 0
        try ORTNative.check(ORTNative.api.pointee.SessionGetOutputCount?(handle, &count), action: "SessionGetOutputCount")
        return try ORTNative.withAllocator { allocator in
            var results: [String] = []
            results.reserveCapacity(count)
            for index in 0..<count {
                var pointer: UnsafeMutablePointer<CChar>?
                try ORTNative.check(
                    ORTNative.api.pointee.SessionGetOutputName?(handle, index, allocator, &pointer),
                    action: "SessionGetOutputName"
                )
                results.append(try ORTNative.copyAllocatedString(pointer, allocator: allocator))
            }
            return results
        }
    }
}

public final class ORTSessionOptions {
    public private(set) var executionProviders: [(name: String, options: [String: String])] = []
    public private(set) var intraOpNumThreads: Int32 = 0
    public private(set) var graphOptimizationLevel: ORTGraphOptimizationLevel = .all
    public private(set) var optimizedModelFilePath: String?
    public private(set) var logID: String?
    public private(set) var logSeverityLevel: ORTLoggingLevel?
    public private(set) var configEntries: [String: String] = [:]
    public private(set) var customOpRegistrationFunctionName: String?
    public private(set) var customOpRegistrationFunctionPointer: ORTCAPIRegisterCustomOpsFnPtr?
    public private(set) var ortExtensionsCustomOpsEnabled = false
    public private(set) var coreMLExecutionProviderOptions: ORTCoreMLExecutionProviderOptions?
    public private(set) var coreMLExecutionProviderV2Options: [String: String]?
    public private(set) var xnnpackExecutionProviderOptions: ORTXnnpackExecutionProviderOptions?

    internal let handle: OpaquePointer

    public init() throws {
        var options: OpaquePointer?
        try ORTNative.check(ORTNative.api.pointee.CreateSessionOptions?(&options), action: "CreateSessionOptions")
        guard let options else {
            throw ORTError.runtimeFailure(action: "CreateSessionOptions", message: "Session options were not returned")
        }
        handle = options
    }

    deinit {
        ORTNative.api.pointee.ReleaseSessionOptions?(handle)
    }

    public func appendExecutionProvider(
        _ providerName: String,
        providerOptions: [String: String]
    ) throws {
        let sorted = providerOptions.keys.sorted().map { ($0, providerOptions[$0]!) }
        let keyStorage = sorted.map { Array($0.0.utf8CString) }
        let valueStorage = sorted.map { Array($0.1.utf8CString) }
        let keys = keyStorage.map { $0.withUnsafeBufferPointer { $0.baseAddress } }
        let values = valueStorage.map { $0.withUnsafeBufferPointer { $0.baseAddress } }

        try providerName.withCString { providerNamePointer in
            try keys.withUnsafeBufferPointer { keyBuffer in
                try values.withUnsafeBufferPointer { valueBuffer in
                    try ORTNative.check(
                        ORTNative.api.pointee.SessionOptionsAppendExecutionProvider?(
                            handle,
                            providerNamePointer,
                            keyBuffer.baseAddress,
                            valueBuffer.baseAddress,
                            sorted.count
                        ),
                        action: "SessionOptionsAppendExecutionProvider"
                    )
                }
            }
        }

        executionProviders.append((providerName, providerOptions))
    }

    public func setIntraOpNumThreads(_ count: Int32) throws {
        try ORTNative.check(
            ORTNative.api.pointee.SetIntraOpNumThreads?(handle, Int32(count)),
            action: "SetIntraOpNumThreads"
        )
        intraOpNumThreads = count
    }

    public func setGraphOptimizationLevel(_ level: ORTGraphOptimizationLevel) throws {
        try ORTNative.check(
            ORTNative.api.pointee.SetSessionGraphOptimizationLevel?(handle, level.native),
            action: "SetSessionGraphOptimizationLevel"
        )
        graphOptimizationLevel = level
    }

    public func setOptimizedModelFilePath(_ path: String) throws {
        try path.withCString { pointer in
            try ORTNative.check(
                ORTNative.api.pointee.SetOptimizedModelFilePath?(handle, pointer),
                action: "SetOptimizedModelFilePath"
            )
        }
        optimizedModelFilePath = path
    }

    public func setLogID(_ logID: String) throws {
        try logID.withCString { pointer in
            try ORTNative.check(
                ORTNative.api.pointee.SetSessionLogId?(handle, pointer),
                action: "SetSessionLogId"
            )
        }
        self.logID = logID
    }

    public func setLogSeverityLevel(_ level: ORTLoggingLevel) throws {
        try ORTNative.check(
            ORTNative.api.pointee.SetSessionLogSeverityLevel?(handle, level.rawValue),
            action: "SetSessionLogSeverityLevel"
        )
        logSeverityLevel = level
    }

    public func addConfigEntry(key: String, value: String) throws {
        try key.withCString { keyPointer in
            try value.withCString { valuePointer in
                try ORTNative.check(
                    ORTNative.api.pointee.AddSessionConfigEntry?(handle, keyPointer, valuePointer),
                    action: "AddSessionConfigEntry"
                )
            }
        }
        configEntries[key] = value
    }

    public func registerCustomOps(usingFunction name: String) throws {
        try name.withCString { pointer in
            try ORTNative.check(
                ORTNative.api.pointee.RegisterCustomOpsUsingFunction?(handle, pointer),
                action: "RegisterCustomOpsUsingFunction"
            )
        }
        customOpRegistrationFunctionName = name
    }

    public func registerCustomOps(usingFunctionPointer fn: ORTCAPIRegisterCustomOpsFnPtr) throws {
        let nativeFunction = unsafeBitCast(fn, to: RegisterCustomOpsFn.self)
        try ORTNative.check(
            nativeFunction(handle, ORTNative.apiBase),
            action: "RegisterCustomOpsUsingFunctionPointer"
        )
        customOpRegistrationFunctionPointer = fn
    }

    public func enableOrtExtensionsCustomOps() throws {
        try ORTNative.check(
            ORTNative.api.pointee.EnableOrtCustomOps?(handle),
            action: "EnableOrtCustomOps"
        )
        ortExtensionsCustomOpsEnabled = true
    }

    public func appendCoreMLExecutionProvider(options: ORTCoreMLExecutionProviderOptions) throws {
        try appendCoreMLExecutionProviderV2(providerOptions: options.providerOptions)
        coreMLExecutionProviderOptions = options
    }

    public func appendCoreMLExecutionProviderV2(providerOptions: [String: String]) throws {
        try appendExecutionProvider("CoreML", providerOptions: providerOptions)
        coreMLExecutionProviderV2Options = providerOptions
    }

    public func appendXnnpackExecutionProvider(options: ORTXnnpackExecutionProviderOptions) throws {
        try appendExecutionProvider(
            "XNNPACK",
            providerOptions: ["intra_op_num_threads": String(options.intraOpNumThreads)]
        )
        xnnpackExecutionProviderOptions = options
    }
}

public final class ORTRunOptions {
    public private(set) var logTag: String?
    public private(set) var logSeverityLevel: ORTLoggingLevel?
    public private(set) var configEntries: [String: String] = [:]

    internal let handle: OpaquePointer

    public init() throws {
        var options: OpaquePointer?
        try ORTNative.check(ORTNative.api.pointee.CreateRunOptions?(&options), action: "CreateRunOptions")
        guard let options else {
            throw ORTError.runtimeFailure(action: "CreateRunOptions", message: "Run options were not returned")
        }
        handle = options
    }

    deinit {
        ORTNative.api.pointee.ReleaseRunOptions?(handle)
    }

    public func setLogTag(_ tag: String) throws {
        try tag.withCString { pointer in
            try ORTNative.check(
                ORTNative.api.pointee.RunOptionsSetRunTag?(handle, pointer),
                action: "RunOptionsSetRunTag"
            )
        }
        logTag = tag
    }

    public func setLogSeverityLevel(_ level: ORTLoggingLevel) throws {
        try ORTNative.check(
            ORTNative.api.pointee.RunOptionsSetRunLogSeverityLevel?(handle, level.rawValue),
            action: "RunOptionsSetRunLogSeverityLevel"
        )
        logSeverityLevel = level
    }

    public func addConfigEntry(key: String, value: String) throws {
        try key.withCString { keyPointer in
            try value.withCString { valuePointer in
                try ORTNative.check(
                    ORTNative.api.pointee.AddRunConfigEntry?(handle, keyPointer, valuePointer),
                    action: "AddRunConfigEntry"
                )
            }
        }
        configEntries[key] = value
    }
}

public final class ORTValue {
    private let storage: ORTValueStorage

    internal var handle: OpaquePointer {
        storage.handle
    }

    public init(
        tensorData: NSMutableData,
        elementType: ORTTensorElementDataType,
        shape: [Int]
    ) throws {
        let elementCount = try ORTNative.tensorElementCount(shape: shape)
        let expectedLength = try ORTNative.elementByteCount(for: elementType) * elementCount
        guard tensorData.length >= expectedLength else {
            throw ORTError.invalidArgument("Tensor buffer is too small for shape \(shape) and element type \(elementType)")
        }

        var memoryInfo: OpaquePointer?
        try ORTNative.check(
            ORTNative.api.pointee.CreateCpuMemoryInfo?(OrtArenaAllocator, OrtMemTypeDefault, &memoryInfo),
            action: "CreateCpuMemoryInfo"
        )
        guard let memoryInfo else {
            throw ORTError.runtimeFailure(action: "CreateCpuMemoryInfo", message: "Memory info was not returned")
        }
        defer {
            ORTNative.api.pointee.ReleaseMemoryInfo?(memoryInfo)
        }

        var value: OpaquePointer?
        var nativeShape = ORTNative.makeShapeBuffer(shape)
        let nativeShapeCount = nativeShape.count
        try nativeShape.withUnsafeMutableBufferPointer { shapeBuffer in
            try ORTNative.check(
                ORTNative.api.pointee.CreateTensorWithDataAsOrtValue?(
                    memoryInfo,
                    tensorData.mutableBytes,
                    tensorData.length,
                    shapeBuffer.baseAddress,
                    nativeShapeCount,
                    elementType.native,
                    &value
                ),
                action: "CreateTensorWithDataAsOrtValue"
            )
        }

        guard let value else {
            throw ORTError.runtimeFailure(action: "CreateTensorWithDataAsOrtValue", message: "Value was not returned")
        }
        storage = ORTValueStorage(handle: value, retainedBuffer: tensorData)
    }

    public init(
        tensorStringData: [String],
        shape: [Int]
    ) throws {
        let elementCount = try ORTNative.tensorElementCount(shape: shape)
        guard elementCount == tensorStringData.count else {
            throw ORTError.invalidArgument("String tensor element count does not match shape")
        }

        var value: OpaquePointer?
        var nativeShape = ORTNative.makeShapeBuffer(shape)
        let nativeShapeCount = nativeShape.count
        try ORTNative.withAllocator { allocator in
            try nativeShape.withUnsafeMutableBufferPointer { shapeBuffer in
                try ORTNative.check(
                    ORTNative.api.pointee.CreateTensorAsOrtValue?(
                        allocator,
                        shapeBuffer.baseAddress,
                        nativeShapeCount,
                        ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING,
                        &value
                    ),
                    action: "CreateTensorAsOrtValue"
                )
            }
        }

        guard let value else {
            throw ORTError.runtimeFailure(action: "CreateTensorAsOrtValue", message: "String tensor value was not returned")
        }

        let stringStorage = tensorStringData.map { Array($0.utf8CString) }
        let stringPointers = stringStorage.map { $0.withUnsafeBufferPointer { $0.baseAddress } }
        try stringPointers.withUnsafeBufferPointer { pointerBuffer in
            try ORTNative.check(
                ORTNative.api.pointee.FillStringTensor?(value, pointerBuffer.baseAddress, stringPointers.count),
                action: "FillStringTensor"
            )
        }

        storage = ORTValueStorage(handle: value)
    }

    internal init(adopting handle: OpaquePointer) {
        storage = ORTValueStorage(handle: handle)
    }

    public func typeInfo() throws -> ORTValueTypeInfo {
        try ORTNative.fetchValueTypeInfo(from: handle)
    }

    public func tensorTypeAndShapeInfo() throws -> ORTTensorTypeAndShapeInfo {
        try ORTNative.fetchTensorInfo(from: handle)
    }

    public func tensorData() throws -> NSMutableData {
        let tensorInfo = try tensorTypeAndShapeInfo()
        guard tensorInfo.elementType != .string else {
            throw ORTError.unsupported("tensorData() is not available for string tensors")
        }

        var dataPointer: UnsafeMutableRawPointer?
        try ORTNative.check(
            ORTNative.api.pointee.GetTensorMutableData?(handle, &dataPointer),
            action: "GetTensorMutableData"
        )
        guard let dataPointer else {
            throw ORTError.runtimeFailure(action: "GetTensorMutableData", message: "Tensor data pointer was not returned")
        }

        let elementCount = try ORTNative.tensorElementCount(shape: tensorInfo.shape)
        let length = try ORTNative.elementByteCount(for: tensorInfo.elementType) * elementCount
        return NSMutableData(bytesNoCopy: dataPointer, length: length, freeWhenDone: false)
    }

    public func tensorStringData() throws -> [String] {
        let tensorInfo = try tensorTypeAndShapeInfo()
        guard tensorInfo.elementType == .string else {
            throw ORTError.unsupported("tensorStringData() is only available for string tensors")
        }

        let elementCount = try ORTNative.tensorElementCount(shape: tensorInfo.shape)
        var totalLength = 0
        try ORTNative.check(
            ORTNative.api.pointee.GetStringTensorDataLength?(handle, &totalLength),
            action: "GetStringTensorDataLength"
        )

        var content = Array(repeating: UInt8(0), count: totalLength)
        var offsets = Array(repeating: Int(0), count: elementCount)
        try ORTNative.check(
            ORTNative.api.pointee.GetStringTensorContent?(
                handle,
                &content,
                totalLength,
                &offsets,
                elementCount
            ),
            action: "GetStringTensorContent"
        )

        return (0..<elementCount).map { index in
            let start = offsets[index]
            let end = index + 1 < offsets.count ? offsets[index + 1] : totalLength
            return String(decoding: content[start..<end], as: UTF8.self)
        }
    }
}

public final class ORTCheckpoint {
    internal let handle: OpaquePointer

    public init(path: String) throws {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }

        var checkpoint: OpaquePointer?
        try path.withCString { pointer in
            try ORTNative.check(
                trainingApi.pointee.LoadCheckpoint?(pointer, &checkpoint),
                action: "LoadCheckpoint"
            )
        }

        guard let checkpoint else {
            throw ORTError.runtimeFailure(action: "LoadCheckpoint", message: "Checkpoint state was not returned")
        }
        handle = checkpoint
    }

    deinit {
        guard let trainingApi = ORTNative.trainingApi else {
            return
        }
        trainingApi.pointee.ReleaseCheckpointState?(handle)
    }

    public func saveCheckpoint(
        to path: String,
        includeOptimizerState: Bool
    ) throws {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }

        try path.withCString { pointer in
            try ORTNative.check(
                trainingApi.pointee.SaveCheckpoint?(handle, pointer, includeOptimizerState),
                action: "SaveCheckpoint"
            )
        }
    }

    public func addIntProperty(name: String, value: Int64) throws {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }

        var mutableValue = value
        try name.withCString { pointer in
            try ORTNative.check(
                trainingApi.pointee.AddProperty?(handle, pointer, OrtIntProperty, &mutableValue),
                action: "AddProperty"
            )
        }
    }

    public func addFloatProperty(name: String, value: Float) throws {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }

        var mutableValue = value
        try name.withCString { pointer in
            try ORTNative.check(
                trainingApi.pointee.AddProperty?(handle, pointer, OrtFloatProperty, &mutableValue),
                action: "AddProperty"
            )
        }
    }

    public func addStringProperty(name: String, value: String) throws {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }

        try name.withCString { namePointer in
            try value.withCString { valuePointer in
                let mutablePointer = UnsafeMutableRawPointer(mutating: valuePointer)
                try ORTNative.check(
                    trainingApi.pointee.AddProperty?(handle, namePointer, OrtStringProperty, mutablePointer),
                    action: "AddProperty"
                )
            }
        }
    }

    public func intProperty(name: String) throws -> Int64 {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }

        return try ORTNative.withAllocator { allocator in
            var propertyType = OrtPropertyType(rawValue: 0)
            var propertyValue: UnsafeMutableRawPointer?
            try name.withCString { pointer in
                try ORTNative.check(
                    trainingApi.pointee.GetProperty?(handle, pointer, allocator, &propertyType, &propertyValue),
                    action: "GetProperty"
                )
            }

            defer {
                if let propertyValue {
                    _ = ORTNative.api.pointee.AllocatorFree?(allocator, propertyValue)
                }
            }

            guard propertyType == OrtIntProperty, let propertyValue else {
                throw ORTError.invalidArgument("Checkpoint property \(name) is not an Int64")
            }
            return propertyValue.assumingMemoryBound(to: Int64.self).pointee
        }
    }

    public func floatProperty(name: String) throws -> Float {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }

        return try ORTNative.withAllocator { allocator in
            var propertyType = OrtPropertyType(rawValue: 0)
            var propertyValue: UnsafeMutableRawPointer?
            try name.withCString { pointer in
                try ORTNative.check(
                    trainingApi.pointee.GetProperty?(handle, pointer, allocator, &propertyType, &propertyValue),
                    action: "GetProperty"
                )
            }

            defer {
                if let propertyValue {
                    _ = ORTNative.api.pointee.AllocatorFree?(allocator, propertyValue)
                }
            }

            guard propertyType == OrtFloatProperty, let propertyValue else {
                throw ORTError.invalidArgument("Checkpoint property \(name) is not a Float")
            }
            return propertyValue.assumingMemoryBound(to: Float.self).pointee
        }
    }

    public func stringProperty(name: String) throws -> String? {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }

        return try ORTNative.withAllocator { allocator in
            var propertyType = OrtPropertyType(rawValue: 0)
            var propertyValue: UnsafeMutableRawPointer?
            try name.withCString { pointer in
                try ORTNative.check(
                    trainingApi.pointee.GetProperty?(handle, pointer, allocator, &propertyType, &propertyValue),
                    action: "GetProperty"
                )
            }

            guard propertyType == OrtStringProperty else {
                if let propertyValue {
                    _ = ORTNative.api.pointee.AllocatorFree?(allocator, propertyValue)
                }
                throw ORTError.invalidArgument("Checkpoint property \(name) is not a String")
            }

            let string = propertyValue.map { String(cString: $0.assumingMemoryBound(to: CChar.self)) }
            if let propertyValue {
                _ = ORTNative.api.pointee.AllocatorFree?(allocator, propertyValue)
            }
            return string
        }
    }
}

public final class ORTTrainingSession {
    private let handle: OpaquePointer
    private let checkpoint: ORTCheckpoint

    public init(
        env: ORTEnv,
        sessionOptions: ORTSessionOptions? = nil,
        checkpoint: ORTCheckpoint,
        trainModelPath: String,
        evalModelPath: String? = nil,
        optimizerModelPath: String? = nil
    ) throws {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }

        self.checkpoint = checkpoint

        let createdOptions = sessionOptions == nil ? try ORTSessionOptions() : nil
        let optionsHandle = sessionOptions?.handle ?? createdOptions!.handle
        var trainingSession: OpaquePointer?

        try trainModelPath.withCString { trainPointer in
            try evalModelPath?.withCString { evalPointer in
                try optimizerModelPath?.withCString { optimizerPointer in
                    try ORTNative.check(
                        trainingApi.pointee.CreateTrainingSession?(
                            env.handle,
                            optionsHandle,
                            checkpoint.handle,
                            trainPointer,
                            evalPointer,
                            optimizerPointer,
                            &trainingSession
                        ),
                        action: "CreateTrainingSession"
                    )
                }
            } ?? {
                try optimizerModelPath?.withCString { optimizerPointer in
                    try ORTNative.check(
                        trainingApi.pointee.CreateTrainingSession?(
                            env.handle,
                            optionsHandle,
                            checkpoint.handle,
                            trainPointer,
                            nil,
                            optimizerPointer,
                            &trainingSession
                        ),
                        action: "CreateTrainingSession"
                    )
                } ?? {
                    try ORTNative.check(
                        trainingApi.pointee.CreateTrainingSession?(
                            env.handle,
                            optionsHandle,
                            checkpoint.handle,
                            trainPointer,
                            nil,
                            nil,
                            &trainingSession
                        ),
                        action: "CreateTrainingSession"
                    )
                }()
            }()
        }

        guard let trainingSession else {
            throw ORTError.runtimeFailure(action: "CreateTrainingSession", message: "Training session was not returned")
        }
        handle = trainingSession
    }

    deinit {
        ORTNative.trainingApi?.pointee.ReleaseTrainingSession?(handle)
    }

    public func trainStep(inputValues: [ORTValue]) throws -> [ORTValue] {
        try runStep(
            countAction: "TrainingSessionGetTrainingModelOutputCount",
            countCall: { out in ORTNative.trainingApi!.pointee.TrainingSessionGetTrainingModelOutputCount?(handle, out) },
            stepAction: "TrainStep",
            stepCall: { outputCount, outputBuffer in
                let inputs = inputValues.map(\.handle)
                let inputPointers = inputs.map(Optional.init)
                return try inputPointers.withUnsafeBufferPointer { inputBuffer in
                    try ORTNative.check(
                        ORTNative.trainingApi!.pointee.TrainStep?(
                            handle,
                            nil,
                            inputs.count,
                            inputBuffer.baseAddress,
                            outputCount,
                            outputBuffer
                        ),
                        action: "TrainStep"
                    )
                }
            }
        )
    }

    public func evalStep(inputValues: [ORTValue]) throws -> [ORTValue] {
        try runStep(
            countAction: "TrainingSessionGetEvalModelOutputCount",
            countCall: { out in ORTNative.trainingApi!.pointee.TrainingSessionGetEvalModelOutputCount?(handle, out) },
            stepAction: "EvalStep",
            stepCall: { outputCount, outputBuffer in
                let inputs = inputValues.map(\.handle)
                let inputPointers = inputs.map(Optional.init)
                return try inputPointers.withUnsafeBufferPointer { inputBuffer in
                    try ORTNative.check(
                        ORTNative.trainingApi!.pointee.EvalStep?(
                            handle,
                            nil,
                            inputs.count,
                            inputBuffer.baseAddress,
                            outputCount,
                            outputBuffer
                        ),
                        action: "EvalStep"
                    )
                }
            }
        )
    }

    public func lazyResetGrad() throws {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }
        try ORTNative.check(trainingApi.pointee.LazyResetGrad?(handle), action: "LazyResetGrad")
    }

    public func optimizerStep() throws {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }
        try ORTNative.check(trainingApi.pointee.OptimizerStep?(handle, nil), action: "OptimizerStep")
    }

    public func trainInputNames() throws -> [String] {
        try trainingNames(
            countAction: "TrainingSessionGetTrainingModelInputCount",
            countCall: { out in ORTNative.trainingApi!.pointee.TrainingSessionGetTrainingModelInputCount?(handle, out) },
            nameAction: "TrainingSessionGetTrainingModelInputName",
            nameCall: { index, allocator, out in
                ORTNative.trainingApi!.pointee.TrainingSessionGetTrainingModelInputName?(handle, index, allocator, out)
            }
        )
    }

    public func evalInputNames() throws -> [String] {
        try trainingNames(
            countAction: "TrainingSessionGetEvalModelInputCount",
            countCall: { out in ORTNative.trainingApi!.pointee.TrainingSessionGetEvalModelInputCount?(handle, out) },
            nameAction: "TrainingSessionGetEvalModelInputName",
            nameCall: { index, allocator, out in
                ORTNative.trainingApi!.pointee.TrainingSessionGetEvalModelInputName?(handle, index, allocator, out)
            }
        )
    }

    public func trainOutputNames() throws -> [String] {
        try trainingNames(
            countAction: "TrainingSessionGetTrainingModelOutputCount",
            countCall: { out in ORTNative.trainingApi!.pointee.TrainingSessionGetTrainingModelOutputCount?(handle, out) },
            nameAction: "TrainingSessionGetTrainingModelOutputName",
            nameCall: { index, allocator, out in
                ORTNative.trainingApi!.pointee.TrainingSessionGetTrainingModelOutputName?(handle, index, allocator, out)
            }
        )
    }

    public func evalOutputNames() throws -> [String] {
        try trainingNames(
            countAction: "TrainingSessionGetEvalModelOutputCount",
            countCall: { out in ORTNative.trainingApi!.pointee.TrainingSessionGetEvalModelOutputCount?(handle, out) },
            nameAction: "TrainingSessionGetEvalModelOutputName",
            nameCall: { index, allocator, out in
                ORTNative.trainingApi!.pointee.TrainingSessionGetEvalModelOutputName?(handle, index, allocator, out)
            }
        )
    }

    public func registerLinearLRScheduler(
        warmupStepCount: Int64,
        totalStepCount: Int64,
        initialLr: Float
    ) throws {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }
        try ORTNative.check(
            trainingApi.pointee.RegisterLinearLRScheduler?(handle, warmupStepCount, totalStepCount, initialLr),
            action: "RegisterLinearLRScheduler"
        )
    }

    public func schedulerStep() throws {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }
        try ORTNative.check(trainingApi.pointee.SchedulerStep?(handle), action: "SchedulerStep")
    }

    public func learningRate() throws -> Float {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }
        var learningRate: Float = 0
        try ORTNative.check(trainingApi.pointee.GetLearningRate?(handle, &learningRate), action: "GetLearningRate")
        return learningRate
    }

    public func setLearningRate(_ lr: Float) throws {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }
        try ORTNative.check(trainingApi.pointee.SetLearningRate?(handle, lr), action: "SetLearningRate")
    }

    public func fromBuffer(_ buffer: ORTValue) throws {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }
        try ORTNative.check(trainingApi.pointee.CopyBufferToParameters?(handle, buffer.handle, false), action: "CopyBufferToParameters")
    }

    public func toBuffer(onlyTrainable: Bool) throws -> ORTValue {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }

        var parameterCount = 0
        try ORTNative.check(
            trainingApi.pointee.GetParametersSize?(handle, &parameterCount, onlyTrainable),
            action: "GetParametersSize"
        )

        let byteCount = parameterCount * MemoryLayout<Float>.stride
        let buffer = NSMutableData(length: byteCount) ?? NSMutableData()
        let value = try ORTValue(tensorData: buffer, elementType: .float, shape: [parameterCount])
        try ORTNative.check(
            trainingApi.pointee.CopyParametersToBuffer?(handle, value.handle, onlyTrainable),
            action: "CopyParametersToBuffer"
        )
        return value
    }

    public func exportModelForInference(
        outputPath: String,
        graphOutputNames: [String]
    ) throws {
        guard let trainingApi = ORTNative.trainingApi else {
            throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
        }

        try outputPath.withCString { pathPointer in
            try ORTNative.withCStringArray(graphOutputNames) { namePointers in
                try namePointers.withUnsafeBufferPointer { nameBuffer in
                    try ORTNative.check(
                        trainingApi.pointee.ExportModelForInferencing?(
                            handle,
                            pathPointer,
                            graphOutputNames.count,
                            nameBuffer.baseAddress
                        ),
                        action: "ExportModelForInferencing"
                    )
                }
            }
        }
    }

    private func trainingNames(
        countAction: String,
        countCall: (_ out: UnsafeMutablePointer<Int>) -> OpaquePointer?,
        nameAction: String,
        nameCall: (
            _ index: Int,
            _ allocator: UnsafeMutablePointer<OrtAllocator>,
            _ out: UnsafeMutablePointer<UnsafeMutablePointer<CChar>?>
        ) -> OpaquePointer?
    ) throws -> [String] {
        var count = 0
        try ORTNative.check(countCall(&count), action: countAction)

        return try ORTNative.withAllocator { allocator in
            var names: [String] = []
            names.reserveCapacity(count)
            for index in 0..<count {
                var namePointer: UnsafeMutablePointer<CChar>?
                try ORTNative.check(nameCall(index, allocator, &namePointer), action: nameAction)
                names.append(try ORTNative.copyAllocatedString(namePointer, allocator: allocator))
            }
            return names
        }
    }

    private func runStep(
        countAction: String,
        countCall: (_ out: UnsafeMutablePointer<Int>) -> OpaquePointer?,
        stepAction: String,
        stepCall: (_ outputCount: Int, _ outputBuffer: UnsafeMutablePointer<OpaquePointer?>?) throws -> Void
    ) throws -> [ORTValue] {
        var outputCount = 0
        try ORTNative.check(countCall(&outputCount), action: countAction)

        var outputs = Array<OpaquePointer?>(repeating: nil, count: outputCount)
        try outputs.withUnsafeMutableBufferPointer { buffer in
            try stepCall(outputCount, buffer.baseAddress)
        }

        return try outputs.enumerated().map { index, handle in
            guard let handle else {
                throw ORTError.runtimeFailure(action: stepAction, message: "Output \(index) was not returned")
            }
            return ORTValue(adopting: handle)
        }
    }
}