prompt/Sources/Prompt/Store.swift

import AppKit

enum ItemContent: Codable, Equatable {
    case text(String)
    case image(String) // filename inside the Images directory
}

struct PromptItem: Codable, Identifiable, Equatable {
    var id: UUID = UUID()
    var content: ItemContent
    var done: Bool = false
    var createdAt: Date = Date()
    var sourceApp: String? = nil
}

struct Section: Codable, Identifiable, Equatable {
    var id: UUID = UUID()
    var name: String
    var items: [PromptItem] = []
}

private struct StoreFile: Codable {
    var sections: [Section]
    var activeSectionID: UUID?
}

@MainActor
final class Store: ObservableObject {
    static let shared = Store()

    @Published private(set) var sections: [Section] = []
    @Published var activeSectionID: UUID? {
        didSet { save() }
    }

    let directory: URL
    let imagesDirectory: URL
    private let storeFile: URL
    private let imageCache = NSCache<NSString, NSImage>()

    /// `directory` is injectable so the store can be exercised against a
    /// scratch folder instead of the user's real data.
    init(directory: URL? = nil) {
        if let directory {
            self.directory = directory
        } else {
            let appSupport = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0]
            self.directory = appSupport.appendingPathComponent("Prompt", isDirectory: true)
        }
        imagesDirectory = self.directory.appendingPathComponent("Images", isDirectory: true)
        storeFile = self.directory.appendingPathComponent("store.json")

        try? FileManager.default.createDirectory(at: imagesDirectory, withIntermediateDirectories: true)
        load()
        if sections.isEmpty {
            let workspace = Section(name: "Default")
            sections = [workspace]
            activeSectionID = workspace.id
            save()
        }
        if activeSectionID == nil || !sections.contains(where: { $0.id == activeSectionID }) {
            activeSectionID = sections.first?.id
        }
    }

    var activeSection: Section {
        if let id = activeSectionID, let section = sections.first(where: { $0.id == id }) {
            return section
        }
        return sections[0]
    }

    // MARK: - Persistence

    private func load() {
        guard let data = try? Data(contentsOf: storeFile),
              let file = try? JSONDecoder().decode(StoreFile.self, from: data) else { return }
        sections = file.sections
        activeSectionID = file.activeSectionID
    }

    private func save() {
        let file = StoreFile(sections: sections, activeSectionID: activeSectionID)
        guard let data = try? JSONEncoder().encode(file) else { return }
        try? data.write(to: storeFile, options: .atomic)
        onDidChange?()
    }

    // MARK: - Sync support

    /// Called after every save so `SyncService` can push the change.
    var onDidChange: (() -> Void)?

    func replaceAll(sections newSections: [Section], activeSectionID newActiveSectionID: UUID?) {
        sections = newSections
        if let newActiveSectionID, newSections.contains(where: { $0.id == newActiveSectionID }) {
            activeSectionID = newActiveSectionID
        } else {
            activeSectionID = newSections.first?.id
        }
        save()
    }

    func imageData(for filename: String) -> Data? {
        try? Data(contentsOf: imagesDirectory.appendingPathComponent(filename))
    }

    func importImage(filename: String, data: Data) {
        let url = imagesDirectory.appendingPathComponent(filename)
        if let existing = try? Data(contentsOf: url), existing == data { return }
        try? data.write(to: url)
        imageCache.removeObject(forKey: filename as NSString)
    }

    func cleanupOrphanedImages(referenced: Set<String>) {
        let files = (try? FileManager.default.contentsOfDirectory(at: imagesDirectory, includingPropertiesForKeys: nil)) ?? []
        for file in files where !referenced.contains(file.lastPathComponent) {
            try? FileManager.default.removeItem(at: file)
            imageCache.removeObject(forKey: file.lastPathComponent as NSString)
        }
    }

    // MARK: - Sections

    @discardableResult
    func createSection(named rawName: String) -> Section {
        let name = rawName.trimmingCharacters(in: .whitespacesAndNewlines)
        if let existing = sections.first(where: { $0.name.caseInsensitiveCompare(name) == .orderedSame }) {
            activeSectionID = existing.id
            return existing
        }
        let section = Section(name: name)
        sections.append(section)
        activeSectionID = section.id
        save()
        return section
    }

    /// Reorders a workspace within the list. The UI calls these "workspaces";
    /// the model and the persisted JSON keep the original `sections` naming so
    /// existing stores and sync payloads keep working.
    func moveSection(_ id: UUID, by offset: Int) {
        guard let index = sections.firstIndex(where: { $0.id == id }) else { return }
        let target = index + offset
        guard sections.indices.contains(target) else { return }
        let section = sections.remove(at: index)
        sections.insert(section, at: target)
        save()
    }

    // MARK: - Items

    @discardableResult
    func addText(_ text: String, sourceApp: String? = nil) -> PromptItem {
        let item = PromptItem(content: .text(text), sourceApp: sourceApp)
        append(item)
        return item
    }

    @discardableResult
    func addImage(_ image: NSImage) -> PromptItem? {
        guard let tiff = image.tiffRepresentation,
              let rep = NSBitmapImageRep(data: tiff),
              let png = rep.representation(using: .png, properties: [:]) else { return nil }
        let filename = "\(UUID().uuidString).png"
        do {
            try png.write(to: imagesDirectory.appendingPathComponent(filename))
        } catch {
            return nil
        }
        let item = PromptItem(content: .image(filename))
        append(item)
        return item
    }

    private func append(_ item: PromptItem) {
        guard let index = sections.firstIndex(where: { $0.id == activeSection.id }) else { return }
        sections[index].items.append(item)
        save()
    }

    func toggleDone(_ id: UUID) {
        mutateItem(id) { $0.done.toggle() }
    }

    func updateText(_ id: UUID, newText: String) {
        mutateItem(id) { item in
            if case .text = item.content { item.content = .text(newText) }
        }
    }

    func delete(_ id: UUID) {
        delete(ids: [id])
    }

    /// Deletes every item in `ids` in one pass, cleaning up their image files.
    func delete(ids: Set<UUID>) {
        guard !ids.isEmpty else { return }
        var removedAny = false
        for sectionIndex in sections.indices {
            for item in sections[sectionIndex].items where ids.contains(item.id) {
                removedAny = true
                if case .image(let filename) = item.content {
                    try? FileManager.default.removeItem(at: imagesDirectory.appendingPathComponent(filename))
                    imageCache.removeObject(forKey: filename as NSString)
                }
            }
            sections[sectionIndex].items.removeAll { ids.contains($0.id) }
        }
        if removedAny { save() }
    }

    func setDone(_ done: Bool, ids: Set<UUID>) {
        guard !ids.isEmpty else { return }
        for sectionIndex in sections.indices {
            for itemIndex in sections[sectionIndex].items.indices
            where ids.contains(sections[sectionIndex].items[itemIndex].id) {
                sections[sectionIndex].items[itemIndex].done = done
            }
        }
        save()
    }

    /// Moves every item in `ids` into `targetSectionID`, keeping their order.
    func move(ids: Set<UUID>, to targetSectionID: UUID) {
        guard !ids.isEmpty,
              let targetIndex = sections.firstIndex(where: { $0.id == targetSectionID }) else { return }
        var moved: [PromptItem] = []
        for sectionIndex in sections.indices where sectionIndex != targetIndex {
            moved.append(contentsOf: sections[sectionIndex].items.filter { ids.contains($0.id) })
            sections[sectionIndex].items.removeAll { ids.contains($0.id) }
        }
        guard !moved.isEmpty else { return }
        sections[targetIndex].items.append(contentsOf: moved)
        save()
    }

    func move(_ id: UUID, to targetSectionID: UUID) {
        var moved: PromptItem?
        for sectionIndex in sections.indices {
            if let itemIndex = sections[sectionIndex].items.firstIndex(where: { $0.id == id }) {
                moved = sections[sectionIndex].items.remove(at: itemIndex)
                break
            }
        }
        guard let item = moved,
              let targetIndex = sections.firstIndex(where: { $0.id == targetSectionID }) else { return }
        sections[targetIndex].items.append(item)
        save()
    }

    func merge(ids: Set<UUID>) {
        let merged = sections.flatMap(\.items)
            .filter { ids.contains($0.id) }
            .sorted { $0.createdAt < $1.createdAt }
        let texts = merged.compactMap { item -> String? in
            if case .text(let text) = item.content { return text }
            return nil
        }
        guard texts.count > 1, let first = merged.first else { return }
        let targetSectionID = sectionID(of: first.id)
        for item in merged where !item.content.isImage {
            delete(item.id)
        }
        let combined = PromptItem(content: .text(texts.joined(separator: "\n\n")))
        if let index = sections.firstIndex(where: { $0.id == targetSectionID }) {
            sections[index].items.append(combined)
            save()
        }
    }

    func clearDoneInActiveSection() {
        guard let index = sections.firstIndex(where: { $0.id == activeSection.id }) else { return }
        for item in sections[index].items where item.done {
            if case .image(let filename) = item.content {
                try? FileManager.default.removeItem(at: imagesDirectory.appendingPathComponent(filename))
                imageCache.removeObject(forKey: filename as NSString)
            }
        }
        sections[index].items.removeAll { $0.done }
        save()
    }

    func sectionID(of itemID: UUID) -> UUID? {
        for section in sections where section.items.contains(where: { $0.id == itemID }) {
            return section.id
        }
        return nil
    }

    func item(_ id: UUID) -> PromptItem? {
        sections.flatMap(\.items).first { $0.id == id }
    }

    private func mutateItem(_ id: UUID, _ mutate: (inout PromptItem) -> Void) {
        for sectionIndex in sections.indices {
            if let itemIndex = sections[sectionIndex].items.firstIndex(where: { $0.id == id }) {
                mutate(&sections[sectionIndex].items[itemIndex])
                save()
                return
            }
        }
    }

    // MARK: - Clipboard

    func copy(_ item: PromptItem) {
        let pasteboard = NSPasteboard.general
        pasteboard.clearContents()
        switch item.content {
        case .text(let text):
            pasteboard.setString(text, forType: .string)
        case .image(let filename):
            if let image = image(for: filename) {
                pasteboard.writeObjects([image])
            }
        }
    }

    func copyAsList(ids: Set<UUID>) {
        let texts = selectedTexts(ids: ids)
        guard !texts.isEmpty else { return }
        let list = texts.enumerated().map { "\($0.offset + 1). \($0.element)" }.joined(separator: "\n")
        writeToPasteboard(list)
    }

    /// Same selection as `copyAsList`, but joined with plain newlines so it
    /// drops into a prompt without numbering.
    func copyAsLines(ids: Set<UUID>) {
        let texts = selectedTexts(ids: ids)
        guard !texts.isEmpty else { return }
        writeToPasteboard(texts.joined(separator: "\n"))
    }

    private func selectedTexts(ids: Set<UUID>) -> [String] {
        sections.flatMap(\.items)
            .filter { ids.contains($0.id) }
            .sorted { $0.createdAt < $1.createdAt }
            .compactMap { item in
                if case .text(let text) = item.content { return text }
                return nil
            }
    }

    private func writeToPasteboard(_ string: String) {
        let pasteboard = NSPasteboard.general
        pasteboard.clearContents()
        pasteboard.setString(string, forType: .string)
    }

    // MARK: - Images

    func image(for filename: String) -> NSImage? {
        if let cached = imageCache.object(forKey: filename as NSString) { return cached }
        let url = imagesDirectory.appendingPathComponent(filename)
        guard let image = NSImage(contentsOf: url) else { return nil }
        imageCache.setObject(image, forKey: filename as NSString)
        return image
    }

    func imageURL(for filename: String) -> URL {
        imagesDirectory.appendingPathComponent(filename)
    }
}

private extension ItemContent {
    var isImage: Bool {
        if case .image = self { return true }
        return false
    }
}