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

import Foundation

/// One JSON-lines event from a PyBridge helper script.
/// All scripts emit `{"event": "...", ...}` objects, one per line
/// (protocol defined in `PyBridge/scripts/` — docs/TRAINING-RESEARCH.md §5).
struct PythonEvent: @unchecked Sendable {
    let event: String
    /// JSONSerialization plist values (NSString/NSNumber/…) — value types in
    /// practice, hence the unchecked-Sendable annotation.
    let payload: [String: Any]

    func double(_ key: String) -> Double? {
        (payload[key] as? Double) ?? (payload[key] as? Int).map(Double.init)
    }
    func int(_ key: String) -> Int? {
        (payload[key] as? Int) ?? (payload[key] as? Double).map(Int.init)
    }
    func string(_ key: String) -> String? { payload[key] as? String }
}

enum PythonRunnerError: LocalizedError {
    case scriptNotFound(String)
    case environmentNotProvisioned
    case processFailed(exitCode: Int32, stderr: String)

    var errorDescription: String? {
        switch self {
        case .scriptNotFound(let name):
            "Bundled Python script missing: \(name)"
        case .environmentNotProvisioned:
            "The Python environment is not set up yet. It is provisioned automatically on first use."
        case .processFailed(let code, let stderr):
            "Python pipeline failed (exit \(code)). \(stderr.suffix(500))"
        }
    }
}

/// All Python interaction goes through here (charter engineering standard):
/// spawns the pinned venv's interpreter on a bundled script and streams the
/// JSON-lines protocol back. Cancellation terminates the process.
actor PythonRunner {

    static let shared = PythonRunner()

    /// Resolve a bundled helper script by name.
    static func scriptURL(named name: String) throws -> URL {
        guard
            let url = Bundle.module.url(
                forResource: name, withExtension: "py", subdirectory: "scripts")
        else {
            throw PythonRunnerError.scriptNotFound("\(name).py")
        }
        return url
    }

    /// Run a script and stream its events. The stream throws on nonzero exit
    /// (unless an `error` event already surfaced the failure) and finishes
    /// after the process exits. Terminating the stream kills the process.
    func stream(script: String, arguments: [String]) async throws
        -> AsyncThrowingStream<PythonEvent, Error>
    {
        let env = PythonEnvironment.shared
        if await !env.isProvisioned {
            try await env.provision()
        }
        let python = await env.pythonURL
        let scriptURL = try Self.scriptURL(named: script)

        let process = Process()
        process.executableURL = python
        process.arguments = ["-u", scriptURL.path] + arguments
        // Isolated, reproducible interpreter environment. Homebrew paths are
        // included for external tools some pipelines shell out to (ffmpeg
        // for mlx-whisper audio decoding).
        process.environment = [
            "PATH": "/opt/homebrew/bin:/usr/local/bin:/usr/bin:/bin",
            "HOME": NSHomeDirectory(),
            "PYTHONUNBUFFERED": "1",
        ]

        let stdout = Pipe()
        let stderr = Pipe()
        process.standardOutput = stdout
        process.standardError = stderr

        return AsyncThrowingStream { continuation in
            let stderrBuffer = LockedBuffer()

            stderr.fileHandleForReading.readabilityHandler = { handle in
                stderrBuffer.append(handle.availableData)
            }

            let lineBuffer = LockedBuffer()
            stdout.fileHandleForReading.readabilityHandler = { handle in
                let data = handle.availableData
                guard !data.isEmpty else { return }
                for line in lineBuffer.appendAndExtractLines(data) {
                    if let event = Self.decode(line: line) {
                        continuation.yield(event)
                    }
                }
            }

            process.terminationHandler = { proc in
                stdout.fileHandleForReading.readabilityHandler = nil
                stderr.fileHandleForReading.readabilityHandler = nil
                // Flush any trailing output.
                let rest = stdout.fileHandleForReading.readDataToEndOfFile()
                for line in lineBuffer.appendAndExtractLines(rest, flush: true) {
                    if let event = Self.decode(line: line) {
                        continuation.yield(event)
                    }
                }
                stderrBuffer.append(stderr.fileHandleForReading.readDataToEndOfFile())
                if proc.terminationStatus == 0 {
                    continuation.finish()
                } else {
                    continuation.finish(
                        throwing: PythonRunnerError.processFailed(
                            exitCode: proc.terminationStatus,
                            stderr: stderrBuffer.string))
                }
            }

            continuation.onTermination = { termination in
                if case .cancelled = termination, process.isRunning {
                    process.terminate()
                }
            }

            do {
                try process.run()
            } catch {
                continuation.finish(throwing: error)
            }
        }
    }

    private static func decode(line: Data) -> PythonEvent? {
        guard
            let object = try? JSONSerialization.jsonObject(with: line) as? [String: Any],
            let event = object["event"] as? String
        else { return nil }
        return PythonEvent(event: event, payload: object)
    }
}

/// Thread-safe byte buffer for pipe readability handlers (they run on a
/// private DispatchQueue while the stream lives elsewhere).
private final class LockedBuffer: @unchecked Sendable {
    private let lock = NSLock()
    private var data = Data()

    func append(_ chunk: Data) {
        lock.lock()
        defer { lock.unlock() }
        data.append(chunk)
    }

    /// Append then split out complete `\n`-terminated lines (keeps the tail).
    func appendAndExtractLines(_ chunk: Data, flush: Bool = false) -> [Data] {
        lock.lock()
        defer { lock.unlock() }
        data.append(chunk)
        var lines: [Data] = []
        while let newline = data.firstIndex(of: UInt8(ascii: "\n")) {
            lines.append(data.subdata(in: data.startIndex..<newline))
            data.removeSubrange(data.startIndex...newline)
        }
        if flush, !data.isEmpty {
            lines.append(data)
            data.removeAll()
        }
        return lines
    }

    var string: String {
        lock.lock()
        defer { lock.unlock() }
        return String(data: data, encoding: .utf8) ?? ""
    }
}