from typing import override

import jax.numpy as jnp

from envelope.environment import Info
from envelope.struct import field
from envelope.typing import Key, PyTree
from envelope.wrappers.wrapper import WrappedState, Wrapper


class TruncationWrapper(Wrapper):
    max_steps: int = field(kw_only=True)

    class TruncationState(WrappedState):
        steps: jnp.ndarray | int = field(default=0)

    @override
    def init(self, key: Key) -> tuple[WrappedState, Info]:
        inner_state, info = self.env.init(key)
        state = self.TruncationState(inner_state=inner_state, steps=0)
        return state, info.update(truncated=self.max_steps <= 0)

    @override
    def reset(self, state: WrappedState, key: Key) -> tuple[WrappedState, Info]:
        inner_state, info = self.env.reset(state.inner_state, key)
        state = state.replace(inner_state=inner_state, steps=0)
        return state, info.update(truncated=self.max_steps <= 0)

    @override
    def step(self, state: WrappedState, action: PyTree) -> tuple[WrappedState, Info]:
        next_inner_state, info = self.env.step(state.inner_state, action)
        steps = state.steps + 1
        state = self.TruncationState(inner_state=next_inner_state, steps=steps)
        truncated = jnp.asarray(steps) >= self.max_steps
        return state, info.update(truncated=truncated)
