Caps-iOS/Caps/Data/ImageCache.swift

import Foundation
import UIKit

final class ImageCache: ObservableObject {

    let folder: URL

    let server: URL

    let thumbnailSize: CGFloat

    private let fm: FileManager = .default

    private let session: URLSession = .shared

    private let thumbnailQuality: CGFloat = 0.7
    
    private let imageQuality: CGFloat = 0.3
    
    private let imageSize: CGSize = .init(width: 360, height: 360)
    
    @Published
    private(set) var imageCount = 0

    init(folder: URL, server: URL, thumbnailSize: CGFloat) throws {
        self.folder = folder
        self.server = server
        self.thumbnailSize = thumbnailSize * UIScreen.main.scale

        if !fm.fileExists(atPath: folder.path) {
            try fm.createDirectory(at: folder, withIntermediateDirectories: true)
        }
        self.imageCount = countLocalImages()
    }

    private func localImageUrl(_ image: CapImage) -> URL {
        folder.appendingPathComponent(String(format: "%04d-%02d.jpg", image.cap, image.version))
    }

    private func remoteImageUrl(_ image: CapImage) -> URL {
        server.appendingPathComponent(String(format: "images/%04d/%04d-%02d.jpg", image.cap, image.cap, image.version))
    }
    
    @discardableResult
    func save(_ image: UIImage, for cap: CapImage) -> Bool {
        guard let data = image.resize(to: imageSize).jpegData(compressionQuality: imageQuality) else {
            return false
        }
        let localUrl = localImageUrl(cap)
        guard removePossibleFile(localUrl) else {
            return false
        }
        do {
            try data.write(to: localUrl)
            return true
        } catch {
            print("failed to save image \(localUrl.lastPathComponent): \(error)")
            return false
        }
    }
    
    func availableImageUrl(_ image: CapImage) -> URL? {
        let url = localImageUrl(image)
        guard exists(url) else {
            return nil
        }
        return url
    }

    func image(_ image: CapImage, completion: @escaping (UIImage?) -> ()) {
        Task {
            let image = await self.image(image)
            completion(image)
        }
    }

    func image(_ image: CapImage, download: Bool = true) async -> UIImage? {
        if let localUrl = existingLocalImageUrl(image) {
            return UIImage(at: localUrl)
        }
        guard download else {
            return nil
        }
        guard let downloadedImageUrl = await loadRemoteImage(image) else {
            return nil
        }
        guard saveImage(image, at: downloadedImageUrl) else {
            return UIImage(at: downloadedImageUrl)
        }
        let localUrl = localImageUrl(image)
        return UIImage(at: localUrl)
    }

    func cachedImage(_ image: CapImage) -> UIImage? {
        guard let localUrl = existingLocalImageUrl(image) else {
            return nil
        }
        return UIImage(at: localUrl)
    }

    @discardableResult
    func refreshImage(_ image: CapImage) async -> Bool {
        guard let downloadedImageUrl = await loadRemoteImage(image) else {
            return false
        }
        return saveImage(image, at: downloadedImageUrl)
    }

    @discardableResult
    func removeCachedImages(for cap: Int) -> Bool {
        let prefix = String(format: "%04d-", cap)
        let files: [URL]
        do {
            files = try fm.contentsOfDirectory(atPath: folder.path)
                .filter { $0.hasPrefix(prefix) }
                .map { folder.appendingPathComponent($0) }
        } catch {
            log("Failed to get image cache file list: \(error)")
            return false
        }
        var isSuccessful = true
        for url in files {
            do {
                try fm.removeItem(at: url)
                DispatchQueue.main.async {
                    self.imageCount -= 1
                }
            } catch {
                isSuccessful = false
                log("Failed to remove image \(url.lastPathComponent) from cache: \(error)")
            }
        }
        return isSuccessful
    }

    private func loadRemoteImage(_ image: CapImage) async -> URL? {
        let remoteURL = remoteImageUrl(image)
        return await loadRemoteImage(at: remoteURL)
    }

    private func loadRemoteImage(at url: URL) async -> URL? {
        let tempUrl: URL
        let response: URLResponse
        do {
            (tempUrl, response) = try await session.download(from: url)
        } catch {
            print("Failed to download image \(url.lastPathComponent): \(error)")
            return nil
        }
        guard let httpResponse = response as? HTTPURLResponse else {
            print("Failed to download image \(url.lastPathComponent): Not a HTTP response: \(response)")
            return nil
        }
        guard httpResponse.statusCode == 200 else {
            print("Failed to download image \(url.path): Response \(httpResponse.statusCode)")
            return nil
        }
        return tempUrl
    }

    private func saveImage(_ image: CapImage, at tempUrl: URL) -> Bool {
        let localUrl = localImageUrl(image)
        let isOverwrite = exists(localUrl)
        guard removePossibleFile(localUrl) else {
            return false
        }
        do {
            try fm.moveItem(at: tempUrl, to: localUrl)
            if !isOverwrite {
                DispatchQueue.main.async {
                    self.imageCount += 1
                }
            }
            return true
        } catch {
            print("Failed to save image \(localUrl.lastPathComponent): \(error)")
            return false
        }
    }

    private func existingLocalImageUrl(_ image: CapImage) -> URL? {
        let localFile = localImageUrl(image)
        guard exists(localFile) else {
            return nil
        }
        return localFile
    }

    private func exists(_ url: URL) -> Bool {
        fm.fileExists(atPath: url.path)
    }

    @discardableResult
    private func removePossibleFile(_ file: URL) -> Bool {
        guard exists(file) else {
            return true
        }
        return remove(file)
    }

    @discardableResult
    private func remove(_ url: URL) -> Bool {
        do {
            try fm.removeItem(at: url)
            return true
        } catch {
            print("Failed to remove \(url.lastPathComponent): \(error)")
            return false
        }
    }

    // MARK: Thumbnails

    private func localThumbnailUrl(cap: Int) -> URL {
        folder.appendingPathComponent(String(format: "%04d.jpg", cap))
    }

    func thumbnail(for image: CapImage, download: Bool = true) async -> UIImage? {
        let localUrl = localThumbnailUrl(cap: image.cap)
        if exists(localUrl) {
            return UIImage(at: localUrl)
        }

        guard let mainImage = await self.image(image, download: download) else {
            return nil
        }
        let thumbnail = await createThumbnail(mainImage)
        save(thumbnail: thumbnail, for: image.cap)
        return thumbnail
    }

    func cachedThumbnail(for image: CapImage) -> UIImage? {
        let localUrl = localThumbnailUrl(cap: image.cap)
        guard exists(localUrl) else {
            return nil
        }
        return UIImage(at: localUrl)
    }

    @discardableResult
    func createThumbnail(for image: CapImage, download: Bool = false) async -> Bool {
        await thumbnail(for: image, download: download) != nil
    }

    private func createThumbnail(_ image: UIImage) async -> UIImage {
        let size = thumbnailSize
        return await withCheckedContinuation { continuation in
            DispatchQueue.global(qos: .background).async {
                let small = image.resize(to: CGSize(width: size, height: size))
                continuation.resume(returning: small)
            }
        }
    }

    @discardableResult
    private func save(thumbnail: UIImage, for cap: Int) -> Bool {
        guard let data = thumbnail.jpegData(compressionQuality: thumbnailQuality) else {
            print("Failed to get thumbnail JPEG data")
            return false
        }
        let localUrl = localThumbnailUrl(cap: cap)
        let isOverwrite = exists(localUrl)
        do {
            try data.write(to: localUrl)
            if !isOverwrite {
                DispatchQueue.main.async {
                    self.imageCount += 1
                }
            }
            return true
        } catch {
            print("Failed to save thumbnail \(cap): \(error)")
            return false
        }
    }
    
    private func countLocalImages() -> Int {
        cachedImages.count
    }
    
    private var cachedImages: [URL] {
        do {
            return try fm.contentsOfDirectory(at: folder, includingPropertiesForKeys: nil)
        } catch {
            log("Failed to get cached images: \(error)")
            return []
        }
    }
    
    private func imageId(from url: URL) -> CapImage? {
        let parts = url.deletingPathExtension().lastPathComponent.components(separatedBy: "-")
        guard parts.count == 2 else {
            log("File \(url.lastPathComponent) is not a cap image")
            return nil
        }
        guard let capId = Int(parts.first!),
              let version = Int(parts.last!) else {
            log("File \(url.lastPathComponent) is not a cap image")
            return nil
        }
        return .init(cap: capId, version: version)
    }

    func clearImageCache(keeping: Set<CapImage>) {
        let allImages = cachedImages
        for url in allImages {
            if let id = imageId(from: url), keeping.contains(id) {
                continue // Skip image
            }
            do {
                try fm.removeItem(at: url)
            } catch {
                log("Failed to delete cached image \(url.lastPathComponent): \(error)")
            }
        }
        let newCount = countLocalImages()
        log("Deleted \(allImages.count - newCount) of \(allImages.count) cached images, leaving \(newCount)")
        DispatchQueue.main.async {
            self.imageCount = newCount
        }
    }
}