import CoreML
import Foundation
import NucleicCore
import OSLog

/// Small seam between composer routing and the bundled task-specific model. Tests can inject a
/// deterministic implementation without loading Core ML; production uses the shared actor below.
protocol PurposeModelClassifying: Sendable {
    func classify(prompt: String) async -> PurposeVerdict?
}

/// Lazy, actor-isolated Core ML inference for the Intelligence Bar's purpose router.
///
/// The model is loaded once on first meaningful draft, configured for CPU + Neural Engine, and
/// kept warm for subsequent composer changes. Every failure returns nil so the layered service can
/// continue through AFM, the delegated agent, and finally the deterministic heuristic.
actor PurposeMLClassifier: PurposeModelClassifying {
    static let shared = PurposeMLClassifier()

    static let enabledKey = "nucleic.purposeClassifierEnabled"
    static let modelKey = "nucleic.purposeClassifierModel"
    static let artifactVersion = "purpose-lite-v1-w8a8"

    private let logger = Logger(subsystem: "com.blakeslee.Nucleic", category: "PurposeClassifier")
    private var runtime: Runtime?
    private var loadAttempted = false
    private var didLogFailure = false

    func classify(prompt: String) async -> PurposeVerdict? {
        guard Self.isEnabled, Self.selectedModelIsAvailable else { return nil }
        do {
            let runtime = try loadRuntime()
            let encoded = runtime.tokenizer.encode(prompt)
            let logits = try Self.predict(model: runtime.model, input: encoded)
            return runtime.calibration.verdict(logits: logits)
        } catch {
            if !didLogFailure {
                didLogFailure = true
                logger.error("Purpose classifier failed: \(String(describing: error), privacy: .public)")
            }
            return nil
        }
    }

    private static var isEnabled: Bool {
        let defaults = UserDefaults.standard
        return (defaults.object(forKey: enabledKey) as? Bool) ?? true
    }

    private static var selectedModelIsAvailable: Bool {
        let selected = UserDefaults.standard.string(forKey: modelKey)?
            .trimmingCharacters(in: .whitespacesAndNewlines).lowercased() ?? "auto"
        return selected.isEmpty || selected == "auto" || selected == "lite"
    }

    private func loadRuntime() throws -> Runtime {
        if let runtime { return runtime }
        guard !loadAttempted else { throw PurposeMLClassifierError.priorLoadFailure }
        loadAttempted = true

        guard let vocabularyURL = PurposeClassifierResources.url(
            named: "vocab", extensions: ["txt"]),
              let calibrationURL = PurposeClassifierResources.url(
                named: "calibration", extensions: ["json"]),
              let modelURL = PurposeClassifierResources.url(
                named: Self.artifactVersion, extensions: ["mlmodelc", "mlpackage", "mlmodel"])
        else {
            throw PurposeMLClassifierError.missingResource
        }

        let tokenizer = try PurposeWordPieceTokenizer(
            vocabularyText: String(contentsOf: vocabularyURL, encoding: .utf8))
        let calibration = try JSONDecoder().decode(
            PurposeCalibration.self, from: Data(contentsOf: calibrationURL))
        guard calibration.modelVersion == "purpose-lite-v1",
              calibration.labels == PromptPurpose.classifiable.map(\.rawValue)
        else {
            throw PurposeMLClassifierError.invalidCalibration
        }
        let configuration = MLModelConfiguration()
        configuration.computeUnits = .cpuAndNeuralEngine
        let loadURL: URL
        if modelURL.pathExtension == "mlmodelc" {
            loadURL = modelURL
        } else {
            loadURL = try MLModel.compileModel(at: modelURL)
        }
        let loaded = Runtime(
            model: try MLModel(contentsOf: loadURL, configuration: configuration),
            tokenizer: tokenizer,
            calibration: calibration)
        runtime = loaded
        logger.info("Loaded \(Self.artifactVersion, privacy: .public) on CPU + Neural Engine")
        return loaded
    }

    private static func predict(model: MLModel, input: PurposeModelInput) throws -> [Double] {
        let provider = try MLDictionaryFeatureProvider(dictionary: [
            "input_ids": MLFeatureValue(multiArray: try multiArray(input.inputIDs)),
            "attention_mask": MLFeatureValue(multiArray: try multiArray(input.attentionMask)),
            "token_type_ids": MLFeatureValue(multiArray: try multiArray(input.tokenTypeIDs)),
        ])
        let output = try model.prediction(from: provider)
        guard let values = output.featureValue(for: "logits")?.multiArrayValue,
              values.count == PromptPurpose.classifiable.count
        else {
            throw PurposeMLClassifierError.invalidLogits
        }
        return (0..<values.count).map { values[$0].doubleValue }
    }

    private static func multiArray(_ values: [Int32]) throws -> MLMultiArray {
        guard values.count == PurposeWordPieceTokenizer.sequenceLength else {
            throw PurposeMLClassifierError.invalidInputShape
        }
        let result = try MLMultiArray(
            shape: [1, NSNumber(value: PurposeWordPieceTokenizer.sequenceLength)],
            dataType: .int32)
        for (index, value) in values.enumerated() {
            result[index] = NSNumber(value: value)
        }
        return result
    }

    private struct Runtime {
        let model: MLModel
        let tokenizer: PurposeWordPieceTokenizer
        let calibration: PurposeCalibration
    }
}

private enum PurposeMLClassifierError: Error {
    case missingResource
    case priorLoadFailure
    case invalidInputShape
    case invalidLogits
    case invalidCalibration
}

private struct PurposeCalibration: Decodable {
    struct Confidence: Decodable {
        struct Threshold: Decodable { let minimumScore: Double }
        let high: Threshold
        let medium: Threshold
    }

    let confidence: Confidence
    let labels: [String]
    let modelVersion: String
    let temperature: Double

    func verdict(logits: [Double]) -> PurposeVerdict? {
        guard logits.count == labels.count, logits.count >= 2, temperature > 0 else { return nil }
        let adjusted = logits.map { $0 / temperature }
        guard let maximum = adjusted.max() else { return nil }
        let exponentials = adjusted.map { Foundation.exp($0 - maximum) }
        let total = exponentials.reduce(0, +)
        guard total.isFinite, total > 0 else { return nil }
        let probabilities = exponentials.map { $0 / total }
        let ranked = probabilities.indices.sorted { probabilities[$0] > probabilities[$1] }
        let top = ranked[0]
        let topProbability = probabilities[top]
        let margin = topProbability - probabilities[ranked[1]]
        let score = topProbability * (0.5 + 0.5 * margin)
        let confidenceLevel: PurposeVerdict.Confidence
        if score >= confidence.high.minimumScore {
            confidenceLevel = .high
        } else if score >= confidence.medium.minimumScore {
            confidenceLevel = .medium
        } else {
            confidenceLevel = .low
        }
        guard let purpose = PromptPurpose(rawValue: labels[top]), purpose != .general else {
            return nil
        }
        let reason = String(
            format: "%@ score %.3f (p=%.3f)",
            PurposeMLClassifier.artifactVersion, score, topProbability)
        return PurposeVerdict(
            purpose: purpose, confidence: confidenceLevel, source: .custom, reason: reason)
    }
}

/// Resource lookup supports both SwiftPM's compiled `.mlmodelc` output and a copied source model
/// package. The latter is compiled lazily on the classifier actor if the package manager leaves it
/// untouched, which keeps launch and the send path free of model work.
enum PurposeClassifierResources {
    static func url(named name: String, extensions: [String]) -> URL? {
        let bundle = Bundle.module
        for fileExtension in extensions {
            if let direct = bundle.url(
                forResource: name, withExtension: fileExtension,
                subdirectory: "PurposeClassifier")
                ?? bundle.url(forResource: name, withExtension: fileExtension) {
                return direct
            }
        }
        guard let root = bundle.resourceURL,
              let enumerator = FileManager.default.enumerator(
                at: root, includingPropertiesForKeys: nil)
        else { return nil }
        let accepted = Set(extensions)
        for case let candidate as URL in enumerator
        where candidate.deletingPathExtension().lastPathComponent == name
            && accepted.contains(candidate.pathExtension) {
            return candidate
        }
        return nil
    }
}
