"""
ml/callbacks.py - Stable-Baselines3 Custom Callbacks for Logging & Curriculum Advancement
"""
import os
import numpy as np
from stable_baselines3.common.callbacks import BaseCallback


class RewardLoggerCallback(BaseCallback):
    """Logs individual reward component averages to TensorBoard."""

    def __init__(self, verbose: int = 0):
        super().__init__(verbose)

    def _on_step(self) -> bool:
        # Pull component averages from the environment vector
        for env_idx, env in enumerate(self.training_env.envs):
            if hasattr(env, "get_reward_component_averages"):
                averages = env.get_reward_component_averages()
                for key, val in averages.items():
                    self.logger.record(f"reward_components/{key}", val)
        return True


class CurriculumCallback(BaseCallback):
    """Monitors evaluation performance and automatically manages curriculum progression."""

    def __init__(self, eval_freq: int = 10000, verbose: int = 1):
        super().__init__(verbose)
        self.eval_freq = eval_freq

    def _on_step(self) -> bool:
        if self.n_calls % self.eval_freq == 0:
            for env in self.training_env.envs:
                if hasattr(env, "get_current_robot_metrics"):
                    metrics = env.get_current_robot_metrics()
                    if metrics:
                        phase = metrics[0].get("phase_name", "UNKNOWN")
                        dist = metrics[0].get("distance_from_start", 0.0)
                        self.logger.record("curriculum/phase_idx", phase)
                        self.logger.record("curriculum/max_distance", dist)
                        if self.verbose > 0:
                            print(f"[CurriculumCallback] Step {self.num_timesteps}: Current Phase = {phase}, Max Dist = {dist:.2f}m")
        return True