"""Main entry point for LLM-driven demonstration generation."""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np

logger = logging.getLogger(__name__)


@dataclass
class GenerationResult:
    """Result of demonstration generation."""

    task: str
    simulator: str
    total_attempts: int
    successful_episodes: int
    failed_episodes: int
    trajectories: list = field(default_factory=list)
    output_dir: str = ""
    cost_estimate: str = ""


def generate_demonstrations(
    task_description: str,
    simulator: str = "mujoco",
    sim_config: dict | None = None,
    num_episodes: int = 10,
    max_attempts_per_episode: int = 20,
    output_format: str = "lerobot-v3",
    output_dir: str = "./generated/",
    model_path: str | None = None,
) -> GenerationResult:
    """Generate robot demonstrations using an LLM agent in simulation.

    The agent:
    1. Reads the task description and simulator API docs
    2. Writes a Python control script
    3. Executes it in the simulator
    4. Observes the result (success/failure + state trajectory)
    5. Iterates until success
    6. Records the successful trajectory as a demonstration

    Args:
        task_description: Natural language description of the task.
        simulator: Simulator backend ("mujoco" or "isaac").
        sim_config: Optional simulator configuration dict.
        num_episodes: Number of demonstrations to generate.
        max_attempts_per_episode: Max LLM attempts per episode.
        output_format: Output dataset format.
        output_dir: Directory to save generated data.
        model_path: Path to simulator model file (e.g., .xml for MuJoCo).

    Returns:
        GenerationResult with statistics and output path.

    Important caveats (from FAEA research):
        - Uses privileged state info (object positions), not camera images
        - 2-8 seconds per decision (not real-time)
        - $0.51-$5.60 per successful episode
        - 0% on precision tasks (peg insertion)
        - Best use: generating seed data for VLA fine-tuning
    """
    from orbit.generator.agent import DemoGeneratorAgent

    # Initialize simulator
    sim = _create_simulator(simulator, model_path, sim_config)

    # Initialize agent
    agent = DemoGeneratorAgent(
        max_attempts=max_attempts_per_episode,
    )

    result = GenerationResult(
        task=task_description,
        simulator=simulator,
        total_attempts=0,
        successful_episodes=0,
        failed_episodes=0,
        output_dir=output_dir,
    )

    out_path = Path(output_dir)
    out_path.mkdir(parents=True, exist_ok=True)

    for ep_idx in range(num_episodes):
        logger.info("Generating episode %d/%d for task: %s", ep_idx + 1, num_episodes, task_description)

        trajectory = agent.generate_episode(task_description, sim)

        if trajectory is not None and trajectory.success:
            result.successful_episodes += 1
            result.trajectories.append(trajectory)

            # Save trajectory
            _save_trajectory(trajectory, ep_idx, out_path, output_format)
        else:
            result.failed_episodes += 1

        result.total_attempts += agent.max_attempts

    # Cost estimate (Gemini Flash-Lite: ~$0.01 per call)
    estimated_calls = result.total_attempts
    result.cost_estimate = f"~${estimated_calls * 0.01:.2f} (estimated)"

    # Save metadata
    meta = {
        "task": task_description,
        "simulator": simulator,
        "successful_episodes": result.successful_episodes,
        "failed_episodes": result.failed_episodes,
        "total_attempts": result.total_attempts,
        "output_format": output_format,
        "cost_estimate": result.cost_estimate,
    }
    (out_path / "generation_metadata.json").write_text(json.dumps(meta, indent=2))

    return result


def _create_simulator(
    simulator: str,
    model_path: str | None,
    config: dict | None,
):
    """Create a simulator instance."""
    if simulator == "mujoco":
        if model_path is None:
            raise ValueError("--model path required for MuJoCo simulator")
        from orbit.generator.mujoco_sim import MuJoCoSimulator
        return MuJoCoSimulator(model_path)
    elif simulator == "isaac":
        raise NotImplementedError(
            "Isaac Lab simulator support is planned but not yet implemented. "
            "Use --sim mujoco for now."
        )
    else:
        raise ValueError(f"Unknown simulator: {simulator}. Choose: mujoco, isaac")


def _save_trajectory(trajectory, episode_idx: int, output_dir: Path, fmt: str) -> None:
    """Save a trajectory to disk."""
    states_arr, actions_arr = trajectory.to_arrays()

    if fmt == "lerobot-v3":
        # Save as parquet-compatible numpy arrays
        ep_dir = output_dir / f"episode_{episode_idx:06d}"
        ep_dir.mkdir(exist_ok=True)

        if actions_arr.size > 0:
            np.save(ep_dir / "actions.npy", actions_arr)
        if states_arr.size > 0:
            np.save(ep_dir / "states.npy", states_arr)

        # Metadata
        meta = {
            "episode_index": episode_idx,
            "num_frames": len(actions_arr) if actions_arr.size > 0 else 0,
            "success": trajectory.success,
        }
        (ep_dir / "metadata.json").write_text(json.dumps(meta))
    else:
        # Generic numpy save
        np.savez(
            output_dir / f"episode_{episode_idx:06d}.npz",
            actions=actions_arr,
            states=states_arr,
        )
