//
//  ModelManager.swift
//  transcriptor
//
//  Downloads and locates qwen-asr model files — the one thing "one-click
//  install" deliberately leaves at runtime (see CLAUDE.md, "Goal": models
//  are multi-GB and don't belong baked into the app bundle, same reasoning
//  Ollama/LM Studio don't ship them inside their own app either).
//
//  File list and HuggingFace URL pattern mirror
//  transcriptor-ai/third_party/qwen-asr/download_model.sh exactly — keep in
//  sync by hand if that script's file list ever changes.
//

import Foundation

enum ModelChoice: String, CaseIterable, Identifiable {
    case small
    case large

    var id: String { rawValue }

    var displayName: String {
        switch self {
        case .small: "Qwen3-ASR 0.6B (faster, lower quality)"
        case .large: "Qwen3-ASR 1.7B (recommended)"
        }
    }

    fileprivate var modelID: String {
        switch self {
        case .small: "Qwen/Qwen3-ASR-0.6B"
        case .large: "Qwen/Qwen3-ASR-1.7B"
        }
    }

    var directoryName: String {
        switch self {
        case .small: "qwen3-asr-0.6b"
        case .large: "qwen3-asr-1.7b"
        }
    }

    fileprivate var files: [String] {
        switch self {
        case .small:
            ["config.json", "generation_config.json", "model.safetensors", "vocab.json", "merges.txt"]
        case .large:
            [
                "config.json", "generation_config.json", "model.safetensors.index.json",
                "model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors",
                "vocab.json", "merges.txt",
            ]
        }
    }
}

enum ModelManagerError: Error {
    case invalidDownloadedFile
}

@Observable
final class ModelManager: NSObject {
    private(set) var isDownloading = false
    private(set) var statusText = ""
    private(set) var fileProgress: Double = 0 // 0...1, current file only
    private(set) var currentFileIndex = 0
    private(set) var totalFiles = 0

    private struct DownloadContext {
        let destination: URL
        let completion: (Result<Void, Error>) -> Void
    }
    private var contexts: [Int: DownloadContext] = [:]
    // delegateQueue: .main so delegate callbacks land on the main actor
    // directly — this class's @Observable properties are read by SwiftUI,
    // which needs that.
    // @Observable's macro doesn't support `lazy var` (init-accessor
    // synthesis conflict) — this is an internal implementation detail
    // anyway, not UI-relevant state, so excluding it from observation
    // tracking is the right fix, not a workaround.
    @ObservationIgnored
    private lazy var session = URLSession(configuration: .default, delegate: self, delegateQueue: .main)

    static var modelsRootDirectory: URL {
        let appSupport = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0]
        return appSupport.appendingPathComponent("eu.sttlab.transcriptor/models", isDirectory: true)
    }

    func localDirectory(for choice: ModelChoice) -> URL {
        Self.modelsRootDirectory.appendingPathComponent(choice.directoryName, isDirectory: true)
    }

    func isDownloaded(_ choice: ModelChoice) -> Bool {
        let dir = localDirectory(for: choice)
        return choice.files.allSatisfy {
            FileManager.default.fileExists(atPath: dir.appendingPathComponent($0).path)
        }
    }

    /// Downloads whatever files are missing (already-present files are left
    /// alone, so an interrupted download resumes rather than restarting)
    /// and returns the model's local directory.
    func ensureModel(_ choice: ModelChoice) async throws -> URL {
        let dir = localDirectory(for: choice)
        if isDownloaded(choice) { return dir }

        try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)

        isDownloading = true
        defer { isDownloading = false }

        totalFiles = choice.files.count
        for (index, file) in choice.files.enumerated() {
            currentFileIndex = index + 1
            let destination = dir.appendingPathComponent(file)
            if FileManager.default.fileExists(atPath: destination.path) {
                continue
            }
            statusText = "Downloading \(file) (\(currentFileIndex)/\(totalFiles))..."
            fileProgress = 0
            let url = URL(string: "https://huggingface.co/\(choice.modelID)/resolve/main/\(file)")!
            try await downloadFile(from: url, to: destination)
        }
        statusText = "Model ready."
        return dir
    }

    private func downloadFile(from url: URL, to destination: URL) async throws {
        try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
            let task = session.downloadTask(with: url)
            contexts[task.taskIdentifier] = DownloadContext(destination: destination) { result in
                continuation.resume(with: result)
            }
            task.resume()
        }
    }
}

extension ModelManager: URLSessionDownloadDelegate {
    func urlSession(
        _ session: URLSession, downloadTask: URLSessionDownloadTask,
        didWriteData bytesWritten: Int64, totalBytesWritten: Int64, totalBytesExpectedToWrite: Int64
    ) {
        guard totalBytesExpectedToWrite > 0 else { return }
        fileProgress = Double(totalBytesWritten) / Double(totalBytesExpectedToWrite)
    }

    func urlSession(
        _ session: URLSession, downloadTask: URLSessionDownloadTask, didFinishDownloadingTo location: URL
    ) {
        // The file at `location` only exists for the duration of this
        // callback — the move must happen synchronously here, not be
        // deferred to later (e.g. inside the completion closure), or the
        // system may have already deleted it.
        guard let context = contexts.removeValue(forKey: downloadTask.taskIdentifier) else { return }
        do {
            if FileManager.default.fileExists(atPath: context.destination.path) {
                try FileManager.default.removeItem(at: context.destination)
            }
            try FileManager.default.moveItem(at: location, to: context.destination)
            context.completion(.success(()))
        } catch {
            context.completion(.failure(error))
        }
    }

    func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
        // A successful download is already resolved in
        // didFinishDownloadingTo (which removes the context) — this only
        // has work to do for a task that failed before reaching that point
        // (network error, etc).
        guard let error, let context = contexts.removeValue(forKey: task.taskIdentifier) else { return }
        context.completion(.failure(error))
    }
}