import json
import pandas as pd
from transformers import TrainerCallback


class ContinuousLogCallback(TrainerCallback):
    def __init__(self, json_path, csv_path):
        self.json_path = json_path
        self.csv_path = csv_path
        self.logs = []

    def save_logs(self):
        with open(self.json_path, "w", encoding="utf-8") as f:
            json.dump(self.logs, f, ensure_ascii=False, indent=2)

        pd.DataFrame(self.logs).to_csv(self.csv_path, index=False)

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is None:
            return

        item = {
            "step": int(state.global_step),
            "epoch": float(state.epoch) if state.epoch is not None else None,
            "event": "log"
        }

        for key, value in logs.items():
            if isinstance(value, (int, float)):
                item[key] = float(value)

        self.logs.append(item)
        self.save_logs()

    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
        if metrics is None:
            return

        item = {
            "step": int(state.global_step),
            "epoch": float(state.epoch) if state.epoch is not None else None,
            "event": "evaluate"
        }

        for key, value in metrics.items():
            if isinstance(value, (int, float)):
                item[key] = float(value)

        self.logs.append(item)
        self.save_logs()