notes/Notes/Storage/SyncEngine.swift

import Foundation
import IndieSync
import SwiftData

enum ICloudStatus: Equatable, Sendable {
    case checking
    case available
    case unavailable
}

/// Orchestrates iCloud JSON files (the sole source of truth) and the
/// SwiftData cache. Also exposes the domain operations the UI calls: create,
/// update, delete a note.
///
/// The file layer lives in the IndieSync package: `DocumentFileStore` for
/// coordinated, conflict-resolved, eviction-safe I/O; `TombstoneStore` for
/// soft-delete stubs; `MetadataObserver` for change events. Data flows one
/// way — writes touch files only; the observers are the sole writers of the
/// cache, so a local save reaches it through the same pipeline as a change
/// from another device.
///
/// Layout inside the container (record paths in this engine are
/// Records-relative and get prefixed via `dataPath(_:)` at the store
/// boundary):
/// ```
/// Documents/
///   Records/YYYY/MM/<ULID>.json
///   Stubs/YYYY/MM/<ULID>.json      ← soft-delete markers
/// ```
@Observable
@MainActor
final class SyncEngine {
    private(set) var iCloudStatus: ICloudStatus = .checking
    private(set) var isSyncing = false

    /// Last non-fatal sync failure, surfaced for the UI to render. Clears on
    /// the next successful write or rebuild.
    private(set) var lastSyncError: String?

    static let recordsDirectory = "Records"
    static let stubsDirectory = "Stubs"

    /// Store-rooted path of a record (record paths are Records-relative
    /// everywhere inside this engine).
    static func dataPath(_ recordsRelativePath: String) -> String {
        "\(recordsDirectory)/\(recordsRelativePath)"
    }

    private let modelContainer: ModelContainer
    private var store: DocumentFileStore?
    private var tombstones: TombstoneStore?
    private var dataMonitor: MetadataObserver?
    private var stubMonitor: MetadataObserver?
    private var dataMonitorTask: Task<Void, Never>?
    private var stubMonitorTask: Task<Void, Never>?

    private var cacheContext: ModelContext { modelContainer.mainContext }

    init(modelContainer: ModelContainer) {
        self.modelContainer = modelContainer
    }

    // MARK: - Connection

    private var connectAttempt = 0

    /// Resolves the ubiquity container and connects. Safe to call repeatedly —
    /// the gate's "Try Again" button and foreground reactivation reuse it.
    func connect() async {
        guard iCloudStatus != .available else { return }
        connectAttempt += 1
        let attempt = connectAttempt
        iCloudStatus = .checking

        // Blocking call — never on the main thread.
        let url = await Task.detached {
            FileManager.default.url(forUbiquityContainerIdentifier: nil)
        }.value

        guard let containerURL = url else {
            if attempt == connectAttempt {
                iCloudStatus = .unavailable
            }
            return
        }

        let store = DocumentFileStore(
            root: containerURL.appendingPathComponent("Documents", isDirectory: true))

        // File ops inside the container can hang indefinitely while the
        // iCloud account is in a bad state — flip the gate to "unavailable"
        // after a timeout, but let the attempt keep running and finish
        // connecting if it eventually succeeds.
        let timeout = Task { [weak self] in
            try? await Task.sleep(for: .seconds(20))
            guard let self, !Task.isCancelled else { return }
            if self.iCloudStatus == .checking, attempt == self.connectAttempt {
                self.iCloudStatus = .unavailable
            }
        }

        await store.prepareDirectories([Self.recordsDirectory, Self.stubsDirectory])
        timeout.cancel()
        guard attempt == connectAttempt else { return }   // superseded by a retry

        self.store = store
        self.tombstones = TombstoneStore(store: store, stubsDirectory: Self.stubsDirectory)
        iCloudStatus = .available

        startMonitoring()
        await syncCacheOnConnect()
        performLaunchMaintenance()
    }

    // MARK: - Domain operations

    @discardableResult
    func createNote(
        text: String,
        source: Note.Source = .typed,
        tags: [String] = [],
        context: CaptureContext
    ) async throws -> Note {
        let note = Note.create(text: text, source: source, tags: tags, context: context)
        try await save(note)
        return note
    }

    func updateNote(_ note: Note) async throws {
        var updated = note
        updated.meta.modifiedAt = Date()
        try await save(updated)
    }

    func deleteNote(id: ULID) async throws {
        // Without the file layer only the cache row could go — and the iCloud
        // file would re-import on the next rebuild, resurrecting the note.
        // Refuse instead.
        guard let tombstones else {
            report("Delete failed — iCloud not connected")
            throw SyncError.iCloudUnavailable
        }
        let path = Note.relativePath(forID: id)
        try await tombstones.softDelete(dataRelativePath: Self.dataPath(path), at: path)
        lastSyncError = nil
        // Cache cleanup happens via the monitors (file removal + stub
        // appearance both drop the row), but do it eagerly too so the UI
        // updates instantly even while iCloud event delivery lags.
        removeFromCache(jsonRelativePath: path)
    }

    /// Reads the JSON file — the authority — so read-after-write flows never
    /// act on a stale cache. Falls back to the near-lossless cache copy only
    /// if the file is unreadable (offline or not yet downloaded).
    func loadNote(id: ULID) async -> Note? {
        let path = Note.relativePath(forID: id)
        if let store,
           let note = try? await store.read(Note.self, from: Self.dataPath(path)),
           note.isReadable {
            return migrateIfNeeded(note)
        }
        guard let entity = NoteMapper.fetch(id: id.stringValue, in: cacheContext) else { return nil }
        return NoteMapper.note(from: entity)
    }

    // MARK: - Persistence plumbing

    /// Writes touch the file layer ONLY. iCloud JSON is the source of truth:
    /// with no file store the write would exist nowhere durable (a cache-only
    /// copy is wiped by the next rebuild), so refuse the save. The UI gate is
    /// a courtesy; this throw is the invariant.
    private func save(_ note: Note) async throws {
        guard let store else {
            report("Note not saved — iCloud not connected")
            throw SyncError.iCloudUnavailable
        }
        do {
            _ = try await store.write(note, to: Self.dataPath(note.relativePath))
        } catch {
            report("Failed to save note", error)
            throw error
        }
        lastSyncError = nil
    }

    /// Record a non-fatal sync failure: log it and publish it as
    /// `lastSyncError` for the UI to render.
    private func report(_ message: String, _ error: Error? = nil) {
        let full = error.map { "\(message): \($0.localizedDescription)" } ?? message
        print("[Sync] \(full)")
        lastSyncError = full
    }

    private func removeFromCache(jsonRelativePath: String) {
        let entities = NoteMapper.fetch(jsonRelativePath: jsonRelativePath, in: cacheContext)
        guard !entities.isEmpty else { return }
        entities.forEach(cacheContext.delete)
        try? cacheContext.save()
    }

    private func removeFromCache(ids: Set<String>) {
        for id in ids {
            if let entity = NoteMapper.fetch(id: id, in: cacheContext) {
                cacheContext.delete(entity)
            }
        }
        try? cacheContext.save()
    }

    // MARK: - Monitoring

    private func startMonitoring() {
        guard let store else { return }

        // Idempotent: tear down any existing monitors first so a reconnect
        // can't leave two live NSMetadataQuery monitors feeding duplicate events.
        stopMonitoring()

        let recordsURL = store.rootURL.appendingPathComponent(Self.recordsDirectory, isDirectory: true)
        let dataMonitor = MetadataObserver(documentsURL: recordsURL)
        self.dataMonitor = dataMonitor
        dataMonitor.start()
        dataMonitorTask = Task { [weak self] in
            for await batch in dataMonitor.events() {
                for event in batch {
                    await self?.handleFileEvent(event)
                }
            }
        }

        let stubsURL = store.rootURL.appendingPathComponent(Self.stubsDirectory, isDirectory: true)
        let stubMonitor = MetadataObserver(documentsURL: stubsURL)
        self.stubMonitor = stubMonitor
        stubMonitor.start()
        stubMonitorTask = Task { [weak self] in
            for await batch in stubMonitor.events() {
                for event in batch {
                    if case .added(let path) = event {
                        await self?.handleStubAdded(relativePath: path)
                    }
                }
            }
        }
    }

    private func stopMonitoring() {
        dataMonitorTask?.cancel()
        stubMonitorTask?.cancel()
        dataMonitorTask = nil
        stubMonitorTask = nil
        dataMonitor?.stop()
        stubMonitor?.stop()
        dataMonitor = nil
        stubMonitor = nil
    }

    private func handleFileEvent(_ event: FileChangeEvent) async {
        guard let store, let tombstones else { return }
        switch event {
        case .added(let path), .modified(let path):
            // Match stubs by document id, not path equality, so an evicted
            // stub still vetoes resurrection.
            let id = DocumentFileStore.documentID(fromRelativePath: path)
            if await tombstones.stubExists(id: id) {
                try? await store.remove(at: Self.dataPath(path))
                return
            }
            await importRecord(at: path, using: store)

        case .removed(let path):
            removeFromCache(jsonRelativePath: path)
        }
    }

    private func handleStubAdded(relativePath: String) async {
        if let store {
            try? await store.remove(at: Self.dataPath(relativePath))
        }
        removeFromCache(jsonRelativePath: relativePath)
    }

    private func importRecord(at path: String, using store: DocumentFileStore) async {
        // The store's read resolves conflicts and force-downloads evicted
        // files (bounded wait) before the coordinated read.
        guard let note = try? await store.read(Note.self, from: Self.dataPath(path)) else { return }
        // Forward gate: a document written by a newer app version is
        // quarantined — Codable would silently drop its unknown fields.
        guard note.isReadable else {
            print("[Sync] Skipping \(path) — schema \(note.schemaVersion) is newer than this app understands")
            return
        }
        NoteMapper.upsert(migrateIfNeeded(note), relativePath: path, into: cacheContext)
        try? cacheContext.save()
    }

    // MARK: - Cache rebuild / reconcile

    /// Bring the cache in line with iCloud on connect: full rebuild when the
    /// cache is empty (first launch, post-wipe), incremental reconcile
    /// otherwise — records synced down while the app was closed must appear.
    private func syncCacheOnConnect() async {
        guard store != nil else { return }
        let count = (try? cacheContext.fetchCount(FetchDescriptor<NoteEntity>())) ?? 0
        if count == 0 {
            await rebuildCache()
        } else {
            await reconcile()
        }
    }

    /// Diff iCloud's files against the cache and settle the difference:
    /// import records that synced down while the app was closed/backgrounded,
    /// drop rows whose file was deleted elsewhere, and clean up files a stub
    /// has since tombstoned. Runs on every connect and foreground return.
    func reconcile() async {
        guard let store, let tombstones else { return }

        let filePaths = await listRecordPaths(in: store)
        let stubIDs = await tombstones.listStubIDs()

        var livePathsByID: [String: String] = [:]
        var orphanedPaths: [String] = []
        for path in filePaths {
            let id = DocumentFileStore.documentID(fromRelativePath: path)
            if stubIDs.contains(id) {
                orphanedPaths.append(path)
            } else {
                livePathsByID[id] = path
            }
        }

        var descriptor = FetchDescriptor<NoteEntity>()
        descriptor.propertiesToFetch = [\.id]
        let cachedIDs = Set(((try? cacheContext.fetch(descriptor)) ?? []).map(\.id))

        let toImport = Set(livePathsByID.keys).subtracting(cachedIDs)
        let toRemove = cachedIDs.subtracting(livePathsByID.keys)

        guard !toImport.isEmpty || !toRemove.isEmpty || !orphanedPaths.isEmpty else {
            return  // already in sync — the common foreground case, kept cheap
        }
        print("[Sync] Reconcile: +\(toImport.count)/-\(toRemove.count) notes, \(orphanedPaths.count) orphaned files")

        isSyncing = true
        defer { isSyncing = false }

        // Clean up files a stub has tombstoned (mirrors the rebuild path).
        for path in orphanedPaths {
            try? await store.remove(at: Self.dataPath(path))
        }

        removeFromCache(ids: toRemove)

        // Import the missing records through the same conflict-resolved,
        // eviction-safe, forward-gated path the monitors use.
        for id in toImport.sorted() {
            guard let path = livePathsByID[id] else { continue }
            await importRecord(at: path, using: store)
        }
    }

    func rebuildCache() async {
        guard let store, let tombstones else { return }
        isSyncing = true
        defer { isSyncing = false }

        await store.downloadEvictedFiles(inSubdirectory: Self.recordsDirectory)
        for _ in 0..<30 {
            if await store.countEvictedFiles(inSubdirectory: Self.recordsDirectory) == 0 { break }
            try? await Task.sleep(for: .seconds(2))
        }

        do {
            try cacheContext.delete(model: NoteEntity.self)
            try cacheContext.save()
        } catch {
            report("Cache rebuild failed", error)
            return
        }

        let stubIDs = await tombstones.listStubIDs()
        for path in await listRecordPaths(in: store) {
            if stubIDs.contains(DocumentFileStore.documentID(fromRelativePath: path)) {
                try? await store.remove(at: Self.dataPath(path))
                continue
            }
            await importRecord(at: path, using: store)
        }
        lastSyncError = nil
    }

    /// Records-relative paths of all live record files (placeholder-aware —
    /// includes evicted iCloud-only files).
    private func listRecordPaths(in store: DocumentFileStore) async -> [String] {
        let prefix = Self.recordsDirectory + "/"
        return await store.list(inSubdirectory: Self.recordsDirectory)
            .map { String($0.dropFirst(prefix.count)) }
            .sorted()
    }

    // MARK: - Restore support

    /// Stops the metadata observers so a backup restore's bulk file
    /// replacement isn't observed as live sync events.
    func suspendForRestore() {
        stopMonitoring()
    }

    /// Restarts fresh observers after a restore. The new NSMetadataQuery
    /// re-baselines its known files on gather, so the restored tree does not
    /// replay as a flood of add/remove events.
    func resumeAfterRestore() {
        startMonitoring()
    }

    // MARK: - Maintenance

    private func performLaunchMaintenance() {
        Task(priority: .utility) { [tombstones] in
            _ = try? await tombstones?.prune()
        }
    }

    // MARK: - Schema migration

    /// Read-time upgrade for documents behind the current schema. A no-op
    /// while `currentSchemaVersion` is 1, but wired into both read paths so
    /// the first schema bump migrates instead of silently skipping. Write-back
    /// happens lazily on the next save — never from live monitor events.
    private func migrateIfNeeded(_ note: Note) -> Note {
        guard note.meta.schemaVersion < Note.currentSchemaVersion else { return note }
        var migrated = note
        migrated.meta.schemaVersion = Note.currentSchemaVersion
        return migrated
    }
}