app_for_reddit/App AI for Reddit/ViewModels/CommentWriterViewModel.swift

import AppKit
import SwiftUI

@MainActor
@Observable
final class CommentWriterViewModel {
    var draft = CommentDraft()
    var selectedTab: CommentWriterTab = .compose
    var isGenerating = false
    var errorMessage: String?
    var successMessage: String?
    var variants: [CommentVariant] = []
    var selectedVariantID: UUID?
    var showOverwriteConfirmation = false
    /// True after a successful Generate or History restore. Cleared on Reset.
    /// Controls whether the comment result field is shown in Compose.
    private(set) var hasGeneratedContent = false

    private var hasPremiumAccess = false
    private var hasEverPurchasedPremium = false
    private let injectedGenerationService: (any CommentGenerationServiceProtocol)?
    private let historyManager: AIHistoryManager
    private let freeUsageManager: AIFreeUsageManager
    var onPaywallRequired: (() -> Void)?
    private(set) var remainingFreeAIUses: Int
    private let emptyDraft = CommentDraft()

    init(
        hasPremiumAccess: Bool = false,
        generationService: (any CommentGenerationServiceProtocol)? = nil,
        historyManager: AIHistoryManager = .shared,
        freeUsageManager: AIFreeUsageManager = .shared
    ) {
        self.hasPremiumAccess = hasPremiumAccess
        self.injectedGenerationService = generationService
        self.historyManager = historyManager
        self.freeUsageManager = freeUsageManager
        self.remainingFreeAIUses = freeUsageManager.remainingUses
    }

    func setSubscriptionState(hasPremiumAccess: Bool, hasEverPurchasedPremium: Bool) {
        self.hasPremiumAccess = hasPremiumAccess
        self.hasEverPurchasedPremium = hasEverPurchasedPremium
        refreshFreeUsageState()
    }

    func setPremiumAccess(_ hasPremiumAccess: Bool) {
        setSubscriptionState(
            hasPremiumAccess: hasPremiumAccess,
            hasEverPurchasedPremium: hasEverPurchasedPremium
        )
    }

    func refreshFreeUsageState() {
        remainingFreeAIUses = freeUsageManager.remainingUses
    }

    private var generationService: any CommentGenerationServiceProtocol {
        injectedGenerationService ?? CommentGenerationServiceFactory.make(usesLiveAI: usesLiveAI)
    }

    var formattedSubreddit: String {
        let name = PostDraftValidator.normalizedSubreddit(draft.subreddit)
        guard !name.isEmpty else { return "r/subreddit" }
        return "r/\(name)"
    }

    var usesLiveAI: Bool {
        AIConfiguration.usesLiveAI(
            hasPremiumAccess: hasPremiumAccess,
            canUseFreeTrial: freeUsageManager.canUseFreeLiveAI(
                hasEverPurchasedPremium: hasEverPurchasedPremium
            )
        )
    }

    var generationEngineSubtitle: String {
        if usesLiveAI {
            if hasPremiumAccess {
                return "Write Reddit-ready comments with AI"
            }
            let remaining = remainingFreeAIUses
            let useLabel = remaining == 1 ? "use" : "uses"
            return "Write Reddit-ready comments with AI (\(remaining) free \(useLabel) left)"
        }
        if hasEverPurchasedPremium {
            return "Upgrade to Premium for unlimited AI generations"
        }
        return "Write Reddit-ready comments with AI templates"
    }

    var hasDraftChanges: Bool {
        draft != emptyDraft
    }

    var canGenerate: Bool {
        CommentDraftValidator.canGenerate(from: draft, isGenerating: isGenerating)
    }

    var canExport: Bool {
        let body = activeCommentBody.trimmingCharacters(in: .whitespaces)
        guard !body.isEmpty, body.count <= CommentDraftValidator.maxCommentLength else { return false }
        return (try? CommentDraftValidator.validateForGeneration(from: draft)) != nil
    }

    var activeCommentBody: String {
        if let selectedID = selectedVariantID,
           let variant = variants.first(where: { $0.id == selectedID }) {
            return variant.body
        }
        return draft.body
    }

    var bestVariant: CommentVariant? {
        variants.max(by: { $0.score < $1.score })
    }

    var needsOverwriteConfirmation: Bool {
        !draft.body.trimmingCharacters(in: .whitespaces).isEmpty
    }

    func selectCommentType(_ type: CommentType) {
        draft.commentType = type
        if type == .topLevel {
            draft.parentComment = ""
        }
        clearMessages()
    }

    func selectVariant(_ variant: CommentVariant) {
        selectedVariantID = variant.id
        clearMessages()
    }

    func applyVariant(_ variant: CommentVariant) {
        draft.body = variant.body
        selectedVariantID = variant.id
        successMessage = "Variant applied."
        errorMessage = nil
    }

    func generateComment() async {
        guard canGenerate else { return }

        if requiresPaywallForGeneration() {
            onPaywallRequired?()
            return
        }

        if needsOverwriteConfirmation {
            showOverwriteConfirmation = true
            return
        }

        await performGenerate(replaceExisting: true)
    }

    func confirmOverwriteAndGenerate() async {
        showOverwriteConfirmation = false
        await performGenerate(replaceExisting: true)
    }

    func cancelOverwriteConfirmation() {
        showOverwriteConfirmation = false
    }

    func copyToClipboard() {
        let content = activeCommentBody.trimmingCharacters(in: .whitespaces)
        guard !content.isEmpty else {
            errorMessage = CommentDraftValidationError.emptyComment.localizedDescription
            successMessage = nil
            return
        }
        guard content.count <= CommentDraftValidator.maxCommentLength else {
            errorMessage = CommentDraftValidationError.commentTooLong.localizedDescription
            successMessage = nil
            return
        }

        NSPasteboard.general.clearContents()
        NSPasteboard.general.setString(content, forType: .string)
        successMessage = "Comment copied to clipboard."
        errorMessage = nil
    }

    func resetDraft() {
        draft = CommentDraft()
        hasGeneratedContent = false
        selectedTab = .compose
        variants = []
        selectedVariantID = nil
        clearMessages()
    }

    func restore(from entry: AIHistoryEntry) {
        guard case .commentWriter(let storedDraft, let storedVariants, let selectedID) = entry.payload else {
            return
        }
        draft = storedDraft.commentDraft
        variants = storedVariants.map(\.commentVariant)
        selectedVariantID = selectedID
        hasGeneratedContent = true
        selectedTab = variants.isEmpty ? .compose : .preview
        clearMessages()
        successMessage = "Restored from history."
    }

    func clearMessages() {
        errorMessage = nil
        successMessage = nil
    }

    func notifyDraftEdited() {
        clearMessages()
    }

    private func performGenerate(replaceExisting: Bool) async {
        if requiresPaywallForGeneration() {
            onPaywallRequired?()
            return
        }

        isGenerating = true
        errorMessage = nil
        successMessage = nil

        let usedLiveAI = usesLiveAI

        do {
            let result = try await generationService.generateComment(from: draft)
            if replaceExisting || draft.body.trimmingCharacters(in: .whitespaces).isEmpty {
                draft.body = String(result.body.prefix(CommentDraftValidator.maxCommentLength))
            }
            variants = result.variants
            selectedVariantID = result.variants.first?.id
            selectedTab = .preview
            if usedLiveAI, !hasPremiumAccess, !hasEverPurchasedPremium {
                freeUsageManager.recordUse(for: .commentWriter)
                refreshFreeUsageState()
            }
            historyManager.save(AIHistoryEntry.fromCommentWriter(
                draft: draft,
                variants: variants,
                selectedVariantID: selectedVariantID
            ))
            hasGeneratedContent = true
            let engine = usedLiveAI ? "AI" : "local AI templates"
            successMessage = "Generated comment with \(result.variants.count) variants using \(engine)."
        } catch {
            errorMessage = UserFacingError.message(for: error)
        }

        isGenerating = false
    }

    private func requiresPaywallForGeneration() -> Bool {
        !hasPremiumAccess && (hasEverPurchasedPremium || remainingFreeAIUses == 0)
    }
}