"""Base simulator interface for LLM-driven demonstration generation."""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field

import numpy as np


@dataclass
class Trajectory:
    """A recorded trajectory from simulation."""

    states: list[np.ndarray] = field(default_factory=list)
    actions: list[np.ndarray] = field(default_factory=list)
    observations: list[dict] = field(default_factory=list)
    success: bool = False
    total_steps: int = 0
    metadata: dict = field(default_factory=dict)

    def to_arrays(self) -> tuple[np.ndarray, np.ndarray]:
        """Convert to (states, actions) numpy arrays."""
        return (
            np.array(self.states) if self.states else np.array([]),
            np.array(self.actions) if self.actions else np.array([]),
        )


class BaseSimulator(ABC):
    """Abstract interface for robot simulators."""

    @abstractmethod
    def get_state(self) -> dict:
        """Return current simulator state.

        Returns:
            Dict with at minimum: qpos, qvel, site_xpos (end-effector positions).
        """

    @abstractmethod
    def execute_action(self, action: np.ndarray) -> dict:
        """Execute one action step and return new state."""

    @abstractmethod
    def reset(self) -> dict:
        """Reset simulator to initial state."""

    @abstractmethod
    def get_api_description(self) -> str:
        """Return a description of the simulator API for the LLM agent."""

    def check_success(self, task: str, state: dict) -> bool:
        """Check if the task is complete. Override in subclasses."""
        return False
