LlmManager.swift

//
//  LlmManager.swift
//  TranslateBlue / BlueGemma iOS Voice Session
//
//  Created for TranslateBlue Issue #38: iOS Metal GPU Accelerated LLM Engine
//  Optimized for Google Gemma 4 E2B-it (LiteRT-LM / MediaPipe Tasks GenAI)
//

import Foundation
import SwiftUI
import MediaPipeTasksGenAI

@MainActor
public class LlmManager: ObservableObject {
    @Published public var responseText: String = ""
    @Published public var isGenerating: Bool = false
    @Published public var isInitialized: Bool = false
    @Published public var errorMessage: String? = nil
    
    private var llmInference: LlmInference?
    private let modelFileName = "gemma-4-E2B-it-ios"
    private let modelExtension = "task"
    
    public init() {
        setupEngine()
        setupMemoryLifecycleObservers()
    }
    
    deinit {
        NotificationCenter.default.removeObserver(self)
    }
    
    /// Initializes the LiteRT-LM Metal GPU Accelerated Engine with dynamic KV cache quantization.
    public func setupEngine() {
        guard let modelPath = Bundle.main.path(forResource: modelFileName, ofType: modelExtension) ??
                Bundle.main.path(forResource: "Bluegemma", ofType: "litertlm") else {
            let err = "[ERROR] Model file '\(modelFileName).\(modelExtension)' not found in app bundle or document path."
            print(err)
            self.errorMessage = err
            return
        }
        
        let options = LlmInference.Options()
        options.modelPath = modelPath
        options.maxTokens = 512
        options.temperature = 0.7
        options.topK = 40
        
        // ==========================================
        // METAL GPU & ANE MEMORY OPTIMIZATION (iOS)
        // ==========================================
        // Enable dynamic KV cache quantization to prevent OS Jetsam memory termination
        options.kvCacheConfig.enableQuantization = true
        
        do {
            self.llmInference = try LlmInference(options: options)
            self.isInitialized = true
            self.errorMessage = nil
            print("[OK] LiteRT-LM Engine initialized with Metal GPU hardware acceleration successfully.")
        } catch {
            let err = "[FATAL] Failed to initialize LiteRT-LM Metal Engine: \(error.localizedDescription)"
            print(err)
            self.errorMessage = err
            self.isInitialized = false
        }
    }
    
    /// Asynchronously streams direct translation responses to SwiftUI views.
    public func generateTranslationStream(targetLanguage: String, sourceText: String) async {
        guard let engine = llmInference else {
            self.responseText = "Error: Model engine not loaded."
            return
        }
        
        self.isGenerating = true
        self.responseText = ""
        
        let systemPrompt = "You are a direct professional translator. Output ONLY the translated text in \(targetLanguage). Do not write any intro, greeting, explanation, or options."
        let formattedPrompt = "<start_of_turn>system\n\(systemPrompt)<end_of_turn>\n<start_of_turn>user\nTranslate the following sentence into \(targetLanguage): \(sourceText)<end_of_turn>\n<start_of_turn>model\n"
        
        do {
            // Execute background inference off the main thread with userInitiated priority
            try await Task.detached(priority: .userInitiated) {
                autoreleasepool {
                    try? engine.generateAsync(prompt: formattedPrompt) { [weak self] partialResult, error in
                        guard let self = self else { return }
                        
                        if let error = error {
                            DispatchQueue.main.async {
                                self.responseText = "Inference Error: \(error.localizedDescription)"
                                self.isGenerating = false
                            }
                            return
                        }
                        
                        if let newText = partialResult {
                            DispatchQueue.main.async {
                                self.responseText += newText
                            }
                        }
                    }
                }
            }.value
            
            DispatchQueue.main.async {
                self.isGenerating = false
            }
        } catch {
            self.responseText = "Execution Error: \(error.localizedDescription)"
            self.isGenerating = false
        }
    }
    
    /// Registers background/foreground lifecycle notifications to prevent OS Jetsam memory kills.
    private func setupMemoryLifecycleObservers() {
        #if os(iOS)
        NotificationCenter.default.addObserver(
            self,
            selector: #selector(appDidEnterBackground),
            name: UIApplication.didEnterBackgroundNotification,
            object: nil
        )
        NotificationCenter.default.addObserver(
            self,
            selector: #selector(appWillEnterForeground),
            name: UIApplication.willEnterForegroundNotification,
            object: nil
        )
        #endif
    }
    
    @objc private func appDidEnterBackground() {
        print("[INFO] App entered background. Releasing LLM Metal VRAM to prevent OS Jetsam kill...")
        releaseMemory()
    }
    
    @objc private func appWillEnterForeground() {
        print("[INFO] App returned to foreground. Reloading LLM Metal Engine...")
        if !isInitialized {
            setupEngine()
        }
    }
    
    /// Instantly releases model reference to free system RAM/VRAM.
    public func releaseMemory() {
        self.llmInference = nil
        self.isInitialized = false
        print("[INFO] Model released from RAM/VRAM.")
    }
}