"""Tests for the multi-framework training configuration generator."""

from unittest.mock import MagicMock

import pytest

from orbit.training.base import FRAMEWORKS, TrainingConfig, get_framework_config


def _make_info(n_episodes=50, fps=30.0, action_dim=6):
    info = MagicMock()
    info.repo_id = "test/dataset"
    info.num_episodes = n_episodes
    info.num_frames = n_episodes * 100
    info.fps = fps
    info.robot_type = "so100"
    info.camera_names = ["observation.image"]
    info.shapes = {"action": [action_dim], "observation.image": [3, 480, 640]}
    return info


class TestFrameworkRegistry:
    def test_all_frameworks_listed(self):
        assert "lerobot" in FRAMEWORKS
        assert "openvla" in FRAMEWORKS
        assert "openpi" in FRAMEWORKS
        assert "groot" in FRAMEWORKS
        assert "osmo" in FRAMEWORKS
        assert "custom" in FRAMEWORKS

    def test_unknown_framework_raises(self):
        with pytest.raises(ValueError, match="Unknown framework"):
            get_framework_config(
                framework="nonexistent",
                dataset_repo="test/data",
                policy_type="act",
                info=_make_info(),
                episode_count=50,
                action_dim=6,
                camera_count=1,
                resolution="480p",
            )


class TestLeRobotConfig:
    def test_generates_lerobot_train_command(self):
        config = get_framework_config(
            framework="lerobot",
            dataset_repo="test/dataset",
            policy_type="act",
            info=_make_info(),
            episode_count=50,
            action_dim=6,
            camera_count=1,
            resolution="480p",
            batch_size=32,
            steps=100_000,
            learning_rate=1e-4,
            chunk_size=50,
        )
        assert isinstance(config, TrainingConfig)
        assert config.framework == "lerobot"
        assert "lerobot-train" in config.command
        assert "--dataset.repo_id=test/dataset" in config.command
        assert "--policy.type=act" in config.command
        assert config.dataset_format == "lerobot-v3"
        assert not config.conversion_needed

    def test_pre_checks_include_gate(self):
        config = get_framework_config(
            framework="lerobot",
            dataset_repo="test/dataset",
            policy_type="act",
            info=_make_info(),
            episode_count=50,
            action_dim=6,
            camera_count=1,
            resolution="480p",
        )
        assert any("orbit gate" in check for check in config.pre_checks)


class TestOpenVLAConfig:
    def test_generates_torchrun_command(self):
        config = get_framework_config(
            framework="openvla",
            dataset_repo="test/dataset",
            policy_type="act",
            info=_make_info(),
            episode_count=50,
            action_dim=6,
            camera_count=1,
            resolution="480p",
            gpu_memory_gb=80,
        )
        assert config.framework == "openvla"
        assert "torchrun" in config.command
        assert "finetune.py" in config.command
        assert config.dataset_format == "rlds"
        assert config.conversion_needed

    def test_oft_recommended_for_40gb_plus(self):
        """OFT (Open Fine-Tuning) should be default for 40GB+ VRAM."""
        config = get_framework_config(
            framework="openvla",
            dataset_repo="test/dataset",
            policy_type="act",
            info=_make_info(),
            episode_count=50,
            action_dim=6,
            camera_count=1,
            resolution="480p",
            gpu_memory_gb=80,
        )
        assert "openvla-7b-oft" in config.command
        assert any("OFT" in tip for tip in config.tips)

    def test_lora_fallback_for_low_vram(self):
        """LoRA should be used for GPUs with less than 40GB."""
        config = get_framework_config(
            framework="openvla",
            dataset_repo="test/dataset",
            policy_type="act",
            info=_make_info(),
            episode_count=50,
            action_dim=6,
            camera_count=1,
            resolution="480p",
            gpu_memory_gb=24,
        )
        assert "openvla-7b-oft" not in config.command
        assert "lora_rank" in config.command


class TestOpenPIConfig:
    def test_generates_uv_run_command(self):
        config = get_framework_config(
            framework="openpi",
            dataset_repo="test/my-dataset",
            policy_type="act",
            info=_make_info(),
            episode_count=50,
            action_dim=6,
            camera_count=1,
            resolution="480p",
        )
        assert config.framework == "openpi"
        assert "uv run scripts/train.py" in config.command
        assert config.dataset_format == "lerobot-v3"

    def test_generates_train_config_file(self):
        """OpenPI should generate a TrainConfig Python file template."""
        config = get_framework_config(
            framework="openpi",
            dataset_repo="test/my-dataset",
            policy_type="act",
            info=_make_info(action_dim=6),
            episode_count=50,
            action_dim=6,
            camera_count=1,
            resolution="480p",
        )
        assert config.generated_files
        config_files = list(config.generated_files.keys())
        assert any("config" in f and f.endswith(".py") for f in config_files)

        # Check config content
        config_content = list(config.generated_files.values())[0]
        assert "TrainConfig" in config_content
        assert "action_dim" in config_content
        assert "6" in config_content  # action_dim value

    def test_config_includes_camera_keys(self):
        """Config should include actual camera keys from dataset."""
        info = _make_info()
        info.camera_names = ["cam_left", "cam_right"]
        config = get_framework_config(
            framework="openpi",
            dataset_repo="test/data",
            policy_type="act",
            info=info,
            episode_count=50,
            action_dim=6,
            camera_count=2,
            resolution="480p",
        )
        config_content = list(config.generated_files.values())[0]
        assert "cam_left" in config_content
        assert "cam_right" in config_content


class TestGR00TConfig:
    def test_uses_launch_finetune_command(self):
        """GR00T N1.6 must use gr00t/experiment/launch_finetune.py."""
        config = get_framework_config(
            framework="groot",
            dataset_repo="test/dataset",
            policy_type="act",
            info=_make_info(),
            episode_count=50,
            action_dim=6,
            camera_count=1,
            resolution="480p",
        )
        assert config.framework == "groot"
        assert "launch_finetune.py" in config.command
        assert "GR00T-N1.6-3B" in config.command
        assert "--embodiment-tag" in config.command
        assert "--modality-config-path" in config.command

    def test_requires_lerobot_v2(self):
        """GR00T N1.6 requires LeRobot v2 format, NOT v3."""
        config = get_framework_config(
            framework="groot",
            dataset_repo="test/dataset",
            policy_type="act",
            info=_make_info(),
            episode_count=50,
            action_dim=6,
            camera_count=1,
            resolution="480p",
        )
        assert config.dataset_format == "lerobot-v2"
        assert config.conversion_needed
        assert "lerobot-v2" in config.conversion_command

    def test_generates_modality_config(self):
        """Must generate modality config Python file."""
        config = get_framework_config(
            framework="groot",
            dataset_repo="test/dataset",
            policy_type="act",
            info=_make_info(),
            episode_count=50,
            action_dim=6,
            camera_count=1,
            resolution="480p",
        )
        assert config.generated_files
        files = config.generated_files

        # Should have modality config Python file
        py_files = [f for f in files if f.endswith("_config.py")]
        assert len(py_files) >= 1
        py_content = files[py_files[0]]
        assert "ModalityConfig" in py_content

        # Should have modality.json
        json_files = [f for f in files if f.endswith("modality.json")]
        assert len(json_files) >= 1


class TestOSMOConfig:
    def test_generates_yaml_workflow(self):
        config = get_framework_config(
            framework="osmo",
            dataset_repo="test/dataset",
            policy_type="act",
            info=_make_info(),
            episode_count=50,
            action_dim=6,
            camera_count=1,
            resolution="480p",
            batch_size=32,
            steps=100_000,
        )
        assert config.framework == "osmo"
        assert "workflow:" in config.command
        assert "orbit-gate" in config.command
        assert "train-policy" in config.command
        assert "orbit-verify" in config.command

    def test_osmo_includes_all_pipeline_stages(self):
        config = get_framework_config(
            framework="osmo",
            dataset_repo="test/dataset",
            policy_type="act",
            info=_make_info(),
            episode_count=50,
            action_dim=6,
            camera_count=1,
            resolution="480p",
        )
        assert "orbit-analyze" in config.command
        assert "orbit-gate" in config.command
        assert "train-policy" in config.command
        assert "orbit-verify" in config.command


class TestCustomConfig:
    def test_generates_generic_command(self):
        config = get_framework_config(
            framework="custom",
            dataset_repo="test/dataset",
            policy_type="act",
            info=_make_info(),
            episode_count=50,
            action_dim=6,
            camera_count=1,
            resolution="480p",
            batch_size=8,
            steps=100_000,
            learning_rate=1e-4,
            fps=30.0,
        )
        assert config.framework == "custom"
        assert "python train.py" in config.command
        assert "--dataset test/dataset" in config.command


class TestAllFrameworksGeneration:
    @pytest.mark.parametrize("framework", FRAMEWORKS)
    def test_each_framework_returns_valid_config(self, framework):
        config = get_framework_config(
            framework=framework,
            dataset_repo="test/dataset",
            policy_type="act",
            info=_make_info(),
            episode_count=50,
            action_dim=6,
            camera_count=1,
            resolution="480p",
            gpu_memory_gb=80,
            batch_size=32,
            steps=100_000,
            learning_rate=1e-4,
            chunk_size=50,
            fps=30.0,
        )
        assert isinstance(config, TrainingConfig)
        assert config.framework == framework
        assert len(config.command) > 10
        assert config.estimated_time
        assert config.gpu_requirements
        assert config.dataset_format
