@preconcurrency import COnnxRuntime
import Foundation

internal enum ORTNative {
    nonisolated(unsafe) static let apiBase = OrtGetApiBase()
    nonisolated(unsafe) static let api = apiBase!.pointee.GetApi(UInt32(ORT_API_VERSION))!
    nonisolated(unsafe) static let trainingApi = api.pointee.GetTrainingApi?(UInt32(ORT_API_VERSION))

    static func check(_ status: OpaquePointer?, action: String) throws {
        guard let status else {
            return
        }

        let message = api.pointee.GetErrorMessage?(status).map { String(cString: $0) } ?? "Unknown ONNX Runtime error"
        api.pointee.ReleaseStatus?(status)
        throw ORTError.runtimeFailure(action: action, message: message)
    }

    static func withAllocator<Result>(_ body: (UnsafeMutablePointer<OrtAllocator>) throws -> Result) throws -> Result {
        var allocator: UnsafeMutablePointer<OrtAllocator>?
        try check(api.pointee.GetAllocatorWithDefaultOptions?(&allocator), action: "GetAllocatorWithDefaultOptions")
        guard let allocator else {
            throw ORTError.runtimeFailure(action: "GetAllocatorWithDefaultOptions", message: "Allocator was not returned")
        }
        return try body(allocator)
    }

    static func copyAllocatedString(
        _ stringPointer: UnsafeMutablePointer<CChar>?,
        allocator: UnsafeMutablePointer<OrtAllocator>
    ) throws -> String {
        guard let stringPointer else {
            throw ORTError.runtimeFailure(action: "copyAllocatedString", message: "Expected allocated string")
        }

        defer {
            _ = api.pointee.AllocatorFree?(allocator, stringPointer)
        }

        return String(cString: stringPointer)
    }

    static func tensorElementCount(shape: [Int]) throws -> Int {
        guard shape.allSatisfy({ $0 >= 0 }) else {
            throw ORTError.invalidArgument("Tensor shape contains a negative dimension: \(shape)")
        }

        return shape.reduce(1, *)
    }

    static func elementByteCount(for type: ORTTensorElementDataType) throws -> Int {
        switch type {
        case .float:
            return MemoryLayout<Float>.stride
        case .int8:
            return MemoryLayout<Int8>.stride
        case .uInt8:
            return MemoryLayout<UInt8>.stride
        case .int32:
            return MemoryLayout<Int32>.stride
        case .uInt32:
            return MemoryLayout<UInt32>.stride
        case .int64:
            return MemoryLayout<Int64>.stride
        case .uInt64:
            return MemoryLayout<UInt64>.stride
        case .string:
            throw ORTError.invalidArgument("String tensors do not have a fixed byte width")
        case .undefined:
            throw ORTError.invalidArgument("Undefined tensor element type is not supported")
        }
    }

    static func makeShapeBuffer(_ shape: [Int]) -> [Int64] {
        shape.map(Int64.init)
    }

    static func withCStringArray<Result>(
        _ strings: [String],
        _ body: ([UnsafePointer<CChar>?]) throws -> Result
    ) rethrows -> Result {
        let storage = strings.map { Array($0.utf8CString) }
        let pointers = storage.map { $0.withUnsafeBufferPointer { $0.baseAddress } }
        return try body(pointers)
    }

    static func fetchTensorInfo(from value: OpaquePointer) throws -> ORTTensorTypeAndShapeInfo {
        var info: OpaquePointer?
        try check(api.pointee.GetTensorTypeAndShape?(value, &info), action: "GetTensorTypeAndShape")
        guard let info else {
            throw ORTError.runtimeFailure(action: "GetTensorTypeAndShape", message: "Tensor info was not returned")
        }

        defer {
            api.pointee.ReleaseTensorTypeAndShapeInfo?(info)
        }

        var elementType = ONNXTensorElementDataType(rawValue: 0)
        try check(api.pointee.GetTensorElementType?(info, &elementType), action: "GetTensorElementType")

        var dimensionsCount: Int = 0
        try check(api.pointee.GetDimensionsCount?(info, &dimensionsCount), action: "GetDimensionsCount")

        var dimensions = Array(repeating: Int64(0), count: dimensionsCount)
        if dimensionsCount > 0 {
            try check(
                api.pointee.GetDimensions?(info, &dimensions, dimensionsCount),
                action: "GetDimensions"
            )
        }

        return ORTTensorTypeAndShapeInfo(
            elementType: try ORTTensorElementDataType(native: elementType),
            shape: dimensions.map(Int.init)
        )
    }

    static func fetchValueTypeInfo(from value: OpaquePointer) throws -> ORTValueTypeInfo {
        var typeInfo: OpaquePointer?
        try check(api.pointee.GetTypeInfo?(value, &typeInfo), action: "GetTypeInfo")
        guard let typeInfo else {
            throw ORTError.runtimeFailure(action: "GetTypeInfo", message: "Type info was not returned")
        }

        defer {
            api.pointee.ReleaseTypeInfo?(typeInfo)
        }

        var onnxType = ONNXType(rawValue: 0)
        try check(api.pointee.GetOnnxTypeFromTypeInfo?(typeInfo, &onnxType), action: "GetOnnxTypeFromTypeInfo")

        var tensorInfoPointer: OpaquePointer?
        try check(api.pointee.CastTypeInfoToTensorInfo?(typeInfo, &tensorInfoPointer), action: "CastTypeInfoToTensorInfo")

        let tensorInfo = try tensorInfoPointer.map { pointer in
            try fetchTensorInfoView(from: pointer)
        }

        return ORTValueTypeInfo(
            type: try ORTValueType(native: onnxType),
            tensorTypeAndShapeInfo: tensorInfo
        )
    }

    static func fetchTensorInfoView(from info: OpaquePointer) throws -> ORTTensorTypeAndShapeInfo {
        var elementType = ONNXTensorElementDataType(rawValue: 0)
        try check(api.pointee.GetTensorElementType?(info, &elementType), action: "GetTensorElementType")

        var dimensionsCount: Int = 0
        try check(api.pointee.GetDimensionsCount?(info, &dimensionsCount), action: "GetDimensionsCount")

        var dimensions = Array(repeating: Int64(0), count: dimensionsCount)
        if dimensionsCount > 0 {
            try check(api.pointee.GetDimensions?(info, &dimensions, dimensionsCount), action: "GetDimensions")
        }

        return ORTTensorTypeAndShapeInfo(
            elementType: try ORTTensorElementDataType(native: elementType),
            shape: dimensions.map(Int.init)
        )
    }

    static func boolString(_ value: Bool) -> String {
        value ? "1" : "0"
    }
}