// PolishingService.swift
// OSGKeyboard · Shared
//
// v0.3.0 rewrite: one-step "intelligent" polish that combines ASR
// error correction, filler removal, and tone adaptation in a single
// LLM call. The previous design was two separate steps (correction
// then polish) which doubled latency and token cost; Typeless,
// Wispr Flow, and the "intelligent" rewrite literature all confirm
// the merged prompt performs just as well for everyday Chinese /
// English dictation while halving the network round-trip.
//
// Engine matrix:
//   - `engineMode == "cloud"`  → user's cloud ASR + user's cloud LLM (independent)
//   - `engineMode == "local"`  → on-device ASR + user's LLM (or built-in DeepSeek)
//   - Ultra-short, structure-free utterances skip the LLM entirely
//   - Cloud without API key     → raw + `.missingAPIKey` warning
//   - Local without build key   → raw + `.missingAPIKey` warning
//
// Caller-supplied `PolishContext` carries the per-call signals:
//   - `appContext`     code / email / chat / document / unknown
//   - `intensity`      light / medium / heavy (per-call override)
//   - `precedingText`  optional tail of the cursor's preceding text
//     for reference resolution
//
// The prompt is intentionally a single message; multi-message
// conversation history would let earlier hallucinations pollute
// later calls (see MIT 2026 "Do LLMs Benefit From Their Own Words?")
// and the user expectation is that each take is independent.

import Foundation

public actor PolishingService {

    public enum PolishError: Error, Equatable {
        case noTranscript
        case timeout
        /// Local engine DeepSeek step: `PreconfiguredKeys.deepseek` is
        /// still the repo placeholder, or cloud engine Keychain is empty.
        case missingAPIKey
        /// The keychain was unreadable (device locked before first unlock)
        /// — the key likely EXISTS; treat as transient, never as "please
        /// re-enter your API key".
        case keychainLocked
    }

    /// v0.2.1: what the LLM should do with the raw transcript. The
    /// polish path stays the default so every existing call site keeps
    /// its current behaviour — translation is opt-in via the `translate`
    /// case and gets a target-locale parameter baked into the prompt.
    public enum PolishMode: Equatable, Sendable {
        case polish
        case translate(targetLocaleId: String)
    }

    private let store: any ConfigurationStore
    private let timeout: TimeInterval
    /// Optional injected client (mostly for testing). When nil we build
    /// one from `store.makeClient()` per call.
    private let injectedClient: LLMClient?

    /// `timeout` is the baseline (shortest) per-request HTTP timeout,
    /// used as the floor for `effectiveTimeout(for:)`. It defaults to the
    /// shared `LLMClient.requestTimeout`. The safety-net timer adds its
    /// own slack on top of the length-scaled budget in `polishRemote`, so
    /// no `+1` is baked in here.
    public init(
        store: any ConfigurationStore = AppGroupStore(),
        client: LLMClient? = nil,
        timeout: TimeInterval? = nil
    ) {
        self.store = store
        self.injectedClient = client
        self.timeout = timeout ?? LLMClientFactory.defaultRequestTimeout
    }

    /// v0.3.0: context-aware polish entry point. The optional
    /// `PolishContext` carries per-call signals (app context,
    /// intensity, preceding text). Translation is a separate concept
    /// (see `mode` below) so callers wanting the v0.2.1 translate
    /// flow should keep using the override prompt / providerId
    /// overloads exposed by the host.
    public func polish(
        _ raw: String,
        mode: PolishMode = .polish,
        systemPrompt: String? = nil,
        providerIdOverride: String? = nil,
        context: PolishContext? = nil
    ) async throws -> String {
        let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !trimmed.isEmpty else { throw PolishError.noTranscript }

        let resolvedContext = resolveContext(override: context)

        // Ultra-short, structure-free inputs skip the LLM to save
        // latency (e.g. "好", "OK", "明天见").
        if mode == .polish,
           systemPrompt == nil || systemPrompt?.isEmpty == true,
           TranscriptPostProcessor.shouldSkipLLM(for: trimmed) {
            return TranscriptPostProcessor.localClean(trimmed)
        }

        if injectedClient == nil {
            let providerId = Self.resolvedProviderId(store: store, providerIdOverride: providerIdOverride)
            let hasPolishKey = Self.hasPolishAPIKey(store: store, providerId: providerId)
            guard hasPolishKey else {
                if case .unavailable = Keychain.apiKeyOutcome(for: providerId, preferICloudSync: true) {
                    throw PolishError.keychainLocked
                }
                throw PolishError.missingAPIKey
            }
        }

        let llmResult = try await polishRemote(
            trimmed,
            mode: mode,
            systemPrompt: systemPrompt,
            providerIdOverride: providerIdOverride,
            context: resolvedContext
        )

        // Translation and custom prompts bypass the polish post-processor.
        if mode != .polish || (systemPrompt != nil && !(systemPrompt?.isEmpty ?? true)) {
            return llmResult
        }

        return TranscriptPostProcessor.process(original: trimmed, llmOutput: llmResult)
    }

    private func resolveContext(override: PolishContext?) -> PolishContext {
        guard let override else {
            return PolishContext(
                appContext: store.detectedAppContext?.context ?? .unknown,
                intensity: store.polishIntensity
            )
        }
        return override
    }

    private func polishRemote(
        _ trimmed: String,
        mode: PolishMode,
        systemPrompt: String? = nil,
        providerIdOverride: String? = nil,
        context: PolishContext
    ) async throws -> String {
        let effectiveProviderId = Self.resolvedProviderId(
            store: store,
            providerIdOverride: providerIdOverride
        )
        let client: LLMClient
        if let injectedClient {
            client = injectedClient
        } else {
            let preset = LLMProvider.provider(id: effectiveProviderId)
            let (baseURL, model) = Self.resolveLLMEndpoint(
                store: store,
                preset: preset,
                providerIdOverride: providerIdOverride
            )
            let apiKey: String
            if effectiveProviderId == "deepseek" {
                let userKey = store.apiKey.trimmingCharacters(in: .whitespacesAndNewlines)
                if !userKey.isEmpty {
                    apiKey = userKey
                } else if PreconfiguredKeys.isDeepseekConfigured {
                    apiKey = PreconfiguredKeys.deepseek
                } else {
                    throw PolishError.missingAPIKey
                }
            } else {
                apiKey = store.apiKey
            }
            client = LLMClientFactory.make(
                providerId: effectiveProviderId,
                baseURL: baseURL,
                apiKey: apiKey,
                model: model,
                thinkingEnabled: store.llmThinkingEnabled
            )
        }

        let prompt: String
        if let override = systemPrompt, !override.isEmpty {
            prompt = override
        } else {
            switch mode {
            case .polish:
                prompt = buildPrompt(
                    for: trimmed,
                    context: context,
                    providerId: effectiveProviderId
                )
            case .translate(let targetLocaleId):
                let target = TranslationLanguageCatalog.resolve(targetLocaleId)
                prompt = TranslationPrompt.make(
                    target: target,
                    providerId: effectiveProviderId,
                    appContext: context.appContext
                )
            }
        }
        let budget = effectiveTimeout(for: trimmed)
        // The HTTP request itself uses `budget`; the safety-net timer is
        // given a small slack on top so a clean URL timeout surfaces its
        // (more specific) transport error before the race fires.
        let safetyNet = budget + 2

        return try await withThrowingTaskGroup(of: String.self) { group in
            group.addTask {
                try await client.polish(trimmed, systemPrompt: prompt, timeout: budget)
            }
            group.addTask {
                try await Task.sleep(nanoseconds: UInt64(safetyNet * 1_000_000_000))
                throw PolishError.timeout
            }
            let result = try await group.next()!
            group.cancelAll()
            return result
        }
    }

    /// Shared output contract injected into every polish prompt.
    internal static func globalOutputContract(useChinese: Bool) -> String {
        if useChinese {
            return """
            ## 全局输出契约（所有润色档位均必须遵守，优先级最高）
            1. **禁止新增 emoji**：原文无 emoji 时输出不得出现 emoji；原文有 emoji 时仅可原样保留。
            2. **必须恢复合理标点**：逗号、句号、问号、感叹号；按语义分句，不要输出无标点长段。
            3. **结构服从当前风格**：
               - 保留原文明确表达的顺序、分点、步骤和层级，不得把独立事项揉成一段
               - 是否编号、分组或仅自然分段，由当前风格包的结构规则决定
               - 不得为了视觉整齐而给普通聊天、单一事项或连续叙述强加列表
            4. **数字要结合上下文判断**（重要）：
               - 有意义的数字（价格、日期、数量、时间、电话、版本号）→ 保持不变
               - 但语音里的序号常被误识别成数字或时间，需结合上下文修回并列表化：
                 · 已出现「第一点」，随后的「第2:00 / 第2点0 / 第二零零」多半是「第二点」，「第3:00」多半是「第三点」
                 · 「1、2、3」「一、二、三」在列举语境里就是序号，转成 `1. ` 列表
               - 判断依据是上下文里是否在“分点/列举”，不要机械地保留听错的数字
            5. **改写边界**：具体措辞和改写幅度服从当前风格与力度，但不得新增事实、改变立场或虚构上下文。
            6. **不改**人名、地名、专有名词（除非 ASR 明显错误）。
            7. 输出语言必须与原文一致；不翻译、不扩写成 AI 文案。
            8. 只输出最终文本：不要解释、不要引号包裹、不要前缀说明。
            """
        } else {
            return """
            ## Global output contract (mandatory at every intensity — highest priority)
            1. **No new emojis**: if the original has none, output must have none; preserve originals only.
            2. **Restore proper punctuation**: commas, periods, question marks; break run-on speech into sentences.
            3. **Structure follows the active style**:
               - Preserve explicit ordering, points, steps, and hierarchy; do not collapse independent items.
               - Let the active style decide whether to number, group, or use natural paragraphs.
               - Do not force lists onto ordinary chat, a single item, or continuous narrative.
            4. **Judge numbers by context** (important):
               - Meaningful numbers (prices, dates, quantities, times, phone numbers, versions) → keep unchanged.
               - But spoken ordinals are often misrecognized as digits/times; use context to restore and listify:
                 · after a "first point", a following "2:00 / point 2 / two oh oh" is likely "second point", "3:00" is "third point"
                 · "1, 2, 3" or "one, two, three" in an enumerating context are ordinals → convert to a `1. ` list
               - Decide by whether the context is enumerating; do not mechanically preserve a misheard number.
            5. **Rewrite boundary**: wording and rewrite depth follow the active style and intensity, but never add facts, change the user's position, or invent context.
            6. **Do not** alter person names, places, or proper nouns unless clearly misrecognized.
            7. Output language must match the input; do not translate or expand into marketing copy.
            8. Output the final text only: no explanation, no quotes, no preamble.
            """
        }
    }

    internal func buildPrompt(
        for text: String,
        context: PolishContext,
        providerId: String
    ) -> String {
        let dictionaryBlock = Self.mergedDictionaryBlock(
            dictionary: store.personalDictionary,
            supplement: context.dictionarySupplement
        )
        let useChinese = shouldUseChineseGuidance(providerId: providerId)
        let style = PolishStylePackCatalog.resolve(
            id: store.activePolishStyleId,
            userCatalog: store.polishStyleCatalog
        )
        return PolishPromptComposer.compose(
            text: text,
            style: style,
            context: context,
            dictionaryBlock: dictionaryBlock,
            globalContract: Self.globalOutputContract(useChinese: useChinese),
            useChineseGuidance: useChinese
        )
    }

    internal static func mergedDictionaryBlock(
        dictionary: PersonalDictionary,
        supplement: String?
    ) -> String {
        let base = dictionary.promptFragment()
        let extra = supplement?.trimmingCharacters(in: .whitespacesAndNewlines) ?? ""
        if base.isEmpty { return extra }
        if extra.isEmpty { return base }
        return base + "\n" + extra
    }

    private func shouldUseChineseGuidance(providerId: String) -> Bool {
        switch providerId {
        case "zhipu", "moonshot", "qwen", "deepseek", "ark", "minimax", "siliconflow", "mimo":
            return true
        default:
            return false
        }
    }

    /// Per-request HTTP timeout, scaled with transcript length. This is
    /// the *actual* value handed to `LLMClient.polish(timeout:)`, so long
    /// dictations (which generate long, listified, multi-paragraph output)
    /// are not cut off mid-generation by a fixed 15 s ceiling. Grows by
    /// ~10 s per 100 characters, capped at 120 s.
    ///
    /// Previously this value was computed but only used for the safety-net
    /// timer while the URLRequest stayed pinned at 15 s — the scaling was
    /// dead code and long transcripts timed out, falling back to the raw
    /// (unpolished, unsegmented) ASR text.
    internal func effectiveTimeout(for text: String) -> TimeInterval {
        let scaled = timeout + (Double(text.count) / 100.0) * 10.0
        // The cap participates in the keyboard-watchdog budget — see
        // `FlowSessionKeys.keyboardResultTimeout`. Raising it here without
        // going through that constant would silently break the invariant
        // "keyboard timeout > host worst case".
        return min(max(scaled, timeout), FlowSessionKeys.maxPolishTimeout)
    }

    internal static func resolvedProviderId(
        store: any ConfigurationStore,
        providerIdOverride: String?
    ) -> String {
        if let providerIdOverride {
            return providerIdOverride
        }
        let id = store.providerId
        // Local installs without a user LLM key keep using the built-in DeepSeek path.
        if store.engineMode == "local",
           id != "deepseek",
           store.apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty,
           PreconfiguredKeys.isDeepseekConfigured {
            return "deepseek"
        }
        return id
    }

    internal static func hasPolishAPIKey(store: any ConfigurationStore, providerId: String) -> Bool {
        if !store.apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty {
            return true
        }
        if providerId == "deepseek", PreconfiguredKeys.isDeepseekConfigured {
            return true
        }
        return false
    }

    internal static func resolveLLMEndpoint(
        store: any ConfigurationStore,
        preset: LLMProvider,
        providerIdOverride: String?
    ) -> (baseURL: String, model: String) {
        if providerIdOverride != nil {
            return (preset.defaultBaseURL, preset.defaultModel)
        }
        let baseURL = store.baseURL.isEmpty ? preset.defaultBaseURL : store.baseURL
        let model = store.model.isEmpty ? preset.defaultModel : store.model
        return (baseURL, model)
    }
}

extension PolishingService.PolishError: LocalizedError {
    public var errorDescription: String? {
        switch self {
        case .noTranscript:
            return "No transcript to polish."
        case .timeout:
            return "LLM polish timed out."
        case .missingAPIKey:
            return "Missing API key (cloud: Settings API key; local: build configuration)."
        case .keychainLocked:
            return "API key unavailable while the device is locked — will work after unlock."
        }
    }
}