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

import Foundation

/// Typed training events, decoded from the PyBridge JSON-lines protocol
/// (`zyquo_train.py`; payload fields are the stable mlx-lm callback dicts —
/// docs/TRAINING-RESEARCH.md §5.2).
enum TrainingEvent: Sendable {
    case started(model: String, iterations: Int)
    case metric(TrainingMetric)
    case checkpointSaved(fileName: String, iteration: Int)
    case finished
    case failed(message: String)
}

enum MetricsStream {

    /// Decode one Python event into a training event (nil = ignorable).
    static func decode(_ event: PythonEvent) -> TrainingEvent? {
        switch event.event {
        case "start":
            return .started(
                model: event.string("model") ?? "?",
                iterations: event.int("iters") ?? 0)

        case "train":
            return .metric(
                TrainingMetric(
                    iteration: event.int("iteration") ?? 0,
                    trainLoss: event.double("train_loss"),
                    valLoss: nil,
                    learningRate: event.double("learning_rate"),
                    iterationsPerSecond: event.double("iterations_per_second"),
                    tokensPerSecond: event.double("tokens_per_second"),
                    trainedTokens: event.int("trained_tokens"),
                    peakMemoryGB: event.double("peak_memory"),
                    timestamp: .now))

        case "val":
            return .metric(
                TrainingMetric(
                    iteration: event.int("iteration") ?? 0,
                    trainLoss: nil,
                    valLoss: event.double("val_loss"),
                    learningRate: nil,
                    iterationsPerSecond: nil,
                    tokensPerSecond: nil,
                    trainedTokens: nil,
                    peakMemoryGB: nil,
                    timestamp: .now))

        case "done":
            return .finished

        case "error":
            return .failed(message: event.string("message") ?? "unknown error")

        default:
            return nil
        }
    }
}