import dataclasses
import sys
import tempfile
from pathlib import Path

import numpy as np
import structlog
import torch
import tqdm
from clearml import Logger
from skrl.agents.torch.ppo import PPO, PPO_DEFAULT_CONFIG
from skrl.memories.torch import RandomMemory
from skrl.resources.preprocessors.torch import RunningStandardScaler
from skrl.trainers.torch import SequentialTrainer

from src.core.runner import BaseRunner
from src.models.mlp import SharedMLP

log = structlog.get_logger()


@dataclasses.dataclass
class TrainerConfig:
    # PPO
    rollout_steps: int = 2048
    learning_epochs: int = 8
    mini_batches: int = 4
    discount_factor: float = 0.99
    gae_lambda: float = 0.95
    learning_rate: float = 3e-4
    clip_ratio: float = 0.2
    value_loss_scale: float = 0.5
    entropy_loss_scale: float = 0.01
    kl_threshold: float = 0.01      # KL-adaptive LR target; 0 = fixed LR

    hidden_sizes: tuple[int, ...] = (64, 64)

    # Policy
    initial_log_std: float = 0.5    # initial exploration noise
    min_log_std: float = -2.0       # minimum exploration noise
    max_log_std: float = 2.0        # maximum exploration noise (2.0 ≈ σ=7.4)

    # Training
    total_timesteps: int = 1_000_000
    log_interval: int = 10
    checkpoint_interval: int = 50_000

    # Video recording (uploaded to ClearML)
    record_video_every: int = 10_000   # 0 = disabled
    record_video_fps: int = 0          # 0 = derive from sim dt×substeps

    # History encoder (implicit adaptation).  The window size comes from
    # the runner (runner.history_length) — single source of truth.
    embedding_dim: int = 32     # history encoder output dimension


# ── Video-recording trainer ──────────────────────────────────────────

class VideoRecordingTrainer(SequentialTrainer):
    """SequentialTrainer with periodic evaluation videos uploaded to ClearML."""

    def __init__(self, env, agents, cfg=None, trainer_config: TrainerConfig | None = None):
        super().__init__(env=env, agents=agents, cfg=cfg)
        self._tcfg = trainer_config
        self._video_dir = Path(tempfile.mkdtemp(prefix="rl_videos_"))

    def single_agent_train(self) -> None:
        assert self.num_simultaneous_agents == 1 and self.env.num_agents == 1

        states, infos = self.env.reset()

        for timestep in tqdm.tqdm(
            range(self.initial_timestep, self.timesteps),
            disable=self.disable_progressbar,
            file=sys.stdout,
        ):
            self.agents.pre_interaction(timestep=timestep, timesteps=self.timesteps)

            with torch.no_grad():
                actions = self.agents.act(states, timestep=timestep, timesteps=self.timesteps)[0]
                next_states, rewards, terminated, truncated, infos = self.env.step(actions)

                self.agents.record_transition(
                    states=states, actions=actions, rewards=rewards,
                    next_states=next_states, terminated=terminated,
                    truncated=truncated, infos=infos,
                    timestep=timestep, timesteps=self.timesteps,
                )

                if self.environment_info in infos:
                    for k, v in infos[self.environment_info].items():
                        if isinstance(v, torch.Tensor) and v.numel() == 1:
                            self.agents.track_data(f"Info / {k}", v.item())

            self.agents.post_interaction(timestep=timestep, timesteps=self.timesteps)

            # Auto-reset for multi-env; single-env resets manually
            if self.env.num_envs > 1:
                states = next_states
            else:
                if terminated.any() or truncated.any():
                    with torch.no_grad():
                        states, infos = self.env.reset()
                else:
                    states = next_states

            # Periodic video recording.  Recording steps the (shared) envs,
            # so it returns a freshly reset observation — the training loop
            # MUST continue from it, otherwise the recorded transitions no
            # longer match the actual env state.
            if (
                self._tcfg
                and self._tcfg.record_video_every > 0
                and (timestep + 1) % self._tcfg.record_video_every == 0
            ):
                fresh_states = self._record_video(timestep + 1)
                if fresh_states is not None:
                    states = fresh_states

    # ── helpers ───────────────────────────────────────────────────────

    def _get_fps(self) -> int:
        if self._tcfg and self._tcfg.record_video_fps > 0:
            return self._tcfg.record_video_fps
        dt = getattr(self.env.config, "dt", 0.02)
        substeps = getattr(self.env.config, "substeps", 1)
        # SerialRunner has dt but no substeps — dt *is* the control period.
        return max(1, int(round(1.0 / (dt * substeps))))

    def _record_video(self, timestep: int) -> torch.Tensor | None:
        """Record an eval episode and upload it to ClearML.

        Returns the freshly reset observation the training loop should
        continue from (the recording steps the shared envs), or ``None``
        if even the final reset failed.
        """
        try:
            import imageio.v3 as iio
        except ImportError:
            iio = None

        # Rendering needs a GL backend (EGL/OSMesa); never let a headless GL
        # failure crash training — log it and skip the video.
        if iio is not None:
            try:
                fps = self._get_fps()
                max_steps = getattr(self.env.env.config, "max_steps", 500)
                frames: list[np.ndarray] = []

                obs, _ = self.env.reset()
                with torch.no_grad():
                    for _ in range(max_steps):
                        action = self.agents.act(obs, timestep=timestep, timesteps=self.timesteps)[0]
                        obs, _, terminated, truncated, _ = self.env.step(action)

                        frame = self.env.render()
                        if frame is not None:
                            frames.append(frame)

                        if (terminated | truncated).any().item():
                            break

                if frames:
                    path = str(self._video_dir / f"step_{timestep}.mp4")
                    iio.imwrite(path, frames, fps=fps)

                    logger = Logger.current_logger()
                    if logger:
                        logger.report_media(
                            "Training Video", f"step_{timestep}",
                            local_path=path, iteration=timestep,
                        )
            except Exception as exc:
                log.warning("video_recording_failed", timestep=timestep, error=str(exc))

        # Always leave the envs freshly reset and hand the new observation
        # back to the training loop.
        try:
            with torch.no_grad():
                states, _ = self.env.reset()
            return states
        except Exception as exc:
            log.warning("post_video_reset_failed", timestep=timestep, error=str(exc))
            return None


# ── Main trainer ─────────────────────────────────────────────────────

class Trainer:
    def __init__(self, runner: BaseRunner, config: TrainerConfig):
        self.runner = runner
        self.config = config
        self._init_agent()

    def _init_agent(self) -> None:
        device = self.runner.device
        obs_space = self.runner.observation_space
        act_space = self.runner.action_space

        self.memory = RandomMemory(
            memory_size=self.config.rollout_steps,
            num_envs=self.runner.num_envs,
            device=device,
        )

        # Determine raw obs dim (without history augmentation) and the
        # history window size — both come from the runner so the model
        # always matches the observation layout it produces.
        raw_obs_dim = self.runner.env.observation_space.shape[0]
        history_length = getattr(self.runner.config, "history_length", 0)

        self.model = SharedMLP(
            observation_space=obs_space,
            action_space=act_space,
            device=device,
            hidden_sizes=self.config.hidden_sizes,
            initial_log_std=self.config.initial_log_std,
            min_log_std=self.config.min_log_std,
            max_log_std=self.config.max_log_std,
            history_length=history_length,
            raw_obs_dim=raw_obs_dim,
            embedding_dim=self.config.embedding_dim,
        )

        models = {"policy": self.model, "value": self.model}

        agent_cfg = PPO_DEFAULT_CONFIG.copy()
        agent_cfg.update({
            "rollouts": self.config.rollout_steps,
            "learning_epochs": self.config.learning_epochs,
            "mini_batches": self.config.mini_batches,
            "discount_factor": self.config.discount_factor,
            "lambda": self.config.gae_lambda,
            "learning_rate": self.config.learning_rate,
            "ratio_clip": self.config.clip_ratio,
            "value_loss_scale": self.config.value_loss_scale,
            "entropy_loss_scale": self.config.entropy_loss_scale,
            # Truncation (time limit) must bootstrap from the value function;
            # without this the value target is biased at every max_steps cut.
            "time_limit_bootstrap": True,
            "state_preprocessor": RunningStandardScaler,
            "state_preprocessor_kwargs": {"size": obs_space, "device": device},
            "value_preprocessor": RunningStandardScaler,
            "value_preprocessor_kwargs": {"size": 1, "device": device},
        })
        if self.config.kl_threshold > 0:
            from skrl.resources.schedulers.torch import KLAdaptiveLR
            agent_cfg["learning_rate_scheduler"] = KLAdaptiveLR
            agent_cfg["learning_rate_scheduler_kwargs"] = {
                "kl_threshold": self.config.kl_threshold,
            }
        # Wire up logging frequency: write_interval is in timesteps.
        # log_interval=1 → log every PPO update (= every rollout_steps timesteps).
        agent_cfg["experiment"]["write_interval"] = self.config.log_interval
        agent_cfg["experiment"]["checkpoint_interval"] = max(
            self.config.checkpoint_interval, self.config.rollout_steps
        )

        self.agent = PPO(
            models=models,
            memory=self.memory,
            observation_space=obs_space,
            action_space=act_space,
            device=device,
            cfg=agent_cfg,
        )

    def train(self) -> None:
        trainer = VideoRecordingTrainer(
            env=self.runner,
            agents=self.agent,
            cfg={"timesteps": self.config.total_timesteps},
            trainer_config=self.config,
        )
        trainer.train()

    def close(self) -> None:
        self.runner.close()