import argparse
import os
from pathlib import Path

from .env import JackBotEnv


def parse_args():
    parser = argparse.ArgumentParser(description="Train a joint-command policy for JackBot.")
    parser.add_argument("--timesteps", type=int, default=500_000, help="Total training timesteps")
    parser.add_argument("--model-path", type=str, default="ml/checkpoints/ppo_joint_command", help="Where to save the trained model")
    parser.add_argument("--seed", type=int, default=0, help="Random seed")
    parser.add_argument("--device", type=str, default="auto", help="Device to use: 'cpu', 'cuda', or 'auto' to autodetect")
    parser.add_argument("--use-gui", action="store_true", help="Enable PyBullet GUI during training")
    parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the training environment")
    parser.add_argument("--robot-spacing", type=float, default=0.5, help="Spacing between robots in meters")
    parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
    return parser.parse_args()


def train(
    total_timesteps: int,
    model_path: str,
    seed: int = 0,
    device: str = "auto",
    use_gui: bool = False,
    num_robots: int = 1,
    robot_spacing: float = 0.5,
    start_pose: str = "init_deg",
):
    try:
        from stable_baselines3 import PPO
        from stable_baselines3.common.vec_env import DummyVecEnv
    except ImportError as exc:
        raise ImportError(
            "stable-baselines3 and gym are required for training. "
            "Install them with: pip install stable-baselines3 gym"
        ) from exc

    env = DummyVecEnv([
        lambda: JackBotEnv(
            use_gui=use_gui,
            random_command=True,
            num_robots=num_robots,
            robot_spacing=robot_spacing,
            start_pose=start_pose,
        )
    ])

    def resolve_device(requested_device: str) -> str:
        try:
            import torch
        except ImportError:
            if requested_device != "cpu":
                raise RuntimeError(
                    "PyTorch is not installed in the active environment. "
                    "Install torch with a GPU-enabled build before using --device cuda."
                )
            return "cpu"

        hip_supported = getattr(torch.version, "hip", None) is not None
        cuda_available = torch.cuda.is_available()
        hip_available = hip_supported and getattr(torch.backends, "hip", None) is not None and torch.backends.hip.is_available()

        if requested_device == "auto":
            if hip_available or cuda_available:
                return "cuda"
            return "cpu"

        if requested_device in {"cuda", "gpu", "hip"}:
            if hip_available or cuda_available:
                return "cuda"
            raise RuntimeError(
                f"GPU device requested ({requested_device}) but no CUDA/ROCm-capable PyTorch is available. "
                f"Installed torch build: {torch.__version__} (hip={getattr(torch.version, 'hip', None)}, cuda={cuda_available})"
            )

        if requested_device == "cpu":
            return "cpu"

        raise ValueError(
            f"Unsupported device '{requested_device}'. Use 'cpu', 'cuda', or 'auto'."
        )

    device = resolve_device(device)

    model = PPO(
        "MlpPolicy",
        env,
        verbose=1,
        seed=seed,
        device=device,
        tensorboard_log=str(Path(__file__).resolve().parent / "tensorboard"),
    )
    model.learn(total_timesteps=total_timesteps)

    Path(model_path).parent.mkdir(parents=True, exist_ok=True)
    model.save(model_path)
    env.close()


if __name__ == "__main__":
    args = parse_args()
    train(
        args.timesteps,
        args.model_path,
        seed=args.seed,
        device=args.device,
        use_gui=args.use_gui,
        num_robots=args.num_robots,
        robot_spacing=args.robot_spacing,
        start_pose=args.start_pose,
    )