//
//  DownloadManager.swift
//  Zyquo Local
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//

import Foundation
import Observation

/// Queued, resumable model downloads.
///
/// - Files download to `<name>.partial` and are moved into place atomically;
///   a model is valid only when every file is present and size-verified.
/// - Resume uses HTTP Range from the partial file's size — surviving pause,
///   cancel-restart, and app relaunch (manifest + partials persist on disk).
/// - Max 2 concurrent file downloads per model.
@MainActor
@Observable
final class DownloadManager {
    private(set) var tasks: [DownloadTask] = []
    /// Live transfer speed per repoID, bytes/sec.
    private(set) var speeds: [String: Double] = [:]

    var hub: HubService
    private let store: ModelStore
    private var workers: [String: Task<Void, Never>] = [:]
    /// Speed window per repoID: last sample time and cumulative bytes.
    private var speedSamples: [String: (time: ContinuousClock.Instant, bytes: Int64)] = [:]

    /// Active + queued downloads (for the sidebar badge).
    var activeCount: Int {
        tasks.filter { $0.state == .downloading || $0.state == .queued }.count
    }

    var overallFraction: Double {
        let active = tasks.filter { $0.state == .downloading || $0.state == .paused }
        guard !active.isEmpty else { return 0 }
        return active.reduce(0.0) { $0 + $1.fractionCompleted } / Double(active.count)
    }

    init(hub: HubService, store: ModelStore) {
        self.hub = hub
        self.store = store
        restoreManifest()
    }

    enum DownloadError: LocalizedError {
        case insufficientDiskSpace(needed: Int64, available: Int64)
        case sizeMismatch(file: String)
        case interrupted

        var errorDescription: String? {
            switch self {
            case .insufficientDiskSpace(let needed, let available):
                "Not enough disk space: \(ByteCountFormatter.string(fromByteCount: needed, countStyle: .file)) needed, \(ByteCountFormatter.string(fromByteCount: available, countStyle: .file)) available."
            case .sizeMismatch(let file):
                "Downloaded file \(file) does not match its expected size. The download will retry from scratch for this file."
            case .interrupted:
                "The download was interrupted. It can be resumed from where it left off."
            }
        }
    }

    // MARK: - Public API

    /// Starts (or restarts) downloading all required files of a repo.
    func download(repoID: String) async {
        if let existing = task(for: repoID),
            existing.state == .downloading || existing.state == .queued
        {
            return
        }
        do {
            let (files, totalBytes) = try await hub.requiredFiles(of: repoID)
            try checkDiskSpace(needed: totalBytes)
            let dir = store.directory(for: repoID)
            try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)

            var progress: [DownloadTask.FileProgress] = files.map { file in
                let partial = dir.appendingPathComponent(file.path + ".partial")
                let final = dir.appendingPathComponent(file.path)
                let already: Int64
                let done: Bool
                if FileManager.default.fileExists(atPath: final.path) {
                    already = file.size ?? 0
                    done = true
                } else {
                    already = (try? FileManager.default.attributesOfItem(atPath: partial.path)[.size] as? Int64) ?? 0
                    done = false
                }
                return DownloadTask.FileProgress(
                    path: file.path,
                    totalBytes: file.size ?? 0,
                    receivedBytes: already,
                    sha256: file.lfs?.oid,
                    completed: done
                )
            }
            progress.sort { $0.totalBytes < $1.totalBytes }  // small config files first
            upsert(DownloadTask(
                repoID: repoID, state: .downloading, files: progress,
                errorDescription: nil, startedAt: Date()
            ))
            startWorker(repoID: repoID)
        } catch {
            upsert(DownloadTask(
                repoID: repoID, state: .failed, files: task(for: repoID)?.files ?? [],
                errorDescription: error.localizedDescription, startedAt: Date()
            ))
        }
    }

    func pause(repoID: String) {
        workers[repoID]?.cancel()
        workers[repoID] = nil
        mutate(repoID) { $0.state = .paused }
    }

    func resume(repoID: String) {
        guard var t = task(for: repoID), t.state == .paused || t.state == .failed else { return }
        t.state = .downloading
        t.errorDescription = nil
        upsert(t)
        startWorker(repoID: repoID)
    }

    /// Cancels and removes all partial data.
    func cancel(repoID: String) {
        workers[repoID]?.cancel()
        workers[repoID] = nil
        let dir = store.directory(for: repoID)
        if let t = task(for: repoID) {
            for f in t.files where !f.completed {
                try? FileManager.default.removeItem(at: dir.appendingPathComponent(f.path + ".partial"))
            }
        }
        // Remove the directory when nothing complete remains in it.
        if let contents = try? FileManager.default.contentsOfDirectory(atPath: dir.path), contents.isEmpty {
            try? FileManager.default.removeItem(at: dir)
        }
        tasks.removeAll { $0.repoID == repoID }
        speeds[repoID] = nil
        speedSamples[repoID] = nil
        persistManifest()
    }

    func task(for repoID: String) -> DownloadTask? {
        tasks.first { $0.repoID == repoID }
    }

    /// ETA in seconds based on current speed.
    func eta(for repoID: String) -> TimeInterval? {
        guard let t = task(for: repoID), let speed = speeds[t.repoID], speed > 1 else { return nil }
        return Double(t.totalBytes - t.receivedBytes) / speed
    }

    // MARK: - Worker

    private func startWorker(repoID: String) {
        workers[repoID]?.cancel()
        workers[repoID] = Task { [weak self] in
            await self?.runWorker(repoID: repoID)
        }
    }

    private func runWorker(repoID: String) async {
        let dir = store.directory(for: repoID)
        guard let snapshot = task(for: repoID) else { return }
        let pending = snapshot.files.filter { !$0.completed }

        do {
            try await withThrowingTaskGroup(of: Void.self) { group in
                var iterator = pending.makeIterator()
                var inFlight = 0
                func addNext(_ group: inout ThrowingTaskGroup<Void, Error>) {
                    guard let file = iterator.next() else { return }
                    inFlight += 1
                    group.addTask { [weak self] in
                        try await self?.downloadFile(repoID: repoID, file: file, directory: dir)
                    }
                }
                addNext(&group)
                addNext(&group)  // 2 concurrent file downloads max
                while inFlight > 0 {
                    try await group.next()
                    inFlight -= 1
                    addNext(&group)
                }
            }
            mutate(repoID) { $0.state = .verifying }
            try verify(repoID: repoID, directory: dir)
            mutate(repoID) { $0.state = .completed }
            speeds[repoID] = nil
            store.rescan()
        } catch is CancellationError {
            // pause() already set the state; on app quit partials just persist.
        } catch {
            if Task.isCancelled { return }
            mutate(repoID) {
                $0.state = .failed
                $0.errorDescription = error.localizedDescription
            }
            speeds[repoID] = nil
        }
        workers[repoID] = nil
    }

    /// Downloads one file with Range-resume into `<path>.partial`, then moves
    /// it into place. Retries transient network drops from the partial offset.
    /// Throws CancellationError on pause.
    private func downloadFile(repoID: String, file: DownloadTask.FileProgress, directory: URL) async throws {
        let partial = directory.appendingPathComponent(file.path + ".partial")
        let final = directory.appendingPathComponent(file.path)
        try FileManager.default.createDirectory(
            at: partial.deletingLastPathComponent(), withIntermediateDirectories: true)
        if !FileManager.default.fileExists(atPath: partial.path) {
            FileManager.default.createFile(atPath: partial.path, contents: nil)
        }

        var request = URLRequest(url: HubService.resolveURL(repoID: repoID, path: file.path))
        request.setValue("ZyquoLocal/1.0.0", forHTTPHeaderField: "User-Agent")
        if let token = hub.token, !token.isEmpty {
            request.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
        }

        let path = file.path
        let transfer = FileTransfer(partial: partial) { [weak self] received in
            Task { @MainActor [weak self] in
                self?.updateProgress(repoID: repoID, path: path, received: received)
            }
        }

        var attempts = 0
        while true {
            attempts += 1
            let offset: Int64 =
                (try? FileManager.default.attributesOfItem(atPath: partial.path)[.size] as? Int64) ?? 0
            do {
                _ = try await transfer.run(request: request, offset: offset)
                break
            } catch let error as FileTransfer.TransferError {
                switch error {
                case .badStatus(401), .badStatus(403):
                    throw HubService.HubError.gatedOrMissing(repoID)
                case .badStatus(416):
                    // Range beyond EOF: complete if the size matches, else restart.
                    let size = (try? FileManager.default.attributesOfItem(atPath: partial.path)[.size] as? Int64) ?? 0
                    if size == file.totalBytes { break }
                    try? FileManager.default.removeItem(at: partial)
                    FileManager.default.createFile(atPath: partial.path, contents: nil)
                    if attempts >= 3 { throw DownloadError.sizeMismatch(file: file.path) }
                    continue
                case .badStatus(let code):
                    throw HubService.HubError.http(code, repoID)
                }
            } catch is CancellationError {
                throw CancellationError()
            } catch {
                // Transient drop (connection lost, timeout…): retry with backoff
                // from whatever the partial already holds.
                guard attempts < 4 else { throw error }
                try await Task.sleep(for: .seconds(Double(attempts)))
                try Task.checkCancellation()
                continue
            }
        }
        try finishFile(partial: partial, final: final, repoID: repoID, path: file.path, expected: file.totalBytes)
    }

    /// Size check + atomic move into place.
    private func finishFile(partial: URL, final: URL, repoID: String, path: String, expected: Int64) throws {
        let size = (try? FileManager.default.attributesOfItem(atPath: partial.path)[.size] as? Int64) ?? -1
        guard expected == 0 || size == expected else {
            try? FileManager.default.removeItem(at: partial)
            updateProgress(repoID: repoID, path: path, received: 0)
            throw DownloadError.sizeMismatch(file: path)
        }
        if FileManager.default.fileExists(atPath: final.path) {
            try FileManager.default.removeItem(at: final)
        }
        try FileManager.default.moveItem(at: partial, to: final)
        mutate(repoID) { t in
            if let i = t.files.firstIndex(where: { $0.path == path }) {
                t.files[i].completed = true
                t.files[i].receivedBytes = t.files[i].totalBytes
            }
        }
    }

    /// Final integrity pass: every required file present with expected size.
    private func verify(repoID: String, directory: URL) throws {
        guard let t = task(for: repoID) else { return }
        for f in t.files {
            let url = directory.appendingPathComponent(f.path)
            let size = (try? FileManager.default.attributesOfItem(atPath: url.path)[.size] as? Int64) ?? -1
            guard f.totalBytes == 0 || size == f.totalBytes else {
                throw DownloadError.sizeMismatch(file: f.path)
            }
        }
    }

    private func checkDiskSpace(needed: Int64) throws {
        let values = try? store.modelsRoot.resourceValues(forKeys: [.volumeAvailableCapacityForImportantUsageKey])
        let available = values?.volumeAvailableCapacityForImportantUsage ?? 0
        // 5 % headroom so a download never fills the disk completely.
        if available < Int64(Double(needed) * 1.05) {
            throw DownloadError.insufficientDiskSpace(needed: needed, available: available)
        }
    }

    // MARK: - State & manifest

    private func upsert(_ t: DownloadTask) {
        if let i = tasks.firstIndex(where: { $0.repoID == t.repoID }) {
            tasks[i] = t
        } else {
            tasks.append(t)
        }
        persistManifest()
    }

    private func mutate(_ repoID: String, _ change: (inout DownloadTask) -> Void) {
        guard let i = tasks.firstIndex(where: { $0.repoID == repoID }) else { return }
        change(&tasks[i])
        persistManifest()
    }

    private func updateProgress(repoID: String, path: String, received: Int64) {
        guard let i = tasks.firstIndex(where: { $0.repoID == repoID }) else { return }
        if let j = tasks[i].files.firstIndex(where: { $0.path == path }) {
            tasks[i].files[j].receivedBytes = received
        }
        // Smoothed speed over a sliding window of the whole repo's bytes.
        let totalReceived = tasks[i].receivedBytes
        let now = ContinuousClock.now
        if let sample = speedSamples[repoID] {
            let d = sample.time.duration(to: now)
            let elapsed = Double(d.components.seconds) + Double(d.components.attoseconds) * 1e-18
            if elapsed >= 1.0 {
                let instant = Double(totalReceived - sample.bytes) / elapsed
                let previous = speeds[repoID] ?? instant
                speeds[repoID] = previous * 0.6 + instant * 0.4
                speedSamples[repoID] = (now, totalReceived)
            }
        } else {
            speedSamples[repoID] = (now, totalReceived)
        }
    }

    private var manifestURL: URL {
        PersistenceService.appSupportDirectory.appendingPathComponent("downloads.json")
    }

    private func persistManifest() {
        let persistable = tasks.filter { $0.state != .completed && $0.state != .cancelled }
        try? JSONEncoder().encode(persistable).write(to: manifestURL, options: .atomic)
    }

    /// Restores paused/interrupted downloads after relaunch.
    private func restoreManifest() {
        guard let data = try? Data(contentsOf: manifestURL),
            var restored = try? JSONDecoder().decode([DownloadTask].self, from: data)
        else { return }
        for i in restored.indices where restored[i].state == .downloading || restored[i].state == .queued {
            restored[i].state = .paused
        }
        tasks = restored
    }
}