//
//  ConversionService.swift
//  Zyquo MLX
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//

import Foundation
import MLX
import MLXLLM
import MLXLMCommon

/// Progress for conversion/quantization/fusion jobs.
enum ConversionEvent: Sendable {
    case stage(String, fraction: Double?)
    case finished(outputDirectory: URL)
    case failed(message: String)
}

enum ConversionServiceError: LocalizedError {
    case alreadyQuantized
    case outputExists(URL)

    var errorDescription: String? {
        switch self {
        case .alreadyQuantized:
            "This model is already quantized — de-quantize first if you want a different recipe."
        case .outputExists(let url):
            "Output folder already exists: \(url.lastPathComponent). Choose another name or delete it."
        }
    }
}

/// Conversion, quantization, and adapter fusion (charter 3.C).
///
/// Execution split per docs/MLX-RESEARCH.md §2/§6:
/// - Quantizing a local safetensors LLM → **native Swift**
///   (`LLMModelFactory.convert`, affine/mxfp4).
/// - Fusing adapters to disk, HF→MLX conversion of arbitrary repos
///   (incl. `.bin`), de-quantize → **Python bridge** (`mlx_lm` pinned).
actor ConversionService {

    static let shared = ConversionService()

    // MARK: - Swift-native quantization

    /// Quantize a local, unquantized safetensors LLM into a new model
    /// directory in the library.
    func quantize(model: LocalModel, config: QuantConfig, outputName: String)
        -> AsyncStream<ConversionEvent>
    {
        let output = PersistenceService.modelsDirectory
            .appendingPathComponent(outputName, isDirectory: true)

        let (stream, continuation) = AsyncStream.makeStream(of: ConversionEvent.self)
        Task {
            do {
                guard model.quantization == nil else { throw ConversionServiceError.alreadyQuantized }
                guard !FileManager.default.fileExists(atPath: output.path) else {
                    throw ConversionServiceError.outputExists(output)
                }
                let options = ModelConversionOptions(
                    bits: config.bits,
                    groupSize: config.groupSize,
                    mode: config.mode == .mxfp4 ? .mxfp4 : .affine
                )
                _ = try await LLMModelFactory.shared.convert(
                    from: model.directory,
                    to: output,
                    options: options,
                    progressHandler: { progress in
                        continuation.yield(
                            .stage(progress.stage.rawValue, fraction: progress.fractionCompleted))
                    }
                )
                Memory.clearCache()
                continuation.yield(.finished(outputDirectory: output))
            } catch {
                continuation.yield(.failed(message: error.localizedDescription))
            }
            continuation.finish()
        }
        return stream
    }

    // MARK: - Python-bridge pipelines

    /// Fuse a trained adapter into its base model, producing a standalone
    /// MLX model directory in the library.
    func fuse(baseModel: LocalModel, adapterDirectory: URL, outputName: String, dequantize: Bool)
        async throws -> AsyncStream<ConversionEvent>
    {
        let output = PersistenceService.modelsDirectory
            .appendingPathComponent(outputName, isDirectory: true)
        guard !FileManager.default.fileExists(atPath: output.path) else {
            throw ConversionServiceError.outputExists(output)
        }

        var arguments = [
            "--model", baseModel.directory.path,
            "--adapter-path", adapterDirectory.path,
            "--save-path", output.path,
        ]
        if dequantize { arguments.append("--dequantize") }

        return try await pythonJob(script: "zyquo_fuse", arguments: arguments, output: output)
    }

    /// Convert an HF repo/directory (any weights format mlx-lm handles) to
    /// MLX, optionally quantized.
    func convert(hfPath: String, outputName: String, quantize: QuantConfig?)
        async throws -> AsyncStream<ConversionEvent>
    {
        let output = PersistenceService.modelsDirectory
            .appendingPathComponent(outputName, isDirectory: true)
        guard !FileManager.default.fileExists(atPath: output.path) else {
            throw ConversionServiceError.outputExists(output)
        }

        var arguments = ["--hf-path", hfPath, "--mlx-path", output.path]
        if let quantize {
            arguments += [
                "--quantize",
                "--q-bits", String(quantize.bits),
                "--q-group-size", String(quantize.groupSize),
                "--q-mode", quantize.mode.rawValue,
            ]
        }

        return try await pythonJob(script: "zyquo_convert", arguments: arguments, output: output)
    }

    private func pythonJob(script: String, arguments: [String], output: URL)
        async throws -> AsyncStream<ConversionEvent>
    {
        let pythonStream = try await PythonRunner.shared.stream(
            script: script, arguments: arguments)
        let (stream, continuation) = AsyncStream.makeStream(of: ConversionEvent.self)
        Task {
            do {
                for try await event in pythonStream {
                    switch event.event {
                    case "start":
                        continuation.yield(.stage(event.string("stage") ?? "working", fraction: nil))
                    case "done":
                        continuation.yield(.finished(outputDirectory: output))
                    case "error":
                        continuation.yield(.failed(message: event.string("message") ?? "unknown"))
                    default:
                        break
                    }
                }
            } catch {
                continuation.yield(.failed(message: error.localizedDescription))
            }
            continuation.finish()
        }
        return stream
    }
}