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

import Foundation

/// Manages the isolated, pinned Python environment in
/// `~/Library/Application Support/ZyquoMLX/py/venv` (docs/BUILD.md §3.4).
/// Provisioned with `uv` on first use; never touches system Python; lives
/// outside the signed bundle (notarization posture, docs/BUILD.md §4).
actor PythonEnvironment {

    static let shared = PythonEnvironment()

    /// Exact pins per docs/MLX-RESEARCH.md §6 (core set; mlx-vlm/whisper/audio
    /// install on demand when those capabilities are first used).
    static let pythonVersion = "3.12"
    static let corePins = ["mlx-lm==0.31.3", "pyyaml==6.0.3"]

    enum EnvironmentError: LocalizedError {
        case uvNotFound
        case provisioningFailed(String)

        var errorDescription: String? {
            switch self {
            case .uvNotFound:
                "The 'uv' tool is required to set up the Python environment. Install it with: brew install uv"
            case .provisioningFailed(let detail):
                "Python environment setup failed: \(detail)"
            }
        }
    }

    var venvURL: URL {
        PersistenceService.pythonDirectory.appendingPathComponent("venv", isDirectory: true)
    }

    var pythonURL: URL {
        venvURL.appendingPathComponent("bin/python")
    }

    private var markerURL: URL {
        PersistenceService.pythonDirectory.appendingPathComponent("provisioned.json")
    }

    var isProvisioned: Bool {
        FileManager.default.fileExists(atPath: pythonURL.path)
            && FileManager.default.fileExists(atPath: markerURL.path)
    }

    /// Locate `uv` in the usual installation paths (the app has no shell PATH).
    static func findUV() -> URL? {
        let candidates = [
            "\(NSHomeDirectory())/.local/bin/uv",
            "/opt/homebrew/bin/uv",
            "/usr/local/bin/uv",
        ]
        return candidates.first { FileManager.default.fileExists(atPath: $0) }
            .map { URL(fileURLWithPath: $0) }
    }

    /// Create the venv and install core pins. Idempotent: a pre-existing
    /// healthy venv (verified by importing the pinned packages) is adopted;
    /// a broken one is recreated in place (`uv venv --clear`).
    func provision(progress: @Sendable (String) -> Void = { _ in }) async throws {
        if isProvisioned { return }
        guard let uv = Self.findUV() else { throw EnvironmentError.uvNotFound }

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

        if FileManager.default.fileExists(atPath: pythonURL.path), verifyImports() {
            progress("Adopting existing Python environment.")
            try writeMarker()
            return
        }

        progress("Installing Python \(Self.pythonVersion)…")
        try runUV(uv, ["python", "install", Self.pythonVersion])

        progress("Creating environment…")
        try runUV(uv, ["venv", venvURL.path, "--clear", "--python", Self.pythonVersion])

        progress("Installing MLX packages…")
        try runUV(uv, ["pip", "install", "--python", pythonURL.path] + Self.corePins)

        try writeMarker()
        progress("Python environment ready.")
    }

    /// True when the pinned core packages import cleanly.
    private func verifyImports() -> Bool {
        let process = Process()
        process.executableURL = pythonURL
        process.arguments = ["-c", "import mlx_lm, yaml"]
        process.standardOutput = Pipe()
        process.standardError = Pipe()
        do {
            try process.run()
            process.waitUntilExit()
            return process.terminationStatus == 0
        } catch {
            return false
        }
    }

    private func writeMarker() throws {
        let marker: [String: String] = [
            "python": Self.pythonVersion,
            "pins": Self.corePins.joined(separator: " "),
            "provisionedAt": ISO8601DateFormatter().string(from: .now),
        ]
        try PersistenceService.saveJSON(marker, to: markerURL)
    }

    /// Install additional pinned packages (e.g. mlx-vlm when VLM Python
    /// pipelines are first needed).
    func install(pins: [String]) throws {
        guard let uv = Self.findUV() else { throw EnvironmentError.uvNotFound }
        try runUV(uv, ["pip", "install", "--python", pythonURL.path] + pins)
    }

    /// Delete and re-provision (Settings › Python Environment › Repair).
    func repair(progress: @Sendable (String) -> Void = { _ in }) async throws {
        try? FileManager.default.removeItem(at: venvURL)
        try? FileManager.default.removeItem(at: markerURL)
        try await provision(progress: progress)
    }

    private func runUV(_ uv: URL, _ arguments: [String]) throws {
        let process = Process()
        process.executableURL = uv
        process.arguments = arguments
        let errPipe = Pipe()
        process.standardError = errPipe
        process.standardOutput = Pipe()
        try process.run()
        process.waitUntilExit()
        if process.terminationStatus != 0 {
            let data = errPipe.fileHandleForReading.readDataToEndOfFile()
            throw EnvironmentError.provisioningFailed(
                String(data: data, encoding: .utf8) ?? "exit \(process.terminationStatus)")
        }
    }
}