kangkang/康康/Services/DiaryAssistService.swift

import Foundation

/// 「健康记录」AI 辅助:让 LLM 从医生角度提 3-4 个追问问题。
///
/// 设计上和 HealthExportService 同款门面,但输出量小(< 400 token),
/// 不流式 —— 直接 await 收完整结果再解析。
///
/// 调用方:DiaryQuickSheet。
@MainActor
struct DiaryAssistService {
    static let shared = DiaryAssistService()
    private init() {}

    /// 单条追问。fill 是带方括号占位符的模板,采纳时追加到原文末尾。
    /// `dim` 是问诊维度(取自 `DiaryAssistPrompts.dimensions`),用于跨轮按维度去重。
    /// `adopted` 由 UI 标记;`round` 由 UI 在 append 前打戳,用于多轮分组显示。
    struct Question: Identifiable, Hashable {
        let id: UUID
        let q: String
        let fill: String
        let dim: String
        var adopted: Bool
        var round: Int

        init(id: UUID = UUID(),
             q: String,
             fill: String,
             dim: String = "",
             adopted: Bool = false,
             round: Int = 0) {
            self.id = id
            self.q = q
            self.fill = fill
            self.dim = dim
            self.adopted = adopted
            self.round = round
        }
    }

    enum AssistError: Error, LocalizedError {
        case modelNotReady
        case empty
        case parseFailed(String)

        var errorDescription: String? {
            switch self {
            case .modelNotReady:        return String(appLoc: "AI 模型尚未准备好")
            case .empty:                return String(appLoc: "AI 没有给出建议,请稍后重试")
            case .parseFailed(let m):   return String(appLoc: "结果解析失败:\(m)")
            }
        }
    }

    /// 返回 3-4 条追问。
    /// - coveredDimensions: 多轮场景下,把之前各轮已覆盖的维度名(取自 question.dim)传进来,
    ///   prompt 会明确要求本轮避开这些维度。第一轮传空数组。
    /// 注意:本方法在 AIRuntime 的 actor 队列里串行排队,与 Capture / Export 互不抢占 GPU。
    func suggest(content: String,
                 coveredDimensions: [String] = []) async throws -> (questions: [Question], decodeRate: Double) {
        do {
            try await AIRuntime.shared.prepare()
        } catch {
            throw AssistError.modelNotReady
        }

        let prompt = DiaryAssistPrompts.suggest(content: content, coveredDimensions: coveredDimensions)

        // 低温采样下 MNN 仍偶发吐非 JSON / 漏掉外层 {"questions":…} 包裹(换 MNN 后比 MLX 更常见)。
        // 首次解析不出就自动重试一次,两次都失败才报错 —— 守 §10.5「失败回退,不让用户卡在 AI 错误屏」。
        var lastRate: Double = 0
        var parsedButEmpty = false
        var lastRaw = ""
        for _ in 0..<2 {
            try Task.checkCancellation()
            var collected = ""
            // 512(原 400):3-4 条含 dim/q/fill 的中文问题 JSON 约 250-320 token,400 在
            // 模型偶尔加前导/换行时易把尾部那条问题截断 → 整体解析失败。留足余量更稳。
            let stream = await AIRuntime.shared.generate(prompt: prompt, maxTokens: 512)
            for try await chunk in stream {
                collected += chunk.text
                if chunk.decodeRate > 0 { lastRate = chunk.decodeRate }
            }
            lastRaw = collected
            if let questions = Self.parseQuestions(from: collected) {
                if !questions.isEmpty {
                    return (Array(questions.prefix(4)), lastRate)
                }
                parsedButEmpty = true   // JSON 合法但没解析出问题:重试一次,仍空就当 .empty
            }
        }
        // 真机若仍偶发,这条日志能抓到模型当时的原始输出,便于定位是截断还是格式漂移。
        #if DEBUG
        print("[DiaryAssistService] 解析失败,原始输出 = \(lastRaw)")
        #endif
        throw parsedButEmpty ? AssistError.empty : AssistError.parseFailed("非 JSON 输出")
    }

    /// 从模型原始输出解析追问数组。容错链(对齐 §3.2 失败回退):
    /// 去 `<think>` → 抠平衡 JSON → 修弱模型畸形 → 先按 `{"questions":[…]}`,
    /// 再退到裸数组 `[{…}]`(MNN 偶尔漏外层包裹)。彻底解析不出返回 nil(调用方据此重试/报错)。
    /// 解析成功但无可用问题返回 `[]`(与 nil 区分:调用方报 .empty 而非 .parseFailed)。
    static func parseQuestions(from raw: String) -> [Question]? {
        let stripped = HealthExportService.stripThinkBlocks(raw)

        var rawQuestions: [[String: Any]]?
        // ① 标准结构 {"questions":[…]}
        let objStr = CaptureService.repairJSON(CaptureService.extractJSONObject(from: stripped))
        if let data = objStr.data(using: .utf8),
           let dict = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any],
           let arr = dict["questions"] as? [[String: Any]] {
            rawQuestions = arr
        }
        // ② 退路:模型漏了外层包裹,直接吐 [{…},{…}]
        if rawQuestions == nil {
            let arrStr = CaptureService.repairJSON(CaptureService.extractBalancedJSON(from: stripped))
            if let data = arrStr.data(using: .utf8),
               let arr = (try? JSONSerialization.jsonObject(with: data)) as? [[String: Any]] {
                rawQuestions = arr
            }
        }
        // ③ 终极兜底:外层包裹畸形 / 尾部被 maxTokens 截断时,①② 的整体解析都会失败。
        //    逐个抠出文本里所有平衡的 {…},凡含 "q" 字段的就当一条问题救回 —— 即便最后一条
        //    被截断,前面已闭合的几条仍能用。守 §3.2 / 红线 #5「失败回退,不让用户卡在 AI 错误屏」。
        if rawQuestions == nil {
            let salvaged = salvageQuestionObjects(from: stripped)
            if !salvaged.isEmpty { rawQuestions = salvaged }
        }
        guard let rawQuestions else { return nil }

        return rawQuestions.compactMap { d -> Question? in
            guard let q = (d["q"] as? String)?
                    .trimmingCharacters(in: .whitespacesAndNewlines), !q.isEmpty else {
                return nil
            }
            let fill = (d["fill"] as? String)?
                .trimmingCharacters(in: .whitespacesAndNewlines) ?? ""
            let dim = (d["dim"] as? String)?
                .trimmingCharacters(in: .whitespacesAndNewlines) ?? ""
            return Question(q: q, fill: fill, dim: dim)
        }
    }

    /// 逐对象救回:扫描原始文本里【所有】平衡的 `{…}` 子串(含嵌套内层),逐个尝试解析,
    /// 凡是含 `"q"` 字段的对象就当作一条问题收集。用于 ①② 整体解析失败的兜底:
    /// - 外层 `{"questions":[…]}` 的 `}` 被截断 → 但内部前几条 question 对象已闭合,照样救回;
    /// - 模型把对象拆在解释文字里 → 也能逐个抠出。
    /// 外层 wrapper 自身不含 `"q"`(只有 `"questions"`),天然被过滤掉,不会误收。
    /// 按 q 文本去重(保序),避免嵌套时同一对象被父子两层重复收集。
    private static func salvageQuestionObjects(from raw: String) -> [[String: Any]] {
        var openStack: [String.Index] = []
        var collected: [[String: Any]] = []
        var seenQ = Set<String>()
        var inString = false
        var escape = false
        var idx = raw.startIndex
        while idx < raw.endIndex {
            let ch = raw[idx]
            if escape { escape = false }
            else if ch == "\\" { escape = true }
            else if ch == "\"" { inString.toggle() }
            else if !inString {
                if ch == "{" {
                    openStack.append(idx)
                } else if ch == "}", let open = openStack.popLast() {
                    let sub = CaptureService.repairJSON(String(raw[open...idx]))
                    if let data = sub.data(using: .utf8),
                       let dict = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any],
                       let q = dict["q"] as? String,
                       !q.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty,
                       seenQ.insert(q).inserted {
                        collected.append(dict)
                    }
                }
            }
            idx = raw.index(after: idx)
        }
        return collected
    }

    /// 把语音转写稿整理成健康日记草稿(spec 2026-06-10-voice-diary)。
    /// 失败(模型未就绪 / 输出为空)抛错,调用方回退为直接使用原话,不卡死。
    /// 与 suggest 同样走 AIRuntime actor 队列,自然与追问/拍照串行。
    func organize(transcript: String) async throws -> (text: String, decodeRate: Double) {
        do {
            try await AIRuntime.shared.prepare()
        } catch {
            throw AssistError.modelNotReady
        }

        let prompt = DiaryAssistPrompts.organize(transcript: transcript)
        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 AssistError.empty }
        return (text, lastRate)
    }
}