import Foundation

/// A speech model the user can pick in Settings.
public struct ModelOption: Identifiable, Equatable, Sendable {
    public let id: String          // WhisperKit variant name
    public let title: String
    public let detail: String
    public let approxSizeMB: Int

    public init(id: String, title: String, detail: String, approxSizeMB: Int) {
        self.id = id
        self.title = title
        self.detail = detail
        self.approxSizeMB = approxSizeMB
    }
}

public enum ModelCatalog {
    public static let defaultVariant = "large-v3-v20240930_turbo"

    /// Sizes measured from the Hugging Face repository `argmaxinc/whisperkit-coreml`.
    public static let all: [ModelOption] = [
        ModelOption(id: defaultVariant, title: "Whisper Large v3 Turbo", detail: "Aanbevolen · nauwkeurigst", approxSizeMB: 1638),
        ModelOption(id: "large-v3-v20240930_turbo_632MB", title: "Whisper Large v3 Turbo (compact)", detail: "Kleiner, iets minder nauwkeurig", approxSizeMB: 646),
        ModelOption(id: "small", title: "Whisper Small", detail: "Snel en licht, minder nauwkeurig", approxSizeMB: 486),
    ]
}

/// Where models live on disk: `<base>/models/argmaxinc/whisperkit-coreml/openai_whisper-<variant>`.
public struct ModelStorage: Sendable {
    public static let repo = "argmaxinc/whisperkit-coreml"
    public static let requiredModels = ["MelSpectrogram", "AudioEncoder", "TextDecoder"]

    public let base: URL

    public init(base: URL = ModelStorage.defaultBase) {
        self.base = base
    }

    /// Application Support, so models are neither in ~/Documents (privacy prompt, iCloud) nor re-fetched.
    public static var defaultBase: URL {
        FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0]
            .appending(path: "Sst/models", directoryHint: .isDirectory)
    }

    var repoDirectory: URL {
        base.appending(path: "models/\(Self.repo)", directoryHint: .isDirectory)
    }

    private func folders(for variant: String) -> [URL] {
        let names = (try? FileManager.default.contentsOfDirectory(atPath: repoDirectory.path)) ?? []
        return names
            .filter { $0 == variant || $0 == "openai_whisper-\(variant)" }
            .map { repoDirectory.appending(path: $0, directoryHint: .isDirectory) }
    }

    /// The folder of a completely downloaded variant, or nil.
    public func folder(for variant: String) -> URL? {
        folders(for: variant).first { folder in
            Self.requiredModels.allSatisfy {
                FileManager.default.fileExists(atPath: folder.appending(path: "\($0).mlmodelc/coremldata.bin").path)
            }
        }
    }

    public func isInstalled(_ variant: String) -> Bool {
        folder(for: variant) != nil
    }

    public func delete(_ variant: String) throws {
        for folder in folders(for: variant) { try FileManager.default.removeItem(at: folder) }
    }
}