"""Tests for the LLM demonstration generator."""

import numpy as np
import pytest

from orbit.generator.base import BaseSimulator, Trajectory
from orbit.generator.agent import DemoGeneratorAgent


class MockSimulator(BaseSimulator):
    """Mock simulator for testing."""

    def __init__(self, succeed_on_attempt=1):
        self._succeed_on = succeed_on_attempt
        self._attempt = 0
        self._state = {
            "qpos": np.zeros(6),
            "qvel": np.zeros(6),
            "site_xpos": np.zeros((1, 3)),
            "time": 0.0,
        }

    def get_state(self):
        return self._state.copy()

    def execute_action(self, action):
        self._state["qpos"] = self._state["qpos"] + action[:6] * 0.01
        self._state["time"] += 0.01
        return self.get_state()

    def reset(self):
        self._state = {
            "qpos": np.zeros(6),
            "qvel": np.zeros(6),
            "site_xpos": np.zeros((1, 3)),
            "time": 0.0,
        }
        self._attempt += 1
        return self.get_state()

    def get_api_description(self):
        return "Mock simulator with 6 DOF. action shape: (6,)"

    def check_success(self, task, state):
        return self._attempt >= self._succeed_on


class TestTrajectory:
    def test_empty_trajectory(self):
        traj = Trajectory()
        assert not traj.success
        assert traj.total_steps == 0
        states, actions = traj.to_arrays()
        assert states.size == 0
        assert actions.size == 0

    def test_trajectory_with_data(self):
        traj = Trajectory()
        traj.states = [np.array([1, 2, 3]), np.array([4, 5, 6])]
        traj.actions = [np.array([0.1, 0.2]), np.array([0.3, 0.4])]
        traj.success = True
        traj.total_steps = 2

        states, actions = traj.to_arrays()
        assert states.shape == (2, 3)
        assert actions.shape == (2, 2)


class TestMockSimulator:
    def test_reset_returns_state(self):
        sim = MockSimulator()
        state = sim.reset()
        assert "qpos" in state
        assert "qvel" in state

    def test_execute_action_updates_state(self):
        sim = MockSimulator()
        sim.reset()
        action = np.ones(6)
        new_state = sim.execute_action(action)
        assert not np.allclose(new_state["qpos"], 0)

    def test_api_description(self):
        sim = MockSimulator()
        desc = sim.get_api_description()
        assert "6 DOF" in desc


class TestDemoGeneratorAgent:
    def test_extract_code_from_python_block(self):
        agent = DemoGeneratorAgent()
        response = "Here's the code:\n```python\ndef run(sim):\n    return True\n```\nDone."
        code = agent._extract_code(response)
        assert "def run(sim):" in code
        assert "return True" in code

    def test_extract_code_from_generic_block(self):
        agent = DemoGeneratorAgent()
        response = "```\ndef run(sim):\n    return True\n```"
        code = agent._extract_code(response)
        assert "def run(sim):" in code

    def test_extract_code_bare(self):
        agent = DemoGeneratorAgent()
        response = "def run(sim):\n    return True"
        code = agent._extract_code(response)
        assert "def run(sim):" in code

    def test_execute_safely_success(self):
        agent = DemoGeneratorAgent()
        sim = MockSimulator(succeed_on_attempt=1)
        sim.reset()

        code = "def run(sim):\n    sim.execute_action(np.ones(6))\n    return True"
        traj, error = agent._execute_safely(code, sim, "test task")

        assert error is None
        assert traj is not None

    def test_execute_safely_error(self):
        agent = DemoGeneratorAgent()
        sim = MockSimulator()
        sim.reset()

        code = "def run(sim):\n    raise ValueError('test error')"
        traj, error = agent._execute_safely(code, sim, "test task")

        assert traj is None
        assert error is not None
        assert "ValueError" in error

    def test_execute_safely_syntax_error(self):
        agent = DemoGeneratorAgent()
        sim = MockSimulator()
        sim.reset()

        code = "def run(sim:\n    return True"  # syntax error
        traj, error = agent._execute_safely(code, sim, "test task")

        assert traj is None
        assert error is not None


class TestGenerationResult:
    def test_generation_result_import(self):
        from orbit.generator.generate import GenerationResult
        result = GenerationResult(
            task="test",
            simulator="mujoco",
            total_attempts=10,
            successful_episodes=5,
            failed_episodes=5,
        )
        assert result.successful_episodes == 5
