"""LLM agent for generating robot control code in a ReAct loop."""

from __future__ import annotations

import logging
import textwrap
import traceback

import numpy as np

from orbit.generator.base import BaseSimulator, Trajectory

logger = logging.getLogger(__name__)


class DemoGeneratorAgent:
    """LLM agent that generates robot control code via iterative refinement.

    Uses a ReAct loop: the LLM writes Python control code, the simulator
    executes it, and the result is fed back for refinement.
    """

    def __init__(
        self,
        max_attempts: int = 20,
        max_steps_per_attempt: int = 500,
        llm_provider: str = "gemini",
    ):
        self.max_attempts = max_attempts
        self.max_steps_per_attempt = max_steps_per_attempt
        self.llm_provider = llm_provider

    def generate_episode(
        self,
        task: str,
        sim: BaseSimulator,
    ) -> Trajectory | None:
        """Try to complete a task through iterative code generation.

        Args:
            task: Natural language task description.
            sim: Simulator instance.

        Returns:
            Successful Trajectory, or None if all attempts failed.
        """
        api_desc = sim.get_api_description()
        system_prompt = self._build_system_prompt(task, api_desc)
        context = ""

        for attempt in range(self.max_attempts):
            sim.reset()

            # Ask LLM to generate control code
            prompt = system_prompt + context
            code = self._call_llm(prompt)
            if code is None:
                logger.warning("LLM returned None on attempt %d", attempt)
                context += f"\n\nAttempt {attempt + 1}: LLM failed to respond."
                continue

            # Extract Python code from response
            code = self._extract_code(code)

            # Execute in sandbox
            trajectory, error = self._execute_safely(code, sim, task)

            if trajectory is not None and trajectory.success:
                logger.info(
                    "Task '%s' succeeded on attempt %d (%d steps)",
                    task, attempt + 1, trajectory.total_steps,
                )
                return trajectory

            # Build feedback for next attempt
            if error:
                context += (
                    f"\n\nAttempt {attempt + 1} FAILED with error:\n{error}\n"
                    f"Fix the error and try again."
                )
            elif trajectory:
                final_state = trajectory.states[-1] if trajectory.states else {}
                context += (
                    f"\n\nAttempt {attempt + 1}: Code ran but task not completed.\n"
                    f"Final state: qpos={final_state.get('qpos', 'N/A')[:6]}...\n"
                    f"Steps taken: {trajectory.total_steps}\n"
                    f"Adjust your approach."
                )

        logger.warning("Task '%s' failed after %d attempts", task, self.max_attempts)
        return None

    def _build_system_prompt(self, task: str, api_desc: str) -> str:
        return textwrap.dedent(f"""\
            You control a robot in simulation. Your goal is to complete a task
            by writing Python code that commands the robot.

            TASK: {task}

            SIMULATOR API:
            {api_desc}

            Write a Python function called `run(sim)` that:
            1. Calls sim.get_state() to observe the current state
            2. Computes actions based on the state
            3. Calls sim.execute_action(action) to move the robot
            4. Repeats until the task is complete

            The function should return True if successful, False otherwise.

            Important:
            - Use numpy (imported as np) for array operations
            - Actions must match the expected shape
            - Move slowly and deliberately
            - Check state after each action to verify progress

            Respond with ONLY the Python code inside ```python``` blocks.
        """)

    def _extract_code(self, response: str) -> str:
        """Extract Python code from LLM response."""
        # Look for ```python ... ``` blocks
        if "```python" in response:
            parts = response.split("```python")
            if len(parts) > 1:
                code = parts[1].split("```")[0]
                return code.strip()

        # Look for ``` ... ``` blocks
        if "```" in response:
            parts = response.split("```")
            if len(parts) >= 3:
                return parts[1].strip()

        # Assume the whole response is code
        return response.strip()

    def _execute_safely(
        self,
        code: str,
        sim: BaseSimulator,
        task: str,
    ) -> tuple[Trajectory | None, str | None]:
        """Execute generated code in a restricted environment.

        Returns:
            (trajectory, error) — one of them will be None.
        """
        trajectory = Trajectory()

        # Create a restricted namespace
        namespace = {
            "sim": sim,
            "np": np,
            "trajectory": trajectory,
        }

        try:
            # Compile and execute the code
            compiled = compile(code, "<llm_generated>", "exec")
            exec(compiled, namespace)  # noqa: S102

            # Call the run function if it exists
            if "run" in namespace and callable(namespace["run"]):
                result = namespace["run"](sim)

                # Record final state
                state = sim.get_state()
                trajectory.states.append(state)

                # Check success
                if result is True or sim.check_success(task, state):
                    trajectory.success = True

            return trajectory, None

        except Exception as e:
            error_msg = f"{type(e).__name__}: {e}\n{traceback.format_exc()[-500:]}"
            return None, error_msg

    def _call_llm(self, prompt: str) -> str | None:
        """Call the LLM to generate code."""
        if self.llm_provider == "gemini":
            return self._call_gemini(prompt)
        else:
            logger.error("Unknown LLM provider: %s", self.llm_provider)
            return None

    def _call_gemini(self, prompt: str) -> str | None:
        """Call Gemini via the existing llm_utils."""
        from orbit.llm_utils import call_gemini, load_api_key

        api_key = load_api_key()
        if not api_key:
            logger.error("GOOGLE_API_KEY not set — cannot generate demonstrations")
            return None

        return call_gemini(prompt, api_key, temperature=0.2, max_tokens=4096)
