import Foundation
import MLX

enum AIRuntimeError: Error, LocalizedError {
    case notReady
    case modelLoadFailed(String)
    case inferenceFailed(String)

    var errorDescription: String? {
        switch self {
        case .notReady:                  return String(appLoc: "AI 模型尚未准备好")
        case .modelLoadFailed(let m):    return String(appLoc: "模型加载失败:\(m)")
        case .inferenceFailed(let m):    return String(appLoc: "推理失败:\(m)")
        }
    }
}

/// 推理优先级。interactive = 用户正在屏幕前等(识别/问答/自检);
/// background = 预生成(报告摘要等),排队让行、解码中可被协作式抢占。
nonisolated enum InferencePriority: Sendable, Equatable {
    case interactive
    case background
}

actor AIRuntime {
    static let shared = AIRuntime()

    enum Status: Sendable, Equatable {
        case notReady
        case loading
        case ready
        case error(String)
    }

    private(set) var status: Status = .notReady
    private(set) var vlStatus: Status = .notReady

    /// 末次文本生成的性能统计(性能自检页消费;两后端归一)。
    private(set) var lastGenerateStats: GenerateStats?

    /// 当前实际生效的后端标签(性能自检 / PPT 截图用)。
    /// 端侧统一走 MLX/GPU(Gemma-3n E2B);MNN/SME2 LLM 路径已移除。
    var activeBackendLabel: String {
        #if targetEnvironment(simulator)
        return "MLX · CPU(模拟器)"
        #else
        return "MLX · GPU"
        #endif
    }

    private var llmSession: LLMSession?
    private var vlSession: VLSession?

    // MARK: - Gemini 云端后端(hybrid:端侧 Gemma 默认,云端按需增强)
    // 云端调用不占本机显存,不进 OOM 闸门,可与端侧推理并发。用于「云端深度解读 / 多语言」
    // 与「拍报告/药盒多模态读图」——后者恢复端侧 Gemma-3n(MLX 文本版)丢掉的真·视觉能力。
    private let gemini = GeminiBackend()
    /// 云端是否可用(用户已开启 + 有 key)。UI 与调用方据此决定走云还是端侧。
    nonisolated var cloudAvailable: Bool { CloudAI.isConfigured }
    /// 云端后端标签(性能自检 / 截图用)。
    nonisolated static var cloudLabel: String { "Gemini · \(CloudAI.model)" }

    // MARK: - 串行推理闸门(§3.1 OOM 防护的真正落地)
    //
    // actor 只串行化「方法入口」,但 generate() 同步返回流、真正解码在内部 Task;
    // analyzeReport 也在 await 期间让出 actor。若不加闸门,LLM 流正在解码时触发 VL,
    // 两个模型会同时在 GPU 上解码 → 冲过单 App 内存上限被 jetsam 杀
    //(MEMORY 记录的「in-flight 流并发窄口」)。
    //
    // 这里用 actor 内信号量(count = 1):所有「会占显存的重活」(解码 + 模型加载)
    // 进入前先 await acquireGate(),结束后 releaseGate()。actor 串行执行保证
    // gateBusy / gateWaiters 的读写天然无并发。
    private struct GateWaiter {
        let priority: InferencePriority
        let cont: CheckedContinuation<Void, Never>
    }
    private var gateBusy = false
    private var gateHolderPriority: InferencePriority = .interactive
    private var preemptRequested = false
    private var gateWaiters: [GateWaiter] = []

    /// interactive 排到所有 background 等待者之前;同优先级保持 FIFO。纯函数,单测覆盖。
    nonisolated static func gateInsertionIndex(of priority: InferencePriority,
                                               in waiting: [InferencePriority]) -> Int {
        guard priority == .interactive else { return waiting.count }
        return waiting.firstIndex(of: .background) ?? waiting.count
    }

    private func acquireGate(_ priority: InferencePriority = .interactive) async {
        if !gateBusy {
            gateBusy = true
            gateHolderPriority = priority
            return
        }
        // 前台请求撞上后台持有者:请其让位 —— 后台解码循环在下一个 token 抛 CancellationError。
        if priority == .interactive, gateHolderPriority == .background {
            preemptRequested = true
        }
        await withCheckedContinuation { (cont: CheckedContinuation<Void, Never>) in
            let idx = Self.gateInsertionIndex(of: priority, in: gateWaiters.map(\.priority))
            gateWaiters.insert(GateWaiter(priority: priority, cont: cont), at: idx)
        }
        // 被 releaseGate 唤醒时即已持有闸门(gateBusy 保持 true)。
    }

    private func releaseGate() {
        preemptRequested = false
        if gateWaiters.isEmpty {
            gateBusy = false
        } else {
            // 把闸门直接交给队首等待者,gateBusy 维持 true,不留空窗。
            let next = gateWaiters.removeFirst()
            gateHolderPriority = next.priority
            next.cont.resume()
        }
    }

    /// 后台持有者每收到一个 token 查一次:前台在排队就让位。
    private func shouldPreempt(_ priority: InferencePriority) -> Bool {
        priority == .background && preemptRequested
    }

    private init() {}

    /// App 启动时调用一次:给 MLX 的 GPU 缓冲池设上限,避免 reuse cache 在大模型常驻之上
    /// 继续膨胀、把峰值推过单 App 内存上限。仅真机生效(模拟器走 CPU,且部分 Metal 路径会 abort)。
    /// 与 increased-memory-limit entitlement + LLM/VL 互斥卸载配合,三管齐下防 jetsam OOM。
    nonisolated static func configureMLXMemory() {
        #if !targetEnvironment(simulator)
        // 256MB cache 上限:够复用、不至于在 3GB 模型之上再囤几百 MB 空闲缓冲。
        MLX.Memory.cacheLimit = 256 * 1024 * 1024
        #endif
    }

    /// 加载文本模型(端侧 MLX/GPU,Gemma-3n E2B)。首次调用会真正加载,后续幂等。
    func prepare() async throws {
        // 已有其他调用方在加载时,轮询等其结束再判定结果。
        // 不能像旧实现那样裸 return:那会让调用方误以为已 ready,随后 generate 的
        // `guard status == .ready` 失败 → 用户撞上「假错误屏」(模型其实正常加载中)。
        while status == .loading {
            try await Task.sleep(nanoseconds: 80_000_000)
        }
        if status == .ready { return }

        // 用 isComplete(逐文件字节校验)而非 isReady(只看 config.json):config.json 最小最先下完,
        // 半下载时 isReady 仍 true 会让加载在残缺 safetensors 上崩溃。与 ModelDownloadService 的
        // 完成判据保持一致(它也用 isComplete)。
        guard ModelStore.shared.isComplete(for: .llm) else {
            status = .error("LLM 模型未就绪")
            throw AIRuntimeError.notReady
        }

        // 进闸门:等所有在跑的推理(可能是 VL 解码)结束,再卸 VL + 载 LLM,
        // 避免「VL 解码 + LLM 加载」内存峰值叠加 OOM。
        await acquireGate()
        defer { releaseGate() }
        // 拿到闸门后复查:排队期间可能已被别的调用方加载好,避免重复 load。
        if status == .ready { return }

        // OOM 闸门(§3.1):LLM(~1GB)与 VL(~3GB)不可同时常驻,叠加会冲过单 App 内存上限被 jetsam 杀。
        unloadVL()

        status = .loading
        do {
            let session = try await LLMSession.load(
                folderURL: ModelStore.shared.localURL(for: .llm)
            )
            self.llmSession = session
            status = .ready
        } catch {
            status = .error("\(error)")
            throw AIRuntimeError.modelLoadFailed("\(error)")
        }
    }

    /// 流式生成。调用前应先 await prepare()。
    /// 注意:返回流是同步创建的,但跨 actor 调用 LLMSession 需要 await。
    /// priority = .background 时排队让行、解码中可被前台请求按 token 抢占(CancellationError 透传)。
    func generate(prompt: String,
                  maxTokens: Int = 256,
                  priority: InferencePriority = .interactive) -> AsyncThrowingStream<TokenChunk, Error> {
        // 在 actor 隔离上下文中捕获快照,Task 内不再访问 self.status / self.llmSession
        let snapshotStatus = status
        let snapshotSession = llmSession

        return AsyncThrowingStream { continuation in
            let task = Task {
                guard snapshotStatus == .ready, let session = snapshotSession else {
                    continuation.finish(throwing: AIRuntimeError.notReady)
                    return
                }
                // 进闸门:保证本次 LLM 解码与任何 VL 解码 / 模型加载串行,绝不并发占显存。
                await self.acquireGate(priority)
                // defer 保证正常结束 / 异常 / 取消都释放闸门;杜绝未来新增 early-return 导致全局推理死锁。
                defer { self.releaseGate() }
                do {
                    // session.generate 跨 actor 边界,需要 await
                    let stream = await session.generate(prompt: prompt, maxTokens: maxTokens)
                    for try await chunk in stream {
                        // 消费者(UI)提前关闭/取消时,下面的 checkCancellation 让本 Task 尽快退出,
                        // 连带丢弃 session 流并触发其 onTermination,停止底层 MLX 解码,不空耗 GPU。
                        try Task.checkCancellation()
                        // 后台任务让位:前台请求在排队时,下一个 token 处主动退出。
                        if self.shouldPreempt(priority) { throw CancellationError() }
                        continuation.yield(chunk)
                    }
                    self.lastGenerateStats = await session.lastStats
                    continuation.finish()
                } catch is CancellationError {
                    // 取消/抢占以 CancellationError 透传,调用方据此区分「让位」与「真失败」。
                    continuation.finish(throwing: CancellationError())
                } catch {
                    continuation.finish(throwing: AIRuntimeError.inferenceFailed("\(error)"))
                }
            }
            // 消费者取消/流终止时取消内部 Task(与 LLMSession / HealthExportService 一致)。
            continuation.onTermination = { _ in task.cancel() }
        }
    }

    // MARK: - Gemini 云端(hybrid 增强)

    /// 云端流式生成(深度解读 / 多语言)。不进 OOM 闸门,可与端侧并发。
    /// 调用前应确认 `cloudAvailable`;失败时调用方回退端侧 `generate`。
    func generateCloud(prompt: String, maxTokens: Int = 512) -> AsyncThrowingStream<TokenChunk, Error> {
        AsyncThrowingStream { continuation in
            let task = Task {
                do {
                    let stream = await self.gemini.generate(prompt: prompt, maxTokens: maxTokens)
                    for try await chunk in stream {
                        try Task.checkCancellation()
                        continuation.yield(chunk)
                    }
                    self.lastGenerateStats = await self.gemini.lastStats
                    continuation.finish()
                } catch is CancellationError {
                    continuation.finish(throwing: CancellationError())
                } catch {
                    continuation.finish(throwing: AIRuntimeError.inferenceFailed("\(error)"))
                }
            }
            continuation.onTermination = { _ in task.cancel() }
        }
    }

    /// 云端多模态识别:图片直传 Gemini 读出结构化文本(通常 JSON)。
    /// 恢复端侧丢掉的真·视觉读图;失败时调用方回退端侧 OCR+文本(§3.2)。
    func analyzeReportCloud(imageURLs: [URL], prompt: String, maxTokens: Int = 1024) async throws -> String {
        do {
            return try await gemini.analyze(imageURLs: imageURLs, prompt: prompt, maxTokens: maxTokens)
        } catch {
            throw AIRuntimeError.inferenceFailed("\(error)")
        }
    }

    // MARK: - 端侧 ASR(SenseVoice via sherpa-mnn)互斥入口

    /// 给端侧语音转写(SenseVoice,占 CPU + 内存)用:进推理闸门串行 + 卸掉常驻 LLM/VL 腾内存,
    /// 避免与文本/视觉模型同时常驻冲过单 App 内存上限被 jetsam 杀(§3.1 OOM 防护)。
    ///
    /// 问诊流程是「先转写 → 再 organize(会重载 LLM)」严格串行,所以这里卸掉 LLM 是安全的:
    /// body 跑完一次转写返回后,调用方走 DiaryAssistService.organizeConsultation 时会自然重新 prepare。
    /// body 内部应自行 hop 到后台线程做阻塞解码(本 actor 在 await 期间不阻塞,但闸门保持持有,
    /// 其它推理请求会排队等待,杜绝并发占内存)。
    func runExclusiveForASR<T: Sendable>(_ body: @Sendable () async throws -> T) async rethrows -> T {
        await acquireGate(.interactive)
        defer { releaseGate() }
        unloadLLM()
        unloadVL()
        return try await body()
    }

    // MARK: - VL

    /// 加载 VL 模型。幂等,首调真正 load。
    func prepareVL() async throws {
        while vlStatus == .loading {
            try await Task.sleep(nanoseconds: 80_000_000)
        }
        if vlStatus == .ready { return }

        // MLX VL 改用 .llm 的 Qwen3.5-2B 多模态(VLMModelFactory 走 qwen3_5 视觉路径),
        // 不再单独需要 Qwen3-VL-4B。用 isComplete 排除半下载,与下载服务判据一致。
        guard ModelStore.shared.isComplete(for: .llm) else {
            vlStatus = .error("VL 模型未就绪")
            throw AIRuntimeError.notReady
        }

        // 进闸门:等所有在跑的推理(可能是 LLM 文本流)结束,再卸 LLM + 载 VL。
        // —— 这正是「指标速记识别时 App 自动退出」的主因防护。
        await acquireGate()
        defer { releaseGate() }
        if vlStatus == .ready { return }

        // OOM 闸门(§3.1):加载 VL(~3GB)前先卸 LLM(~1GB),否则两者常驻叠加冲过内存上限被 jetsam 杀。
        unloadLLM()

        vlStatus = .loading
        do {
            let session = try await VLSession.load(
                folderURL: ModelStore.shared.localURL(for: .llm)
            )
            self.vlSession = session
            vlStatus = .ready
        } catch {
            vlStatus = .error("\(error)")
            throw AIRuntimeError.modelLoadFailed("\(error)")
        }
    }

    // MARK: - 卸载(OOM 闸门)

    /// 卸载 LLM,释放 ModelContainer 引用并清 MLX 显存缓存。幂等。
    /// 注:只在持有推理闸门时调用(prepareVL 内),此刻不会有 LLM 流在解码,卸载即时生效。
    private func unloadLLM() {
        guard llmSession != nil else { return }
        llmSession = nil
        status = .notReady
        MLX.Memory.clearCache()
    }

    /// 卸载 VL,释放 ModelContainer 引用并清 MLX 显存缓存。幂等。
    private func unloadVL() {
        guard vlSession != nil else { return }
        vlSession = nil
        vlStatus = .notReady
        MLX.Memory.clearCache()
    }

    /// 图像 → JSON 字符串(由 VLPrompts.reportExtraction 引导)。
    /// 调用方负责解析 + 失败回退(§3.2)。
    /// 推理闸门保证本调用与 LLM.generate() 的解码串行,不会同时占显存 OOM。
    func analyzeReport(imageURLs: [URL],
                       prompt: String,
                       maxTokens: Int = 512) async throws -> String {
        guard vlStatus == .ready, let session = vlSession else {
            throw AIRuntimeError.notReady
        }
        await acquireGate()
        defer { releaseGate() }
        do {
            return try await session.analyze(
                imageURLs: imageURLs,
                prompt: prompt,
                maxTokens: maxTokens
            )
        } catch {
            throw AIRuntimeError.inferenceFailed("\(error)")
        }
    }
}