OSGKeyboard/OSGKeyboardShared/Models/AIReplyVariant.swift

// AIReplyVariant.swift
// OSGKeyboard · Shared
//
// Typed reply choices produced by one clipboard-reply request. Model output
// never controls SF Symbols or display labels; only these local allowlists do.

import Foundation

public struct AIReplyVariant: Equatable, Identifiable, Sendable {
    public enum Kind: String, CaseIterable, Sendable {
        case ordinary
        case formal
        case playful

        public var systemImage: String {
            switch self {
            case .ordinary:
                return "bubble.left.fill"
            case .formal:
                return "briefcase.fill"
            case .playful:
                return "theatermasks.fill"
            }
        }

        public var titleKey: String {
            switch self {
            case .ordinary:
                return "keyboard.ai.replyVariant.ordinary"
            case .formal:
                return "keyboard.ai.replyVariant.formal"
            case .playful:
                return "keyboard.ai.replyVariant.playful"
            }
        }
    }

    /// A bounded semantic hint for future feedback and presentation work.
    /// Unknown model values are intentionally normalized to `.neutral`.
    public enum Emotion: String, CaseIterable, Sendable {
        case neutral
        case warm
        case celebratory
        case empathetic
        case encouraging
        case grateful
        case apologetic
        case reassuring
        case playful
        case enthusiastic
        case calm

        /// The model selects only this semantic enum. SF Symbol names stay
        /// local and validated so malformed model output cannot control UI.
        public func systemImage(fallback kind: Kind) -> String {
            switch self {
            case .neutral:
                return kind.systemImage
            case .warm, .grateful:
                return "heart.fill"
            case .celebratory:
                return "party.popper.fill"
            case .empathetic:
                return "heart.text.square.fill"
            case .encouraging:
                return "hand.thumbsup.fill"
            case .apologetic:
                return "exclamationmark.bubble.fill"
            case .reassuring:
                return "checkmark.shield.fill"
            case .playful:
                return "theatermasks.fill"
            case .enthusiastic:
                return "sparkles"
            case .calm:
                return "leaf.fill"
            }
        }
    }

    public let id: UUID
    public let kind: Kind
    public let emotion: Emotion
    public let text: String

    public init(
        id: UUID = UUID(),
        kind: Kind,
        emotion: Emotion,
        text: String
    ) {
        self.id = id
        self.kind = kind
        self.emotion = emotion
        self.text = text
    }
}

public enum AIReplyVariantParsingResult: Equatable, Sendable {
    case variants([AIReplyVariant])
    case single(AIReplyVariant)
}

public enum AIReplyVariantParser {
    /// The only accepted multi-reply wire shape is:
    /// `{"variants":[{"kind":"ordinary","emotion":"warm","text":"…"}, ...]}`
    /// with exactly one item of each kind and no additional JSON fields.
    public static func parse(
        _ raw: String,
        sourceText: String? = nil
    ) -> [AIReplyVariant]? {
        guard let data = raw.data(using: .utf8),
              let root = try? JSONSerialization.jsonObject(with: data),
              let object = root as? [String: Any],
              Set(object.keys) == ["variants"],
              let items = object["variants"] as? [[String: Any]],
              items.count == AIReplyVariant.Kind.allCases.count else {
            return nil
        }

        var variantsByKind: [AIReplyVariant.Kind: AIReplyVariant] = [:]
        for item in items {
            guard Set(item.keys) == ["kind", "emotion", "text"],
                  let rawKind = item["kind"] as? String,
                  let kind = AIReplyVariant.Kind(rawValue: rawKind),
                  variantsByKind[kind] == nil,
                  let rawEmotion = item["emotion"] as? String,
                  let rawText = item["text"] as? String else {
                return nil
            }
            let text = rawText.trimmingCharacters(in: .whitespacesAndNewlines)
            guard !text.isEmpty,
                  !isSourceEcho(text, sourceText: sourceText) else {
                return nil
            }
            let emotion = AIReplyVariant.Emotion(rawValue: rawEmotion) ?? .neutral
            variantsByKind[kind] = AIReplyVariant(
                kind: kind,
                emotion: emotion,
                text: text
            )
        }

        let ordered = AIReplyVariant.Kind.allCases.compactMap { variantsByKind[$0] }
        return ordered.count == AIReplyVariant.Kind.allCases.count ? ordered : nil
    }

    /// Strict multi-reply parsing with a conservative single ordinary fallback.
    /// Fenced or malformed JSON is never surfaced verbatim to the insertion UI.
    public static func parseOrFallback(
        _ raw: String,
        sourceText: String? = nil
    ) -> AIReplyVariantParsingResult? {
        if let variants = parse(raw, sourceText: sourceText) {
            return .variants(variants)
        }
        guard let text = fallbackText(from: raw),
              !isSourceEcho(text, sourceText: sourceText) else {
            return nil
        }
        return .single(
            AIReplyVariant(kind: .ordinary, emotion: .neutral, text: text)
        )
    }

    public static func fallbackText(from raw: String) -> String? {
        let unfenced = removingCodeFence(from: raw)
        guard !unfenced.isEmpty else { return nil }

        if let data = unfenced.data(using: .utf8),
           let json = try? JSONSerialization.jsonObject(with: data),
           let extracted = replyText(in: json) {
            return normalized(extracted)
        }
        if let extracted = textValueFromMalformedJSON(unfenced) {
            return normalized(extracted)
        }

        guard !looksLikeJSON(unfenced) else { return nil }
        return normalized(unfenced)
    }

    private static let preferredTextKeys = [
        "text", "reply", "response", "content", "message", "answer"
    ]

    private static let metadataValues = Set(
        AIReplyVariant.Kind.allCases.map(\.rawValue)
            + AIReplyVariant.Emotion.allCases.map(\.rawValue)
            + ["variants", "kind", "emotion", "text"]
    )

    /// Rejects responses that merely report the clipboard message back to its
    /// sender. A bounded LCS catches close paraphrases while still allowing a
    /// reply to mention a necessary name or short keyword.
    private static func isSourceEcho(
        _ candidate: String,
        sourceText: String?
    ) -> Bool {
        guard let sourceText else { return false }
        let source = normalizedComparisonCharacters(sourceText)
        let reply = normalizedComparisonCharacters(candidate)
        guard source.count >= 8, reply.count >= 8 else { return false }

        let overlap = longestCommonSubsequenceLength(source, reply)
        guard overlap >= 8 else { return false }
        return Double(overlap) / Double(source.count) >= 0.68
            && Double(overlap) / Double(reply.count) >= 0.45
    }

    private static func normalizedComparisonCharacters(
        _ text: String
    ) -> [Character] {
        Array(
            text.lowercased().filter { character in
                character.unicodeScalars.contains {
                    CharacterSet.alphanumerics.contains($0)
                }
            }.prefix(1_000)
        )
    }

    private static func longestCommonSubsequenceLength(
        _ lhs: [Character],
        _ rhs: [Character]
    ) -> Int {
        let shorter: [Character]
        let longer: [Character]
        if lhs.count <= rhs.count {
            shorter = lhs
            longer = rhs
        } else {
            shorter = rhs
            longer = lhs
        }
        var previous = Array(repeating: 0, count: shorter.count + 1)
        var current = previous
        for longCharacter in longer {
            for index in shorter.indices {
                if longCharacter == shorter[index] {
                    current[index + 1] = previous[index] + 1
                } else {
                    current[index + 1] = max(
                        previous[index + 1],
                        current[index]
                    )
                }
            }
            swap(&previous, &current)
            current = Array(repeating: 0, count: shorter.count + 1)
        }
        return previous[shorter.count]
    }

    private static func replyText(in value: Any) -> String? {
        if let object = value as? [String: Any] {
            for key in preferredTextKeys {
                if let text = object[key] as? String,
                   let normalized = normalized(text) {
                    return normalized
                }
            }
            for nested in object.values {
                if let text = replyText(in: nested) {
                    return text
                }
            }
        }
        if let array = value as? [Any] {
            for nested in array {
                if let text = replyText(in: nested) {
                    return text
                }
            }
        }
        if let text = value as? String,
           !metadataValues.contains(text),
           let normalized = normalized(text) {
            return normalized
        }
        return nil
    }

    private static func removingCodeFence(from raw: String) -> String {
        var lines = raw
            .trimmingCharacters(in: .whitespacesAndNewlines)
            .components(separatedBy: .newlines)
        if lines.first?.trimmingCharacters(in: .whitespacesAndNewlines)
            .hasPrefix("```") == true {
            lines.removeFirst()
        }
        if lines.last?.trimmingCharacters(in: .whitespacesAndNewlines) == "```" {
            lines.removeLast()
        }
        return lines
            .joined(separator: "\n")
            .trimmingCharacters(in: .whitespacesAndNewlines)
    }

    private static func textValueFromMalformedJSON(_ raw: String) -> String? {
        let keyPattern = #""(?:text|reply|response|content|message|answer)"\s*:\s*("(?:\\.|[^"\\])*")"#
        if let regex = try? NSRegularExpression(pattern: keyPattern),
           let match = regex.firstMatch(
               in: raw,
               range: NSRange(raw.startIndex..., in: raw)
           ),
           let range = Range(match.range(at: 1), in: raw) {
            return decodeJSONString(String(raw[range]))
        }

        let stringPattern = #""(?:\\.|[^"\\])*""#
        guard let regex = try? NSRegularExpression(pattern: stringPattern) else {
            return nil
        }
        let candidates = regex.matches(
            in: raw,
            range: NSRange(raw.startIndex..., in: raw)
        ).compactMap { match -> String? in
            guard let range = Range(match.range, in: raw),
                  let decoded = decodeJSONString(String(raw[range])),
                  !metadataValues.contains(decoded) else {
                return nil
            }
            return normalized(decoded)
        }
        return candidates.max { $0.count < $1.count }
    }

    private static func decodeJSONString(_ quoted: String) -> String? {
        guard let data = quoted.data(using: .utf8) else { return nil }
        return try? JSONDecoder().decode(String.self, from: data)
    }

    private static func looksLikeJSON(_ value: String) -> Bool {
        let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines)
        return trimmed.hasPrefix("{") || trimmed.hasPrefix("[")
    }

    private static func normalized(_ value: String) -> String? {
        let text = value.trimmingCharacters(in: .whitespacesAndNewlines)
        return text.isEmpty ? nil : text
    }
}