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

import AppKit
import SwiftUI

@MainActor
@Observable
final class PostGeneratorViewModel {
    var draft = PostDraft()
    var selectedTab: PostGeneratorTab = .compose
    var isGenerating = false
    var errorMessage: String?
    var successMessage: String?
    var showImageImporter = false
    var showOverwriteConfirmation = false
    /// True after a successful Generate or History restore. Cleared on Reset.
    /// Controls whether title/body result fields are shown in Compose.
    private(set) var hasGeneratedContent = false

    private var hasPremiumAccess = false
    private var hasEverPurchasedPremium = false
    private var generationService: any PostGenerationServiceProtocol
    private let injectedGenerationService: (any PostGenerationServiceProtocol)?
    private let historyManager: AIHistoryManager
    private let freeUsageManager: AIFreeUsageManager
    var onPaywallRequired: (() -> Void)?
    private(set) var remainingFreeAIUses: Int
    private var imageAccessURL: URL?
    private var isAccessingImageResource = false
    private let emptyDraft = PostDraft()

    init(
        hasPremiumAccess: Bool = false,
        generationService: (any PostGenerationServiceProtocol)? = nil,
        historyManager: AIHistoryManager = .shared,
        freeUsageManager: AIFreeUsageManager = .shared
    ) {
        self.hasPremiumAccess = hasPremiumAccess
        self.injectedGenerationService = generationService
        self.historyManager = historyManager
        self.freeUsageManager = freeUsageManager
        self.remainingFreeAIUses = freeUsageManager.remainingUses
        self.hasEverPurchasedPremium = false
        self.generationService = generationService
            ?? PostGenerationServiceFactory.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
        refreshGenerationService()
    }

    private func refreshGenerationService() {
        if injectedGenerationService == nil {
            generationService = PostGenerationServiceFactory.make(usesLiveAI: usesLiveAI)
        }
    }

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

    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 "Create Reddit-ready posts with AI"
            }
            let remaining = remainingFreeAIUses
            let useLabel = remaining == 1 ? "use" : "uses"
            return "Create Reddit-ready posts with AI (\(remaining) free \(useLabel) left)"
        }
        if hasEverPurchasedPremium {
            return "Upgrade to Premium for unlimited AI generations"
        }
        return "Create Reddit-ready posts with AI templates"
    }

    var hasDraftChanges: Bool {
        draft != emptyDraft
    }

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

    var canExport: Bool {
        PostDraftValidator.canExport(draft)
    }

    var canAddPollOption: Bool {
        draft.pollOptions.count < 6
    }

    var canRemovePollOption: Bool {
        draft.pollOptions.count > 2
    }

    var needsOverwriteConfirmation: Bool {
        !draft.title.trimmingCharacters(in: .whitespaces).isEmpty
            || !draft.body.trimmingCharacters(in: .whitespaces).isEmpty
            || draft.pollOptions.contains { !$0.text.trimmingCharacters(in: .whitespaces).isEmpty }
    }

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

    func addPollOption() {
        guard canAddPollOption else { return }
        draft.pollOptions.append(PollOption())
        clearMessages()
    }

    func removePollOption(_ option: PollOption) {
        guard canRemovePollOption else { return }
        draft.pollOptions.removeAll { $0.id == option.id }
        clearMessages()
    }

    func updatePollOption(id: UUID, text: String) {
        guard let index = draft.pollOptions.firstIndex(where: { $0.id == id }) else { return }
        draft.pollOptions[index].text = text
        clearMessages()
    }

    func setImage(from url: URL) {
        releaseImageAccess()

        guard url.startAccessingSecurityScopedResource() else {
            errorMessage = "Couldn't access the selected image. Try choosing the file again."
            successMessage = nil
            return
        }

        imageAccessURL = url
        isAccessingImageResource = true
        draft.imageFileURL = url
        clearMessages()
    }

    func handleImageImportFailure(_ error: Error) {
        errorMessage = UserFacingError.message(for: error)
        successMessage = nil
    }

    func removeImage() {
        releaseImageAccess()
        draft.imageFileURL = nil
        clearMessages()
    }

    func generatePost() 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() {
        do {
            try PostDraftValidator.validateForExport(draft)
            let content = exportText()
            NSPasteboard.general.clearContents()
            NSPasteboard.general.setString(content, forType: .string)
            successMessage = "Copied to clipboard."
            errorMessage = nil
        } catch {
            errorMessage = UserFacingError.message(for: error)
            successMessage = nil
        }
    }

    func resetDraft() {
        releaseImageAccess()
        draft = PostDraft()
        hasGeneratedContent = false
        selectedTab = .compose
        clearMessages()
    }

    func restore(from entry: AIHistoryEntry) {
        guard case .postGenerator(let storedDraft) = entry.payload else { return }
        releaseImageAccess()
        draft = storedDraft.postDraft
        hasGeneratedContent = true
        selectedTab = .preview
        clearMessages()
        successMessage = "Restored from history."
    }

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

    func notifyDraftEdited() {
        clearMessages()
    }

    func exportText() -> String {
        var lines: [String] = []
        lines.append("Subreddit: \(formattedSubreddit)")
        lines.append("Type: \(draft.postType.title)")

        if !draft.title.isEmpty {
            lines.append("Title: \(draft.title)")
        }

        switch draft.postType {
        case .text:
            if !draft.body.isEmpty { lines.append("\n\(draft.body)") }
        case .image:
            if let url = draft.imageFileURL {
                lines.append("Image: \(url.lastPathComponent)")
            }
            if !draft.body.isEmpty { lines.append("Caption: \(draft.body)") }
        case .link:
            if !draft.linkURL.isEmpty { lines.append("URL: \(draft.linkURL)") }
            if !draft.body.isEmpty { lines.append("\n\(draft.body)") }
        case .video:
            if !draft.videoURL.isEmpty { lines.append("Video URL: \(draft.videoURL)") }
            if !draft.body.isEmpty { lines.append("\n\(draft.body)") }
        case .poll:
            lines.append("Duration: \(draft.pollDuration.label)")
            for (index, option) in draft.pollOptions.enumerated() where !option.text.isEmpty {
                lines.append("Option \(index + 1): \(option.text)")
            }
            if !draft.body.isEmpty { lines.append("\n\(draft.body)") }
        }

        var tags: [String] = []
        if draft.isNSFW { tags.append("NSFW") }
        if draft.isSpoiler { tags.append("Spoiler") }
        if draft.isOC { tags.append("OC") }
        if !tags.isEmpty { lines.append("Tags: \(tags.joined(separator: ", "))") }

        if !draft.flair.isEmpty { lines.append("Flair: \(draft.flair)") }

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

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

        isGenerating = true
        errorMessage = nil
        successMessage = nil
        await Task.yield()

        let usedLiveAI = usesLiveAI

        do {
            let result = try await generationService.generatePost(from: draft)
            applyGeneratedPost(result, replaceExisting: replaceExisting)
            selectedTab = .preview
            if usedLiveAI, !hasPremiumAccess, !hasEverPurchasedPremium {
                freeUsageManager.recordUse(for: .postGenerator)
                refreshFreeUsageState()
            }
            let engine = usedLiveAI ? "AI" : "local AI templates"
            successMessage = "Post generated successfully using \(engine)."
            let entry = AIHistoryEntry.fromPostGenerator(draft: draft)
            Task { @MainActor in
                historyManager.save(entry)
            }
        } catch {
            errorMessage = UserFacingError.message(for: error)
        }

        isGenerating = false
    }

    private func applyGeneratedPost(_ result: GeneratedPost, replaceExisting: Bool) {
        if replaceExisting || draft.title.trimmingCharacters(in: .whitespaces).isEmpty {
            draft.title = String(result.title.prefix(PostDraftValidator.maxTitleLength))
        }

        if replaceExisting || draft.body.trimmingCharacters(in: .whitespaces).isEmpty {
            draft.body = String(result.body.prefix(PostDraftValidator.maxBodyLength))
        }

        if let flair = result.suggestedFlair, draft.flair.trimmingCharacters(in: .whitespaces).isEmpty {
            draft.flair = flair
        }

        if draft.postType == .poll, let options = result.pollOptions, !options.isEmpty {
            if replaceExisting || draft.pollOptions.allSatisfy({ $0.text.trimmingCharacters(in: .whitespaces).isEmpty }) {
                draft.pollOptions = options.map { PollOption(text: $0) }
            }
        }

        hasGeneratedContent = true
    }

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

    private func releaseImageAccess() {
        if isAccessingImageResource, let imageAccessURL {
            imageAccessURL.stopAccessingSecurityScopedResource()
        }
        imageAccessURL = nil
        isAccessingImageResource = false
    }
}