vOOice/VoiceInk/Services/ModelPrewarmService.swift

import Foundation
import SwiftData
import os
import AppKit

@MainActor
final class ModelPrewarmService: ObservableObject {
    private let whisperState: WhisperState
    private let modelContext: ModelContext
    private let logger = Logger(subsystem: "com.prakashjoshipax.voiceink", category: "ModelPrewarm")

    // Services (initialized lazily)
    private var localTranscriptionService: LocalTranscriptionService?
    private var parakeetTranscriptionService: ParakeetTranscriptionService?
    private let nativeAppleTranscriptionService = NativeAppleTranscriptionService()
    private let cloudTranscriptionService = CloudTranscriptionService()

    // Sample audio for prewarming
    private let prewarmAudioURL = Bundle.main.url(forResource: "esc", withExtension: "wav")

    // User preference key
    private let prewarmEnabledKey = "PrewarmModelOnWake"

    init(whisperState: WhisperState, modelContext: ModelContext) {
        self.whisperState = whisperState
        self.modelContext = modelContext
        setupNotifications()
        schedulePrewarmOnAppLaunch()
    }

    // MARK: - Notification Setup

    private func setupNotifications() {
        let center = NSWorkspace.shared.notificationCenter

        // Trigger on wake from sleep
        center.addObserver(
            self,
            selector: #selector(schedulePrewarm),
            name: NSWorkspace.didWakeNotification,
            object: nil
        )

        logger.notice("🌅 ModelPrewarmService initialized - listening for wake and app launch")
    }

    // MARK: - Trigger Handlers

    /// Trigger on app launch (cold start)
    private func schedulePrewarmOnAppLaunch() {
        logger.notice("🌅 App launched, scheduling prewarm")
        Task {
            try? await Task.sleep(for: .seconds(3))
            await performPrewarm()
        }
    }

    /// Trigger on wake from sleep or screen unlock
    @objc private func schedulePrewarm() {
        logger.notice("🌅 Mac activity detected (wake/unlock), scheduling prewarm")
        Task {
            try? await Task.sleep(for: .seconds(3))
            await performPrewarm()
        }
    }

    // MARK: - Core Prewarming Logic

    private func performPrewarm() async {
        guard shouldPrewarm() else { return }

        guard let audioURL = prewarmAudioURL else {
            logger.error("❌ Prewarm audio file (esc.wav) not found")
            return
        }

        guard let currentModel = whisperState.currentTranscriptionModel else {
            logger.notice("🌅 No model selected, skipping prewarm")
            return
        }

        logger.notice("🌅 Prewarming \(currentModel.displayName)")
        let startTime = Date()

        do {
            // Initialize services lazily
            if localTranscriptionService == nil {
                localTranscriptionService = LocalTranscriptionService(
                    modelsDirectory: whisperState.modelsDirectory,
                    whisperState: whisperState
                )
            }
            if parakeetTranscriptionService == nil {
                parakeetTranscriptionService = ParakeetTranscriptionService()
            }

            // Run transcription to trigger model loading and ANE compilation
            let transcribedText: String
            switch currentModel.provider {
            case .local:
                transcribedText = try await localTranscriptionService!.transcribe(audioURL: audioURL, model: currentModel)
            case .parakeet:
                transcribedText = try await parakeetTranscriptionService!.transcribe(audioURL: audioURL, model: currentModel)
            case .nativeApple:
                transcribedText = try await nativeAppleTranscriptionService.transcribe(audioURL: audioURL, model: currentModel)
            default:
                transcribedText = try await cloudTranscriptionService.transcribe(audioURL: audioURL, model: currentModel)
            }

            let duration = Date().timeIntervalSince(startTime)

            // Save for telemetry
            let transcription = Transcription(
                text: "[PREWARM] \(transcribedText)",
                duration: 1.0,
                audioFileURL: audioURL.absoluteString,
                transcriptionModelName: currentModel.displayName,
                transcriptionDuration: duration
            )
            modelContext.insert(transcription)
            try? modelContext.save()

            logger.notice("🌅 Prewarm completed in \(String(format: "%.2f", duration))s")

        } catch {
            logger.error("❌ Prewarm failed: \(error.localizedDescription)")
        }
    }

    // MARK: - Validation

    private func shouldPrewarm() -> Bool {
        // Check if user has enabled prewarming
        let isEnabled = UserDefaults.standard.object(forKey: prewarmEnabledKey) as? Bool ?? true
        guard isEnabled else {
            logger.notice("🌅 Prewarm disabled by user")
            return false
        }

        // Only prewarm local models (Parakeet and Whisper need ANE compilation)
        guard let model = whisperState.currentTranscriptionModel else {
            return false
        }

        switch model.provider {
        case .local, .parakeet:
            return true
        default:
            logger.notice("🌅 Skipping prewarm - cloud models don't need it")
            return false
        }
    }

    deinit {
        NSWorkspace.shared.notificationCenter.removeObserver(self)
        logger.notice("🌅 ModelPrewarmService deinitialized")
    }
}