kangkang/康康/Services/DiaryChatService.swift

import Foundation
import SwiftData

/// 「聊着记」功能的唯一 AI 门面(红线 #3:UI 不直接调 AIRuntime,一律经 Service)。
///
/// 场景:全屏动画形象与用户语音多轮对话——康康基于已有健康记录追问,聊完把整段对话蒸馏成日记草稿。
/// 对话上下文(健康数据快照)进场调一次 `makeContextJSON` 生成,整场复用;每轮回应流式吐字。
///
/// 云优先策略(出处:CaptureService.runVL 的「云优先回退端侧」模式):
/// - 云端可用(Gemini)时先走 `generateCloud`(不占显存、不进 OOM 闸门);
/// - 云端离线 / 超时 / 额度耗尽 / 输出为空,一律静默回退端侧 Gemma(MLX/GPU)。
///
/// 失败回退语义(§3.2 / 红线 #5「不让用户卡在 AI 错误屏」):
/// - `openingLine` / `reply`(流式):
///     · 云端【首个有效 token 前】失败 → 静默吞掉,落端侧;
///     · 云端【已产出内容后】失败 → `finish(throwing:)`,View 对「已有部分文本」的流错误按说完处理;
///     · 端侧 `prepare()` 失败 → `ChatError.modelNotReady`;
///     · 两路都走完但全空 → `ChatError.empty`;
///     · 取消(`CancellationError`)一律透传,不吞、不回退。
/// - `distill`(一次性 await):云端失败 / 清理后为空 → 落端侧;端侧 `prepare()` 失败 → `.modelNotReady`;
///     最终为空 → `.empty`。
/// - 纯函数 `roundsUsed` / `isWrappedUp` / `fallbackDraft`:不碰 AI,永不失败。
@MainActor
struct DiaryChatService {
    static let shared = DiaryChatService()
    private init() {}

    enum ChatError: Error, LocalizedError {
        case modelNotReady   // 端侧模型未就绪且云端不可用 / 失败
        case empty           // 两路都生成了但内容为空

        var errorDescription: String? {
            switch self {
            case .modelNotReady: return String(appLoc: "AI 模型尚未准备好")
            case .empty:         return String(appLoc: "AI 没有给出建议,请稍后重试")
            }
        }
    }

    // MARK: - 上下文快照

    /// 生成整场对话复用的健康数据 JSON(profile + 指标 + 日记 + 用药 …)。
    /// 只读 `ctx` 做快照,不写 SwiftData;同步、永不失败。
    /// 调用时机:进场调用一次,把结果缓存在 View 层,后续每轮 `reply` 传同一份,不重复检索。
    func makeContextJSON(in ctx: ModelContext) -> String {
        HealthExportService.serializeData(
            snapshot: HealthExportService.retrieveDialogueSnapshot(ctx: ctx)
        )
    }

    // MARK: - 流式:开场白 / 每轮回应

    /// 开场白:一句暖场引导 + 首个追问。160 token——开场较短,一句话足够。
    func openingLine(dataJSON: String) -> AsyncThrowingStream<TokenChunk, Error> {
        streamWithFallback(prompt: DiaryChatPrompts.opening(dataJSON: dataJSON), maxTokens: 160)
    }

    /// 每轮回应:共情用户最新一句 + 再追问一步(或收尾)。220 token——含共情 + 追问,比开场略长。
    /// - roundsUsed: 已发生的追问轮数(取自 `roundsUsed(in:)`)。
    ///   `roundsLeft = 本轮生成后剩余可问轮数`,为 0 时 prompt 内部自动切收尾指令(不再追问)。
    func reply(transcript: String,
               latest: String,
               dataJSON: String,
               roundsUsed: Int) -> AsyncThrowingStream<TokenChunk, Error> {
        let roundsLeft = max(0, DiaryChatPrompts.maxRounds - roundsUsed - 1)
        let prompt = DiaryChatPrompts.reply(
            transcript: transcript,
            latest: latest,
            dataJSON: dataJSON,
            roundsLeft: roundsLeft
        )
        return streamWithFallback(prompt: prompt, maxTokens: 220)
    }

    // MARK: - 蒸馏:聊完 → 日记草稿

    /// 把整段对话蒸馏成一篇健康日记草稿。一次性 await(非流式)。
    /// 云优先:`cloudAvailable` 时先试云端并收集全文,任何错误 / 清理后为空都静默回退端侧;
    /// 端侧走到时若 `prepare()` 失败,才抛 `.modelNotReady`(即云端也失败 / 不可用)。
    /// 400 token:对齐 `DiaryAssistService.organize`(:193)的产物量级——一篇日记草稿。
    func distill(turns: [HealthExportDialogueTurn]) async throws -> (text: String, decodeRate: Double) {
        let transcript = HealthExportDialogueTurn.transcript(from: turns)
        let prompt = DiaryChatPrompts.distill(transcript: transcript)

        // —— 云端优先(CaptureService.runVL 模式:失败静默回退,绝不卡死) ——
        if AIRuntime.shared.cloudAvailable {
            do {
                var collected = ""
                var rate: Double = 0
                let stream = await AIRuntime.shared.generateCloud(prompt: prompt, maxTokens: 400)
                for try await chunk in stream {
                    collected += chunk.text
                    if chunk.decodeRate > 0 { rate = chunk.decodeRate }
                }
                let text = HealthExportService.stripThinkBlocks(collected)
                    .trimmingCharacters(in: .whitespacesAndNewlines)
                if !text.isEmpty { return (text, rate) }
                // 云端走完但清理后为空 → 落端侧(不抛)。
            } catch {
                // 云端任何失败(离线 / 超时 / 额度)静默吞掉,落端侧(§3.2 失败回退红线)。
            }
        }

        // —— 端侧回退 ——
        do {
            try await AIRuntime.shared.prepare()
        } catch {
            throw ChatError.modelNotReady
        }
        var collected = ""
        var lastRate: Double = 0
        let stream = await AIRuntime.shared.generate(prompt: prompt, maxTokens: 400)
        for try await chunk in stream {
            collected += chunk.text
            if chunk.decodeRate > 0 { lastRate = chunk.decodeRate }
        }
        let text = HealthExportService.stripThinkBlocks(collected)
            .trimmingCharacters(in: .whitespacesAndNewlines)
        guard !text.isEmpty else { throw ChatError.empty }
        return (text, lastRate)
    }

    // MARK: - 核心:云优先流式,失败前段回退端侧

    /// 流式生成 + 云→端回退。返回同步创建的流,内部起 Task 逐 token 吐(去 `<think>` 后的 delta)。
    ///
    /// 「首个有效 token」的界定:经 `ThinkStripper.feed` 清理后【首个非空 delta 被 yield】的那一刻。
    /// 在此之前云端出任何错 → 静默回退端侧;在此之后出错 → `finish(throwing:)`(View 已有部分文本,按说完处理)。
    /// 风格对齐 HealthExportService.answer 的既有流式写法。
    private func streamWithFallback(prompt: String,
                                    maxTokens: Int) -> AsyncThrowingStream<TokenChunk, Error> {
        AsyncThrowingStream { continuation in
            let task = Task { @MainActor in
                // —— 云端优先(CaptureService.runVL 模式) ——
                var cloudYielded = false   // 是否已 yield 过非空内容 = 是否已产出「首个有效 token」
                if AIRuntime.shared.cloudAvailable {
                    do {
                        var stripper = ThinkStripper()
                        let stream = await AIRuntime.shared.generateCloud(prompt: prompt, maxTokens: maxTokens)
                        for try await chunk in stream {
                            try Task.checkCancellation()
                            let delta = stripper.feed(chunk.text)
                            if !delta.isEmpty {
                                cloudYielded = true
                                continuation.yield(TokenChunk(text: delta, decodeRate: chunk.decodeRate))
                            }
                        }
                        if cloudYielded {
                            continuation.finish()   // 云端有内容,正常收尾,不再走端侧。
                            return
                        }
                        // 云端走完但全空 → 落端侧(不 return)。
                    } catch is CancellationError {
                        continuation.finish(throwing: CancellationError())   // 取消透传,不吞、不回退。
                        return
                    } catch {
                        if cloudYielded {
                            // 已产出内容后失败:交给 View(有部分文本即按说完处理)。
                            continuation.finish(throwing: error)
                            return
                        }
                        // 首个有效 token 前失败:静默吞掉,落端侧(CaptureService.runVL 模式)。
                    }
                }

                // —— 端侧回退(MLX/GPU Gemma) ——
                do {
                    try await AIRuntime.shared.prepare()   // OOM 闸门内卸载互斥模型
                } catch {
                    continuation.finish(throwing: ChatError.modelNotReady)
                    return
                }
                do {
                    var localYielded = false
                    var stripper = ThinkStripper()
                    let stream = await AIRuntime.shared.generate(prompt: prompt, maxTokens: maxTokens)
                    for try await chunk in stream {
                        try Task.checkCancellation()
                        let delta = stripper.feed(chunk.text)
                        if !delta.isEmpty {
                            localYielded = true
                            continuation.yield(TokenChunk(text: delta, decodeRate: chunk.decodeRate))
                        }
                    }
                    if localYielded {
                        continuation.finish()
                    } else {
                        continuation.finish(throwing: ChatError.empty)   // 端侧走完全空。
                    }
                } catch is CancellationError {
                    continuation.finish(throwing: CancellationError())   // 取消透传。
                } catch {
                    // 端侧解码中途失败:透传错误(View 已有部分文本时按说完处理)。
                    continuation.finish(throwing: error)
                }
            }
            // 消费者(UI)关闭 / 取消流时取消内部 Task,停止底层解码,不空耗算力。
            continuation.onTermination = { _ in task.cancel() }
        }
    }

    // MARK: - 纯函数(无 AI,无共享状态;nonisolated 便于非隔离单测精确断言)

    /// 已发生的「追问轮数」= assistant 轮数 - 1(开场白不算一轮追问)。无 assistant 轮时为 0。
    nonisolated static func roundsUsed(in turns: [HealthExportDialogueTurn]) -> Int {
        let assistantCount = turns.filter { $0.role == .assistant }.count
        return max(0, assistantCount - 1)
    }

    /// 是否已聊够(达到 `DiaryChatPrompts.maxRounds` 轮追问,该收尾了)。
    nonisolated static func isWrappedUp(turns: [HealthExportDialogueTurn]) -> Bool {
        roundsUsed(in: turns) >= DiaryChatPrompts.maxRounds
    }

    /// 蒸馏失败时的兜底草稿:只取用户原话拼接,康康(assistant)的话绝不出现,数值原样不改写。
    /// - 0 条有效用户话 → `""`;
    /// - 恰 1 条 → 该原话本身(无前缀);
    /// - ≥2 条 → 每条前缀 `"· "`,按 `"\n"` join。
    nonisolated static func fallbackDraft(turns: [HealthExportDialogueTurn]) -> String {
        let userLines = turns
            .filter { $0.role == .user }
            .map { $0.text.trimmingCharacters(in: .whitespacesAndNewlines) }
            .filter { !$0.isEmpty }
        switch userLines.count {
        case 0:  return ""
        case 1:  return userLines[0]
        default: return userLines.map { "· " + $0 }.joined(separator: "\n")
        }
    }
}