import json
from pathlib import Path

import pandas as pd
from transformers import TrainerCallback


class JsonlLogCallback(TrainerCallback):
    def __init__(self, path):
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs:
            with self.path.open("a") as file:
                file.write(json.dumps({"step": state.global_step, **logs}, default=float) + "\n")


def aggregate_train_logs(log_path, step_min=None, step_max=None):
    path = Path(log_path)
    rows = [json.loads(line) for line in path.read_text().splitlines()] if path.exists() else []
    if not rows:
        return {}
    frame = pd.DataFrame(rows)
    if "step" in frame:
        if step_min is not None:
            frame = frame[frame["step"] > step_min]
        if step_max is not None:
            frame = frame[frame["step"] <= step_max]
    if frame.empty:
        return {}

    def find(keys, last=False):
        for key in keys:
            if key in frame and frame[key].notna().any():
                values = pd.to_numeric(frame[key], errors="coerce").dropna()
                if len(values):
                    return float(values.iloc[-1] if last else values.mean())
        return None

    return {
        "train_reward_mean": find(["reward", "rewards/countdown_reward/mean"]),
        "train_reward_std": find(["reward_std"]),
        "kl_mean": find(["kl"]),
        "entropy_mean": find(["entropy"]),
        "avg_completion_length": find(["completions/mean_length"]),
        "completion_length_clip_ratio": find(["completions/clipped_ratio"]),
        "grad_norm": find(["grad_norm"]),
        "last_loss": find(["loss"], last=True),
        "end_learning_rate": find(["learning_rate"], last=True),
    }