from __future__ import annotations

import warnings
from collections import deque
from typing import Iterable, Optional, Sequence


class HippotorchCurriculumCallback:
    """Simple curriculum driver for corridor-like envs with set_corridor_length.

    Compatible with SB3 callback API (on_step) but does not require SB3.
    Tracks a recent window of episode successes and increases length when
    threshold is met.
    """

    def __init__(
        self,
        *,
        min_length: int = 5,
        max_length: int = 30,
        step: int = 5,
        window: int = 10,
        threshold: float = 0.8,
    ) -> None:
        self.min_length = int(min_length)
        self.max_length = int(max_length)
        self.step = int(step)
        self.window = int(window)
        self.threshold = float(threshold)
        self._successes: deque[float] = deque(maxlen=self.window)
        self._current_length: Optional[int] = None
        self._warned_missing_setter = False

    def on_episode_end(self, env, reward: float) -> None:
        self._maybe_init(env)
        self._successes.append(1.0 if reward > 0 else 0.0)
        if len(self._successes) == self.window and self._current_length is not None:
            rate = sum(self._successes) / self.window
            if rate >= self.threshold and self._current_length < self.max_length:
                new_L = min(self.max_length, self._current_length + self.step)
                if self._set_length(env, new_L):
                    self._current_length = new_L
                    self._successes.clear()

    # SB3-compatible hooks (no-ops unless plugged into SB3)
    def _maybe_init(self, env) -> None:
        if self._current_length is None:
            current = int(getattr(env, "length", self.min_length))
            self._current_length = max(current, self.min_length)
            self._set_length(env, self._current_length)

    def on_step(self) -> bool:  # pragma: no cover - SB3 integration stub
        return True

    # Internal helpers -----------------------------------------------------
    def _iter_envs(self, env) -> Iterable[object]:
        if env is None:
            return
        if isinstance(env, Sequence) and not isinstance(env, (str, bytes)):
            for item in env:
                yield from self._iter_envs(item)
            return
        yield env
        nested = getattr(env, "envs", None)
        if nested is not None:
            yield from self._iter_envs(nested)
        unwrapped = getattr(env, "unwrapped", None)
        if unwrapped is not None and unwrapped is not env:
            yield from self._iter_envs(unwrapped)

    def _set_length(self, env, length: int) -> bool:
        if env is None:
            self._warn_missing_setter(env)
            return False
        # VecEnv-compatible hook
        if hasattr(env, "env_method"):
            try:
                env.env_method("set_corridor_length", length)
                return True
            except Exception:
                pass

        applied = False
        for target in self._iter_envs(env):
            if hasattr(target, "set_corridor_length"):
                try:
                    target.set_corridor_length(length)
                    applied = True
                except Exception:
                    continue

        if not applied:
            self._warn_missing_setter(env)
        return applied

    def _warn_missing_setter(self, env) -> None:
        if self._warned_missing_setter:
            return
        name = "<unknown-env>"
        if env is not None:
            name = env.__class__.__name__
        warnings.warn(
            (
                "HippotorchCurriculumCallback: environment "
                f"{name} lacks set_corridor_length(); skipping curriculum update."
            ),
            stacklevel=2,
        )
        self._warned_missing_setter = True
