import CryptoKit
import Foundation
import Observation

struct AgentDraftItem: Codable, Equatable, Identifiable, Sendable {
    enum State: String, Codable, Sendable { case loading, uploading, ready, failed }
    var id: String = UUID().uuidString // Stable multipart uploadId, including across relaunch.
    var filename: String
    var mimeType: String?
    var sizeBytes: Int?
    var state: State = .loading
    var attachment: AgentAttachment?
    var issue: AgentAttachmentIssue?
}
struct AgentDraftSnapshot: Codable, Sendable {
    var caption = ""
    var items: [AgentDraftItem] = []
    var id: String? = nil // Optional for drafts saved before durable outbox handoff tracking.
}
struct AgentDraftStorage: Sendable {
    var read: @Sendable () throws -> AgentDraftSnapshot
    var write: @Sendable (AgentDraftSnapshot) throws -> Void
    var stage: @Sendable (String, Data) throws -> Void
    var bytes: @Sendable (String) throws -> Data
    var remove: @Sendable (String) -> Void

    static func file(baseURL: URL, root: URL = .applicationSupportDirectory) -> Self {
        let key = SHA256.hash(data: Data(baseURL.absoluteString.utf8)).map { String(format: "%02x", $0) }.joined()
        let directory = root.appendingPathComponent("AgentAttachmentDrafts/\(key)", isDirectory: true)
        let manifest = directory.appendingPathComponent("draft.json")
        let file: @Sendable (String) throws -> URL = { id in
            guard UUID(uuidString: id) != nil else { throw AgentAttachmentIssue.local("Invalid staged upload identity.") }
            return directory.appendingPathComponent(id + ".bytes")
        }
        let write: @Sendable (Data, URL) throws -> Void = { data, destination in
            try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)
            try data.write(to: destination, options: [.atomic, .completeFileProtectionUntilFirstUserAuthentication])
        }
        return Self(read: {
            guard FileManager.default.fileExists(atPath: manifest.path) else { return AgentDraftSnapshot() }
            let data = try Data(contentsOf: manifest)
            guard data.count <= 2 * 1024 * 1024 else { throw AgentAttachmentIssue.local("The saved attachment draft is too large.") }
            let snapshot = try JSONDecoder().decode(AgentDraftSnapshot.self, from: data)
            guard snapshot.items.count <= 8, Set(snapshot.items.map(\.id)).count == snapshot.items.count,
                  snapshot.items.allSatisfy({ UUID(uuidString: $0.id) != nil }) else { throw AgentAttachmentIssue.local("The saved attachment draft is invalid.") }
            return snapshot
        }, write: { snapshot in
            let data = try JSONEncoder().encode(snapshot)
            guard data.count <= 2 * 1024 * 1024 else { throw AgentAttachmentIssue.local("The saved attachment draft is too large.") }
            try write(data, manifest)
        }, stage: { id, data in
            guard data.count <= 8 * 1024 * 1024 else { throw AgentAttachmentIssue.local("The staged file exceeds the size limit.") }
            try write(data, file(id))
        }, bytes: { id in
            let url = try file(id)
            let size = try url.resourceValues(forKeys: [.fileSizeKey]).fileSize ?? Int.max
            guard size <= 8 * 1024 * 1024 else { throw AgentAttachmentIssue.local("The staged file exceeds the size limit.") }
            return try Data(contentsOf: url)
        }, remove: { id in if let url = try? file(id) { try? FileManager.default.removeItem(at: url) } })
    }
}

@MainActor @Observable
final class AgentAttachmentDraft {
    private(set) var id = UUID().uuidString
    var caption = "" { didSet { save() } }
    private(set) var items: [AgentDraftItem] = []
    private(set) var capability: AgentAttachmentCapabilities?
    private(set) var error: String?
    private(set) var persistenceAvailable = true
    @ObservationIgnored private var client: (any AgentAPIClienting)?
    @ObservationIgnored private var storage: AgentDraftStorage?
    @ObservationIgnored private var prepared: [String: AgentPreparedAttachment] = [:]
    @ObservationIgnored private var loaders: [String: @Sendable () async throws -> AgentPreparedAttachment] = [:]
    @ObservationIgnored private var worker: Task<Void, Never>?
    @ObservationIgnored private var generation = UUID()
    @ObservationIgnored private var restoring = false
    @ObservationIgnored private var active = true
    var supportsAttachments: Bool { capability?.supported == true && persistenceAvailable }
    var ready: [AgentAttachment] { items.compactMap(\.attachment) }
    var canSend: Bool {
        caption.utf16.count <= 40_000 && (!caption.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty || !items.isEmpty)
            && (items.isEmpty || (supportsAttachments && items.allSatisfy { $0.state == .ready && $0.attachment != nil }))
            && persistenceAvailable
    }

    func configure(client: any AgentAPIClienting, storage: AgentDraftStorage?) {
        suspend()
        self.client = client; self.storage = storage
        prepared = [:]; loaders = [:]; capability = nil; error = nil; persistenceAvailable = true
        restoring = true
        defer { restoring = false; active = true }
        do {
            let saved = try storage?.read() ?? AgentDraftSnapshot()
            id = saved.id ?? UUID().uuidString
            caption = saved.caption
            items = saved.items.map { item in
                var item = item
                if item.state != .ready {
                    item.state = .failed
                    item.issue = .local("Upload interrupted. Retry with the same upload ID. If selection did not finish, remove and select the file again.", code: "upload_interrupted")
                }
                return item
            }
        } catch {
            caption = ""; items = []; persistenceAvailable = false
            self.error = "Could not restore attachment draft: \(error.localizedDescription)"
        }
    }
    func updateCapability(_ value: AgentAttachmentCapabilities?) {
        capability = value?.supported == true ? value : nil
        if capability == nil { suspend(); active = true }
    }
    func resume() { active = true; pump() }
    func suspend() {
        active = false; generation = UUID(); worker?.cancel(); worker = nil
        for index in items.indices where items[index].state == .loading || items[index].state == .uploading {
            items[index].state = .failed
            items[index].issue = .local("Transfer interrupted. Retry to continue.", code: "upload_interrupted")
        }
        save()
    }
    @discardableResult
    func add(filename: String, load: @escaping @Sendable () async throws -> AgentPreparedAttachment) -> String? {
        guard supportsAttachments, let capability else { error = AgentAttachmentIssue.unavailable.localizedDescription; return nil }
        guard items.count < capability.countLimit else { error = "You can attach up to \(capability.countLimit) files."; return nil }
        let item = AgentDraftItem(filename: filename)
        items.append(item); loaders[item.id] = load; error = nil
        save(); pump()
        return item.id
    }
    func remove(_ id: String) {
        // Restart the serial worker; completion from the old generation can never reinsert a removed item.
        let interrupted = items.filter { $0.state == .loading || $0.state == .uploading }.map(\.id)
        generation = UUID(); worker?.cancel(); worker = nil
        items.removeAll { $0.id == id }; loaders[id] = nil; prepared[id] = nil
        save()
        if persistenceAvailable { storage?.remove(id) }
        for other in interrupted where other != id {
            if let index = items.firstIndex(where: { $0.id == other }) { items[index].state = .loading }
        }
        pump()
    }
    func retry(_ id: String) {
        guard supportsAttachments else { error = AgentAttachmentIssue.unavailable.localizedDescription; return }
        guard let index = items.firstIndex(where: { $0.id == id }), items[index].state == .failed else { return }
        items[index].state = .loading; items[index].issue = nil; error = nil
        save(); pump()
    }
    func validateReady() throws -> [AgentAttachment] {
        guard items.isEmpty || supportsAttachments else { throw AgentAttachmentIssue.unavailable }
        guard items.allSatisfy({ $0.state == .ready && $0.attachment != nil }) else { throw AgentAttachmentIssue.local("Wait for every attachment to finish, or retry/remove failed files.") }
        if let capability { try capability.validate(ready) }
        return ready
    }
    @discardableResult func clearAfterSubmission() -> Bool {
        guard persistenceAvailable else { return false }
        let nextID = UUID().uuidString
        do { try storage?.write(AgentDraftSnapshot(id: nextID)) }
        catch {
            persistenceAvailable = false
            self.error = "Message is saved in the outbox, but its draft could not be cleared. Reconnect before retrying: \(error.localizedDescription)"
            return false
        }
        let ids = items.map(\.id)
        generation = UUID(); worker?.cancel(); worker = nil
        restoring = true; caption = ""; items = []; id = nextID; restoring = false
        prepared = [:]; loaders = [:]
        ids.forEach { storage?.remove($0) }
        return true
    }
    func save() {
        guard !restoring, persistenceAvailable else { return }
        do { try storage?.write(AgentDraftSnapshot(caption: caption, items: items, id: id)) }
        catch { persistenceAvailable = false; self.error = "Could not save attachment draft: \(error.localizedDescription)" }
    }
    private func pump() {
        guard active, supportsAttachments, let client, let limits = capability, worker == nil,
              let item = items.first(where: { $0.state == .loading }) else { return }
        let run = generation
        worker = Task { [weak self] in
            guard let self else { return }
            defer { if self.generation == run { self.worker = nil; self.pump() } }
            do {
                let file: AgentPreparedAttachment
                if let existing = self.prepared[item.id] { file = existing }
                else if let mime = item.mimeType, let storage = self.storage {
                    file = try AgentAttachmentPreparation.prepare(storage.bytes(item.id), filename: item.filename, mimeType: mime, limits: limits)
                } else if let loader = self.loaders[item.id] { file = try await loader() }
                else { throw AgentAttachmentIssue.local("Selection did not finish before the app closed. Remove this item and select it again.", code: "selection_unavailable") }
                try Task.checkCancellation()
                guard self.generation == run, let index = self.items.firstIndex(where: { $0.id == item.id }) else { return }
                let otherBytes = self.items.filter { $0.id != item.id }.reduce(0) { $0 + ($1.sizeBytes ?? 0) }
                guard otherBytes + file.data.count <= limits.totalLimit else { throw AgentAttachmentIssue.local("Attachments exceed the total message size limit.", code: "attachment_limit_exceeded") }
                self.prepared[item.id] = file
                try self.storage?.stage(item.id, file.data)
                self.items[index].filename = file.filename; self.items[index].mimeType = file.mimeType
                self.items[index].sizeBytes = file.data.count; self.items[index].state = .uploading
                self.save()
                guard self.persistenceAvailable else { throw AgentAttachmentIssue.local("Save failed. Reconnect before uploading.") }
                let attachment = try await client.agentUpload(file: file, uploadID: item.id)
                try attachment.validate()
                guard attachment.filename == file.filename, attachment.mimeType == file.mimeType,
                      attachment.sizeBytes == file.data.count, attachment.sha256 == file.sha256 else {
                    throw AgentAttachmentIssue.local("Upload receipt does not match the original file. Retry to confirm.", code: "invalid_upstream_response")
                }
                try Task.checkCancellation()
                guard self.generation == run, let index = self.items.firstIndex(where: { $0.id == item.id }) else { return }
                self.items[index].attachment = attachment; self.items[index].state = .ready; self.items[index].issue = nil
                self.save()
                self.loaders[item.id] = nil; self.prepared[item.id] = nil
                // Keep staged bytes until the message is durably handed to the outbox.
            } catch {
                guard self.generation == run, let index = self.items.firstIndex(where: { $0.id == item.id }) else { return }
                self.items[index].state = .failed; self.items[index].issue = .wrapping(error); self.save()
            }
        }
    }
}