"""Real, on-device LoRA fine-tuning via Apple's MLX framework (no CUDA GPU needed —
runs on Apple Silicon). Pipeline: export approved dataset examples to MLX's chat JSONL
format -> mlx_lm.lora trains a real LoRA adapter -> mlx_lm.fuse merges it into the base
weights (dequantized, since MLX's GGUF exporter can't handle quantized weights). The
fused model then goes through the shared llama.cpp GGUF conversion + Ollama import
(see shared.py), identically to every other backend.
"""

import re
import sys
from pathlib import Path
from typing import Callable

from app.services.finetune.shared import patch_tokenizer_class, run_streaming

_TRAIN_ITER_RE = re.compile(r"Iter (\d+): Train loss ([\d.]+)")
_VAL_ITER_RE = re.compile(r"Iter (\d+): Val loss ([\d.]+)")


def run_lora_training(
    base_model: str, data_dir: Path, adapter_path: Path, config: dict, on_line: Callable[[str], None]
) -> None:
    cmd = [
        sys.executable,
        "-m",
        "mlx_lm",
        "lora",
        "--model",
        base_model,
        "--train",
        "--data",
        str(data_dir),
        "--adapter-path",
        str(adapter_path),
        "--iters",
        str(config["iters"]),
        "--batch-size",
        str(config["batch_size"]),
        "--learning-rate",
        str(config["learning_rate"]),
        "--num-layers",
        str(config["num_layers"]),
        "--steps-per-report",
        str(config.get("steps_per_report", 5)),
        "--steps-per-eval",
        str(config.get("steps_per_eval", 20)),
        "--val-batches",
        "1",
    ]
    run_streaming(cmd, on_line)


def run_fuse(base_model: str, adapter_path: Path, fused_dir: Path) -> None:
    cmd = [
        sys.executable,
        "-m",
        "mlx_lm",
        "fuse",
        "--model",
        base_model,
        "--adapter-path",
        str(adapter_path),
        "--save-path",
        str(fused_dir),
        "--dequantize",
    ]
    run_streaming(cmd, lambda _line: None)
    patch_tokenizer_class(fused_dir)


def parse_progress_line(line: str) -> dict | None:
    m = _TRAIN_ITER_RE.search(line)
    if m:
        return {"type": "train", "iter": int(m.group(1)), "loss": float(m.group(2))}
    m = _VAL_ITER_RE.search(line)
    if m:
        return {"type": "val", "iter": int(m.group(1)), "loss": float(m.group(2))}
    return None