import CoreML
import Darwin
import Foundation

/// Load and prediction timing for a compiled Nemotron bundle on macOS or iOS.
/// Invoke one bundle and one compute-unit configuration per app launch.
@available(macOS 15.0, iOS 18.0, *)
public enum NemotronExportProbe {
    public struct Report: Codable, Sendable {
        public let bundle: String
        public let os: String
        public let machine: String
        public let computeUnits: String
        public let models: [ModelRecord]
    }

    public struct ModelRecord: Codable, Sendable {
        public let component: String
        public let firstObserved: Measurement
        public let sameProcessReload: Measurement
        public let plannedOperationsByPreferredDevice: [String: Int]
        public let inputs: [FeatureSignature]
        public let outputs: [FeatureSignature]
    }

    public struct Measurement: Codable, Sendable {
        public let loadMs: Double
        public let firstPredictionMs: Double
        public let medianRepeatedPredictionMs: Double
    }

    public struct FeatureSignature: Codable, Sendable {
        public let name: String
        public let shape: [Int]
        public let dtype: String
    }

    public static func run(bundleURL: URL, computeUnits: MLComputeUnits) async throws -> Report {
        let configuration = MLModelConfiguration()
        configuration.computeUnits = computeUnits
        var records: [ModelRecord] = []
        for name in ["encoder", "decoder", "joint"] {
            let url = bundleURL.appendingPathComponent("\(name).mlmodelc", isDirectory: true)
            print("MODEL_START \(name)")
            let first = try measure(url: url, configuration: configuration)
            print("MODEL_FIRST_COMPLETE \(name) loadMs=\(first.0.loadMs)")
            let repeated = try measure(url: url, configuration: configuration)
            print("MODEL_RELOAD_COMPLETE \(name) loadMs=\(repeated.0.loadMs)")
            let plan = try await MLComputePlan.load(contentsOf: url, configuration: configuration)
            print("MODEL_PLAN_COMPLETE \(name)")
            var devices: [String: Int] = [:]
            if case let .program(program) = plan.modelStructure,
               let main = program.functions["main"] {
                for operation in main.block.operations {
                    guard let usage = plan.deviceUsage(for: operation) else { continue }
                    let device: String
                    switch usage.preferred {
                    case .cpu: device = "cpu"
                    case .gpu: device = "gpu"
                    case .neuralEngine: device = "neuralEngine"
                    @unknown default: device = "unknown"
                    }
                    devices[device, default: 0] += 1
                }
            }
            records.append(ModelRecord(
                component: name, firstObserved: first.0, sameProcessReload: repeated.0,
                plannedOperationsByPreferredDevice: devices, inputs: first.1, outputs: first.2))
        }
        var system = utsname()
        uname(&system)
        let machine = withUnsafePointer(to: &system.machine) {
            String(cString: UnsafeRawPointer($0).assumingMemoryBound(to: CChar.self))
        }
        let units: String
        switch computeUnits {
        case .cpuOnly: units = "cpuOnly"
        case .cpuAndGPU: units = "cpuAndGPU"
        case .cpuAndNeuralEngine: units = "cpuAndNeuralEngine"
        case .all: units = "all"
        @unknown default: units = "unknown"
        }
        return Report(
            bundle: bundleURL.lastPathComponent,
            os: ProcessInfo.processInfo.operatingSystemVersionString,
            machine: machine, computeUnits: units, models: records)
    }

    private static func signatures(_ descriptions: [String: MLFeatureDescription]) throws -> [FeatureSignature] {
        try descriptions.keys.sorted().map { name in
            guard let array = descriptions[name]?.multiArrayConstraint else {
                throw ProbeError.unsupportedFeature(name)
            }
            return FeatureSignature(
                name: name, shape: array.shape.map(\.intValue), dtype: String(describing: array.dataType))
        }
    }

    private static func measure(
        url: URL, configuration: MLModelConfiguration
    ) throws -> (Measurement, [FeatureSignature], [FeatureSignature]) {
        try autoreleasepool {
            let start = ProcessInfo.processInfo.systemUptime
            let model = try MLModel(contentsOf: url, configuration: configuration)
            let loadMs = (ProcessInfo.processInfo.systemUptime - start) * 1_000
            let inputSignatures = try signatures(model.modelDescription.inputDescriptionsByName)
            let outputSignatures = try signatures(model.modelDescription.outputDescriptionsByName)
            var inputs: [String: Any] = [:]
            for (name, description) in model.modelDescription.inputDescriptionsByName {
                guard let constraint = description.multiArrayConstraint else {
                    throw ProbeError.unsupportedFeature(name)
                }
                let array = try MLMultiArray(shape: constraint.shape, dataType: constraint.dataType)
                let itemBytes: Int
                switch array.dataType {
                case .float16: itemBytes = 2
                case .float32, .int32: itemBytes = 4
                case .double: itemBytes = 8
                default: throw ProbeError.unsupportedFeature(name)
                }
                memset(array.dataPointer, 0, array.count * itemBytes)
                if name == "audio_length" { array[0] = 32 }
                if name == "language_mask" { array[0] = 1 }
                if name == "token" { array[0] = 13087 }
                inputs[name] = array
            }
            let features = try MLDictionaryFeatureProvider(dictionary: inputs)
            var milliseconds: [Double] = []
            for _ in 0..<4 {
                let start = ProcessInfo.processInfo.systemUptime
                let output = try model.prediction(from: features)
                milliseconds.append((ProcessInfo.processInfo.systemUptime - start) * 1_000)
                for name in output.featureNames {
                    guard let array = output.featureValue(for: name)?.multiArrayValue else {
                        throw ProbeError.unsupportedFeature(name)
                    }
                    for index in 0..<min(array.count, 16) where !array[index].doubleValue.isFinite {
                        throw ProbeError.nonFiniteOutput(name)
                    }
                }
            }
            let repeated = milliseconds.dropFirst().sorted()
            return (
                Measurement(loadMs: loadMs, firstPredictionMs: milliseconds[0],
                            medianRepeatedPredictionMs: repeated[repeated.count / 2]),
                inputSignatures, outputSignatures)
        }
    }

    private enum ProbeError: Error {
        case unsupportedFeature(String)
        case nonFiniteOutput(String)
    }
}

#if LOAD_PROBE_CLI
@main
struct LoadProbeCLI {
    static func main() async throws {
        guard CommandLine.arguments.count == 4 else {
            throw NSError(domain: "NemotronExportProbe", code: 1, userInfo: [
                NSLocalizedDescriptionKey: "usage: probe bundle-directory ane|cpu|gpu|all output.json",
            ])
        }
        let units: MLComputeUnits
        switch CommandLine.arguments[2] {
        case "ane": units = .cpuAndNeuralEngine
        case "cpu": units = .cpuOnly
        case "gpu": units = .cpuAndGPU
        case "all": units = .all
        default: throw NSError(domain: "NemotronExportProbe", code: 2)
        }
        let report = try await NemotronExportProbe.run(
            bundleURL: URL(fileURLWithPath: CommandLine.arguments[1], isDirectory: true), computeUnits: units)
        let encoder = JSONEncoder()
        encoder.outputFormatting = [.prettyPrinted, .sortedKeys]
        let json = try encoder.encode(report)
        try json.write(to: URL(fileURLWithPath: CommandLine.arguments[3]))
        print(String(decoding: json, as: UTF8.self))
    }
}
#endif