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 = ""
            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 }
            }
            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
            }
        }
        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)
        }
    }

    /// 把语音转写稿整理成健康日记草稿(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)
    }
}