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

import AppKit
import SwiftUI

@MainActor
@Observable
final class TitleOptimizerViewModel {
    var draft = TitleDraft()
    var selectedTab: TitleOptimizerTab = .optimize
    var isOptimizing = false
    var errorMessage: String?
    var successMessage: String?
    var analysis: TitleAnalysis?
    var variants: [TitleVariant] = []
    var selectedVariantID: UUID?

    private var hasPremiumAccess = false
    private var hasEverPurchasedPremium = false
    private var optimizationService: any TitleOptimizationServiceProtocol
    private let injectedOptimizationService: (any TitleOptimizationServiceProtocol)?
    private let historyManager: AIHistoryManager
    private let freeUsageManager: AIFreeUsageManager
    var onPaywallRequired: (() -> Void)?
    private(set) var remainingFreeAIUses: Int

    init(
        hasPremiumAccess: Bool = false,
        optimizationService: (any TitleOptimizationServiceProtocol)? = nil,
        historyManager: AIHistoryManager = .shared,
        freeUsageManager: AIFreeUsageManager = .shared
    ) {
        self.hasPremiumAccess = hasPremiumAccess
        self.injectedOptimizationService = optimizationService
        self.historyManager = historyManager
        self.freeUsageManager = freeUsageManager
        self.remainingFreeAIUses = freeUsageManager.remainingUses
        self.hasEverPurchasedPremium = false
        self.optimizationService = optimizationService
            ?? TitleOptimizationServiceFactory.make(usesLiveAI: Self.liveAIEnabled(
                hasPremiumAccess: hasPremiumAccess,
                hasEverPurchasedPremium: false,
                freeUsageManager: freeUsageManager
            ))
    }

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

    private func refreshOptimizationService() {
        if injectedOptimizationService == nil {
            optimizationService = TitleOptimizationServiceFactory.make(usesLiveAI: usesLiveAI)
        }
    }

    private static func liveAIEnabled(
        hasPremiumAccess: Bool,
        hasEverPurchasedPremium: Bool,
        freeUsageManager: AIFreeUsageManager
    ) -> Bool {
        AIConfiguration.usesLiveAI(
            hasPremiumAccess: hasPremiumAccess,
            canUseFreeTrial: freeUsageManager.canUseFreeLiveAI(
                hasEverPurchasedPremium: hasEverPurchasedPremium
            )
        )
    }

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

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

    var formattedSubreddit: String {
        let trimmed = draft.subreddit.trimmingCharacters(in: .whitespaces)
        guard !trimmed.isEmpty else { return "r/subreddit" }
        return trimmed.hasPrefix("r/") ? trimmed : "r/\(trimmed)"
    }

    var canOptimize: Bool {
        !isOptimizing && missingOptimizeFields.isEmpty
    }

    var missingOptimizeFields: [String] {
        var missing: [String] = []
        if draft.subreddit.trimmingCharacters(in: .whitespaces).isEmpty {
            missing.append("subreddit")
        }
        if draft.topic.trimmingCharacters(in: .whitespaces).isEmpty {
            missing.append("topic")
        }
        if draft.originalTitle.trimmingCharacters(in: .whitespaces).isEmpty {
            missing.append("title")
        }
        return missing
    }

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

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

    var previewDraft: PostDraft {
        PostDraft(
            postType: draft.postType,
            subreddit: draft.subreddit,
            topic: draft.topic,
            tone: draft.tone,
            title: activeTitle,
            isNSFW: draft.isNSFW,
            isSpoiler: draft.isSpoiler,
            isOC: draft.isOC,
            flair: draft.flair
        )
    }

    func selectPostType(_ type: RedditPostType) {
        draft.postType = type
        clearMessages()
    }

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

    func applyVariant(_ variant: TitleVariant) {
        draft.originalTitle = variant.title
        selectedVariantID = variant.id
        successMessage = "Title applied."
        errorMessage = nil
    }

    func optimizeTitles() async {
        let missing = missingOptimizeFields
        guard missing.isEmpty else {
            errorMessage = "Fill in \(missing.joined(separator: ", ")) before optimizing."
            successMessage = nil
            return
        }

        if requiresPaywallForGeneration() {
            onPaywallRequired?()
            return
        }

        isOptimizing = true
        errorMessage = nil
        successMessage = nil
        await Task.yield()
        defer { isOptimizing = false }

        let usedLiveAI = usesLiveAI

        do {
            let result = try await optimizationService.optimizeTitles(from: draft)
            analysis = result.analysis
            variants = result.variants
            selectedVariantID = result.variants.first?.id
            if usedLiveAI, !hasPremiumAccess, !hasEverPurchasedPremium {
                freeUsageManager.recordUse(for: .titleOptimizer)
                refreshFreeUsageState()
            }
            let engine = usedLiveAI ? "AI" : "local AI templates"
            successMessage = "Generated \(result.variants.count) title suggestions using \(engine)."
            let entry = AIHistoryEntry.fromTitleOptimizer(
                draft: draft,
                analysis: analysis,
                variants: variants,
                selectedVariantID: selectedVariantID
            )
            Task { @MainActor in
                historyManager.save(entry)
            }
        } catch {
            errorMessage = UserFacingError.message(for: error)
        }
    }

    func copyToClipboard() {
        let title = activeTitle.trimmingCharacters(in: .whitespaces)
        guard !title.isEmpty else { return }

        NSPasteboard.general.clearContents()
        NSPasteboard.general.setString(title, forType: .string)
        successMessage = "Copied to clipboard."
        errorMessage = nil
    }

    func resetDraft() {
        draft = TitleDraft()
        selectedTab = .optimize
        analysis = nil
        variants = []
        selectedVariantID = nil
        clearMessages()
    }

    func restore(from entry: AIHistoryEntry) {
        guard case .titleOptimizer(let storedDraft, let storedAnalysis, let storedVariants, let selectedID) = entry.payload else {
            return
        }
        draft = storedDraft.titleDraft
        analysis = storedAnalysis?.titleAnalysis
        variants = storedVariants.map(\.titleVariant)
        selectedVariantID = selectedID
        selectedTab = variants.isEmpty ? .optimize : .compare
        clearMessages()
        successMessage = "Restored from history."
    }

    func exportText() -> String {
        var lines: [String] = []
        lines.append("Subreddit: \(formattedSubreddit)")
        lines.append("Topic: \(draft.topic)")
        lines.append("Goal: \(draft.titleGoal.title)")
        lines.append("Tone: \(draft.tone.title)")
        lines.append("")
        lines.append("Original: \(draft.originalTitle)")

        if let analysis {
            lines.append("")
            lines.append("Analysis Score: \(analysis.overallScore)/100")
            lines.append("  Length: \(analysis.lengthScore)")
            lines.append("  Engagement: \(analysis.engagementScore)")
            lines.append("  Clarity: \(analysis.clarityScore)")
        }

        if !variants.isEmpty {
            lines.append("")
            lines.append("Suggestions:")
            for (index, variant) in variants.enumerated() {
                lines.append("\(index + 1). [\(variant.score)] \(variant.title)")
            }
        }

        return lines.joined(separator: "\n")
    }

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

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