import abc
import dataclasses
from typing import TypeVar, Generic, Any

import numpy as np
import torch
from gymnasium import spaces

from src.core.robot import RobotConfig, load_robot_config

T = TypeVar("T")


@dataclasses.dataclass
class BaseEnvConfig:
    max_steps: int = 1000
    robot_path: str = ""  # directory containing robot.yaml + URDF

class BaseEnv(abc.ABC, Generic[T]):
    def __init__(self, config: BaseEnvConfig):
        self.config = config
        self.robot: RobotConfig = load_robot_config(config.robot_path)

    @property
    @abc.abstractmethod
    def observation_space(self) -> spaces.Space:
        ...

    @property
    @abc.abstractmethod
    def action_space(self) -> spaces.Space:
        ...

    @abc.abstractmethod
    def build_state(self, qpos: torch.Tensor, qvel: torch.Tensor) -> Any:
        ...

    @abc.abstractmethod
    def compute_observations(self, state: Any) -> torch.Tensor:
        ...

    @abc.abstractmethod
    def compute_rewards(
        self, state: Any, actions: torch.Tensor, prev_actions: torch.Tensor,
    ) -> torch.Tensor:
        ...

    @abc.abstractmethod
    def compute_terminations(self, state: Any) -> torch.Tensor:
        ...

    def compute_truncations(self, step_counts: torch.Tensor) -> torch.Tensor:
        return step_counts >= self.config.max_steps

    def initial_state_ranges(
        self, nq: int, nv: int,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """Per-DOF uniform ranges for initial-state randomization.

        Returns (qpos_lo, qpos_hi, qvel_lo, qvel_hi) — offsets added to the
        model's default state on every reset.  All runners (CPU MuJoCo and
        MJX) sample from these, so initial-state distributions stay
        identical across backends.  Default: small ±0.05 perturbation.
        """
        return (
            np.full(nq, -0.05), np.full(nq, 0.05),
            np.full(nv, -0.05), np.full(nv, 0.05),
        )

    def is_reset_ready(self, qpos: torch.Tensor, qvel: torch.Tensor) -> bool:
        """Check whether the physical robot has settled enough to start an episode.

        Used by the SerialRunner after driving to center and waiting for the
        pendulum.  Default: always ready (sim doesn't need settling).
        """
        return True
