"""Tests for the pre/post-training hook system."""

from unittest.mock import MagicMock, patch

import pytest

from orbit.hooks import PreTrainingResult, _detect_dataset_from_command


class TestPreTrainingResult:
    def test_dataclass_creation(self):
        result = PreTrainingResult(
            passed=True,
            grade="B",
            score=84,
            summary="Dataset ready for training",
        )
        assert result.passed is True
        assert result.grade == "B"
        assert result.score == 84
        assert result.cleaned_dataset is None
        assert result.curated_dataset is None
        assert result.warnings == []

    def test_failed_result(self):
        result = PreTrainingResult(
            passed=False,
            grade="F",
            score=15,
            summary="Too few episodes",
            warnings=["Only 5 episodes", "Dead joints detected"],
        )
        assert result.passed is False
        assert len(result.warnings) == 2


class TestPreTrainingCheck:
    @patch("orbit.analyzer.dataset_loader.load_dataset_info")
    def test_returns_result_on_load_failure(self, mock_load_info):
        from orbit.hooks import pre_training_check

        mock_load_info.side_effect = ValueError("Dataset not found")

        result = pre_training_check("bad/dataset", "act")
        assert result.passed is False
        assert result.grade == "F"


class TestDetectDatasetFromCommand:
    def test_lerobot_style(self):
        cmd = ["lerobot-train", "--dataset.repo_id=lerobot/my-data", "--policy.type=act"]
        assert _detect_dataset_from_command(cmd) == "lerobot/my-data"

    def test_generic_dataset_flag(self):
        cmd = ["python", "train.py", "--dataset", "/path/to/data"]
        assert _detect_dataset_from_command(cmd) == "/path/to/data"

    def test_data_root_dir(self):
        cmd = ["torchrun", "finetune.py", "--data_root_dir", "./data"]
        assert _detect_dataset_from_command(cmd) == "./data"

    def test_dataset_path(self):
        cmd = ["python", "launch_finetune.py", "--dataset-path", "./my-data"]
        assert _detect_dataset_from_command(cmd) == "./my-data"

    def test_no_dataset_found(self):
        cmd = ["python", "train.py", "--batch_size", "32"]
        assert _detect_dataset_from_command(cmd) is None


class TestRunTrainingPipeline:
    def test_function_exists(self):
        from orbit.hooks import run_training_pipeline
        assert callable(run_training_pipeline)
