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

import Foundation

/// One HTTP transfer of one file into a `.partial` on disk, delegate-backed
/// for chunked throughput. Supports Range resume (server 206), transparent
/// full restarts (server 200), cancellation, and strips the Authorization
/// header when redirected off huggingface.co to the signed CDN.
final class FileTransfer: NSObject, URLSessionDataDelegate, @unchecked Sendable {
    enum TransferError: LocalizedError {
        case badStatus(Int)

        var errorDescription: String? {
            switch self {
            case .badStatus(let code): "Server returned HTTP \(code)."
            }
        }
    }

    private let partial: URL
    /// Cumulative bytes on disk, throttled (~10 Hz).
    private let onProgress: @Sendable (Int64) -> Void

    private var handle: FileHandle?
    private var received: Int64 = 0
    private var status = 0
    private var rejectedByStatus = false
    private var continuation: CheckedContinuation<Int64, Error>?
    private var lastReport = ContinuousClock.now

    init(partial: URL, onProgress: @escaping @Sendable (Int64) -> Void) {
        self.partial = partial
        self.onProgress = onProgress
    }

    /// Runs the transfer starting at `offset` (0 = fresh). Returns final byte
    /// count on disk. Throws CancellationError when the task is cancelled.
    func run(request: URLRequest, offset: Int64) async throws -> Int64 {
        received = offset
        status = 0
        rejectedByStatus = false

        let config = URLSessionConfiguration.default
        config.timeoutIntervalForRequest = 60
        config.networkServiceType = .responsiveData
        let session = URLSession(configuration: config, delegate: self, delegateQueue: nil)
        defer { session.finishTasksAndInvalidate() }

        var request = request
        if offset > 0 {
            request.setValue("bytes=\(offset)-", forHTTPHeaderField: "Range")
        }
        let task = session.dataTask(with: request)

        return try await withTaskCancellationHandler {
            try await withCheckedThrowingContinuation { c in
                continuation = c
                task.resume()
            }
        } onCancel: {
            task.cancel()
        }
    }

    // MARK: - URLSessionDataDelegate

    func urlSession(
        _ session: URLSession, task: URLSessionTask,
        willPerformHTTPRedirection response: HTTPURLResponse,
        newRequest request: URLRequest,
        completionHandler: @escaping (URLRequest?) -> Void
    ) {
        var request = request
        if request.url?.host != task.originalRequest?.url?.host {
            request.setValue(nil, forHTTPHeaderField: "Authorization")
        }
        completionHandler(request)
    }

    func urlSession(
        _ session: URLSession, dataTask: URLSessionDataTask,
        didReceive response: URLResponse,
        completionHandler: @escaping (URLSession.ResponseDisposition) -> Void
    ) {
        guard let http = response as? HTTPURLResponse else {
            rejectedByStatus = true
            completionHandler(.cancel)
            return
        }
        status = http.statusCode
        switch http.statusCode {
        case 200:
            // Fresh body (or the server ignored our Range): restart the file.
            received = 0
            FileManager.default.createFile(atPath: partial.path, contents: nil)
            handle = try? FileHandle(forWritingTo: partial)
            completionHandler(handle == nil ? .cancel : .allow)
        case 206:
            handle = try? FileHandle(forWritingTo: partial)
            _ = try? handle?.seekToEnd()
            completionHandler(handle == nil ? .cancel : .allow)
        default:
            rejectedByStatus = true
            completionHandler(.cancel)
        }
    }

    func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
        guard let handle else { return }
        do {
            try handle.write(contentsOf: data)
            received += Int64(data.count)
        } catch {
            dataTask.cancel()
            return
        }
        let now = ContinuousClock.now
        if lastReport.duration(to: now) > .milliseconds(100) {
            lastReport = now
            onProgress(received)
        }
    }

    func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
        try? handle?.close()
        handle = nil
        onProgress(received)
        let c = continuation
        continuation = nil
        if rejectedByStatus {
            c?.resume(throwing: TransferError.badStatus(status))
        } else if let error {
            if (error as? URLError)?.code == .cancelled {
                c?.resume(throwing: CancellationError())
            } else {
                c?.resume(throwing: error)
            }
        } else {
            c?.resume(returning: received)
        }
    }
}