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

        // 1. 去 <think>...</think>(复用 HealthExportService 的兜底)
        let stripped = HealthExportService.stripThinkBlocks(collected)
        // 2. 抠出第一段平衡 JSON(复用 CaptureService.extractJSONObject)+ 弱模型畸形修复
        let jsonStr = CaptureService.repairJSON(CaptureService.extractJSONObject(from: stripped))
        guard let data = jsonStr.data(using: .utf8),
              let obj = try? JSONSerialization.jsonObject(with: data, options: [.fragmentsAllowed]),
              let dict = obj as? [String: Any] else {
            throw AssistError.parseFailed("非 JSON 输出")
        }
        guard let rawQuestions = dict["questions"] as? [[String: Any]] else {
            throw AssistError.parseFailed("缺少 questions 字段")
        }
        let questions = 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)
        }
        guard !questions.isEmpty else { throw AssistError.empty }
        return (Array(questions.prefix(4)), lastRate)
    }
}