"""
ml/callbacks.py - Stable-Baselines3 callbacks for training diagnostics.

These callbacks extend SB3 training with two responsibilities: logging reward-component
statistics for TensorBoard/console output, and checking whether the curriculum should
advance to a harder set of commands based on recent training performance.
"""

import numpy as np
from stable_baselines3.common.callbacks import BaseCallback


class RewardLoggerCallback(BaseCallback):
    """
    Logs individual reward component averages to TensorBoard and prints 
    the best worker's performance breakdown to the console per iteration.
    """

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

    def _on_step(self) -> bool:
        return True

    def _on_rollout_end(self) -> None:
        self.iteration += 1
        if self.training_env is None:
            return

        try:
            # Safely query method across all parallel worker processes
            all_worker_averages = self.training_env.env_method("get_reward_component_averages")
        except Exception:
            return

        if not all_worker_averages or len(all_worker_averages) == 0:
            return

        # 1. Log mean component values across ALL workers to TensorBoard
        component_keys = all_worker_averages[0].keys()
        for key in component_keys:
            mean_val = float(np.mean([w.get(key, 0.0) for w in all_worker_averages]))
            self.logger.record(f"reward_components/{key}", mean_val)

        # 2. Identify the best performing worker of this iteration
        worker_totals = [sum(w.values()) for w in all_worker_averages]
        best_worker_idx = int(np.argmax(worker_totals))
        best_averages = all_worker_averages[best_worker_idx]
        best_total = worker_totals[best_worker_idx]

        # 3. Print best worker breakdown to console
        if self.verbose > 0:
            print(f"\n" + "=" * 65)
            print(f" ITERATION {self.iteration} | BEST WORKER (#{best_worker_idx}) REWARD BREAKDOWN")
            print(f" Total Avg Reward / Step: {best_total:+.4f}")
            print("-" * 65)
            for key, val in best_averages.items():
                print(f"   • {key:<26}: {val:+.5f}")
            print("=" * 65 + "\n")


class CurriculumCallback(BaseCallback):
    """
    Monitors training metrics using SB3's native ep_info_buffer and 
    dynamically advances curriculum phases across worker processes.
    """

    def __init__(self, reward_threshold: float = 100.0, verbose: int = 1):
        super().__init__(verbose)
        self.reward_threshold = reward_threshold

    def _on_step(self) -> bool:
        return True

    def _on_rollout_end(self) -> None:
        if self.training_env is None:
            return

        # SB3 natively records finished episode stats in self.model.ep_info_buffer
        if hasattr(self.model, "ep_info_buffer") and len(self.model.ep_info_buffer) > 0:
            recent_rewards = [ep_info["r"] for ep_info in self.model.ep_info_buffer]
            mean_reward = float(np.mean(recent_rewards[-50:]))

            try:
                # Query current phase from worker 0
                phases = self.training_env.get_attr("curriculum_phase")
                current_phase = phases[0]

                # Advance curriculum if mean reward exceeds threshold
                if mean_reward >= self.reward_threshold:
                    if hasattr(current_phase, "next"):
                        next_phase = current_phase.next()
                        if next_phase != current_phase:
                            self.training_env.set_attr("curriculum_phase", next_phase)
                            if self.verbose > 0:
                                print(
                                    f"\n[Curriculum] 🚀 Promoted workers to phase: {next_phase.name} "
                                    f"(Mean Reward: {mean_reward:.2f})"
                                )
            except Exception:
                pass  # Keep rollout loop running safely if phase check fails