import Foundation
import SstCore
import WhisperKit

public final class WhisperKitTranscriber: Transcriber {
    public enum TranscriberError: Error, LocalizedError {
        case notLoaded

        public var errorDescription: String? { "Spraakmodel is nog niet geladen." }
    }

    private let modelName: String
    private var pipeline: WhisperKit?

    /// Download progress (0...1) during the one-time model download. May be called from any thread.
    public var onProgress: ((Double) -> Void)?

    public init(model: String = "large-v3-v20240930_turbo") {
        self.modelName = model
    }

    /// Hugging Face hub root; WhisperKit stores models in `<root>/models/<repo>/<variant>`.
    /// Kept in Application Support so it is neither in ~/Documents (TCC prompt, iCloud) nor re-fetched.
    private static var storageBase: URL {
        FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0]
            .appending(path: "Sst/models", directoryHint: .isDirectory)
    }

    private static let repo = "argmaxinc/whisperkit-coreml"

    /// A previously downloaded variant folder, if it contains the compiled Core ML models.
    private func localModelFolder(base: URL) -> URL? {
        let repoDir = base.appending(path: "models/\(Self.repo)", directoryHint: .isDirectory)
        let folders = (try? FileManager.default.contentsOfDirectory(atPath: repoDir.path)) ?? []
        return folders
            .filter { $0.hasSuffix(modelName) }
            .map { repoDir.appending(path: $0, directoryHint: .isDirectory) }
            .first { folder in
                ["MelSpectrogram", "AudioEncoder", "TextDecoder"].allSatisfy {
                    FileManager.default.fileExists(
                        atPath: folder.appending(path: "\($0).mlmodelc/coremldata.bin").path)
                }
            }
    }

    private func loadPipeline(from folder: URL, base: URL) async throws -> WhisperKit {
        try await WhisperKit(WhisperKitConfig(
            downloadBase: base,
            modelFolder: folder.path,
            tokenizerFolder: base,
            prewarm: true,
            load: true,
            download: false
        ))
    }

    /// Loads the model into memory. Downloads it (≈1.5 GB) only if it is not stored locally yet;
    /// once downloaded, launching needs no network.
    public func load() async throws {
        guard pipeline == nil else { return }
        let base = Self.storageBase
        try FileManager.default.createDirectory(at: base, withIntermediateDirectories: true)

        if let folder = localModelFolder(base: base),
           let loaded = try? await loadPipeline(from: folder, base: base) {
            pipeline = loaded
            return
        }
        let folder = try await WhisperKit.download(
            variant: modelName,
            downloadBase: base,
            from: Self.repo
        ) { [onProgress] progress in
            onProgress?(progress.fractionCompleted)
        }
        pipeline = try await loadPipeline(from: folder, base: base)
    }

    public func transcribe(_ samples: [Float]) async throws -> String {
        guard let pipeline else { throw TranscriberError.notLoaded }
        // language nil + detectLanguage: Whisper picks NL or EN per utterance.
        let options = DecodingOptions(
            task: .transcribe,
            language: nil,
            detectLanguage: true,
            skipSpecialTokens: true,
            withoutTimestamps: true
        )
        let results = try await pipeline.transcribe(audioArray: samples, decodeOptions: options)
        return results
            .map(\.text)
            .joined(separator: " ")
            .trimmingCharacters(in: .whitespacesAndNewlines)
    }
}