"""Utilities for recording hash compressor training metrics."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Mapping


def write_training_metrics(
    path: str | Path,
    *,
    epoch: int,
    step: int,
    lr: float,
    components: Mapping[str, float],
) -> None:
    """Append one training metrics row as JSON Lines.

    The JSONL format keeps training logging cheap and easy to resume: every
    training step is an independent row that plotting scripts can stream later.
    """
    metrics_path = Path(path)
    metrics_path.parent.mkdir(parents=True, exist_ok=True)

    row = {
        "epoch": int(epoch),
        "step": int(step),
        "lr": float(lr),
        "total": float(components["total"]),
        "contrastive": float(components["contrastive"]),
        "distill": float(components["distill"]),
        "quantization": float(components["quantization"]),
    }

    with metrics_path.open("a", encoding="utf-8") as f:
        f.write(json.dumps(row, ensure_ascii=False) + "\n")