"""Tests for orbit.analyzer.recommendations."""

from __future__ import annotations

from orbit.analyzer.coverage import FullCoverageReport
from orbit.analyzer.dataset_loader import DatasetInfo
from orbit.analyzer.quality import compute_quality_score
from orbit.analyzer.recommendations import (
    MAX_RECOMMENDATIONS,
    Recommendation,
    generate_recommendations,
)

# ---------------------------------------------------------------------------
# Fixtures / helpers
# ---------------------------------------------------------------------------


def _make_info(num_episodes: int = 50, num_frames: int = 12500) -> DatasetInfo:
    return DatasetInfo(
        repo_id="test/dataset",
        num_episodes=num_episodes,
        num_frames=num_frames,
        fps=30.0,
        robot_type="aloha",
        camera_names=["cam_high"],
        video_keys=["observation.images.cam_high"],
        shapes={"action": [14]},
        duration_seconds=num_frames / 30.0,
    )


def _make_coverage(
    position: float = 1.0,
    action: float = 1.0,
    consistency: float = 1.0,
    temporal: float = 1.0,
    **kwargs,
) -> FullCoverageReport:
    return FullCoverageReport(
        position_diversity=position,
        action_diversity=action,
        episode_consistency=consistency,
        temporal_coverage=temporal,
        **kwargs,
    )


def _make_quality(info, coverage):
    return compute_quality_score(info, coverage)


# ---------------------------------------------------------------------------
# generate_recommendations
# ---------------------------------------------------------------------------


class TestGenerateRecommendations:
    """Tests for generate_recommendations."""

    def test_perfect_dataset_no_recommendations(self):
        """A perfect dataset should generate 0 recommendations."""
        info = _make_info(num_episodes=50)
        coverage = _make_coverage(1.0, 1.0, 0.8, 1.0)
        quality = _make_quality(info, coverage)
        recs = generate_recommendations(info, coverage, quality)

        assert len(recs) == 0

    def test_low_episode_count_generates_recommendation(self):
        """Low episode count should produce an episode_count recommendation."""
        info = _make_info(num_episodes=20)
        coverage = _make_coverage(0.8, 0.8, 0.8, 0.8)
        quality = _make_quality(info, coverage)
        recs = generate_recommendations(info, coverage, quality)

        ep_recs = [r for r in recs if r.category == "episode_count"]
        assert len(ep_recs) == 1
        assert "30" in ep_recs[0].description  # 50 - 20 = 30

    def test_low_position_diversity_generates_recommendation(self):
        """Low position diversity should produce a position_diversity recommendation."""
        info = _make_info(num_episodes=50)
        coverage = _make_coverage(position=0.3, action=0.8, consistency=0.8, temporal=0.8)
        quality = _make_quality(info, coverage)
        recs = generate_recommendations(info, coverage, quality)

        pos_recs = [r for r in recs if r.category == "position_diversity"]
        assert len(pos_recs) == 1
        assert pos_recs[0].priority == "HIGH"

    def test_low_action_diversity_generates_recommendation(self):
        """Low action diversity should produce an action_diversity recommendation."""
        info = _make_info(num_episodes=50)
        coverage = _make_coverage(position=0.8, action=0.3, consistency=0.8, temporal=0.8)
        quality = _make_quality(info, coverage)
        recs = generate_recommendations(info, coverage, quality)

        act_recs = [r for r in recs if r.category == "action_diversity"]
        assert len(act_recs) == 1
        assert act_recs[0].priority == "HIGH"

    def test_recommendations_sorted_by_impact(self):
        """Recommendations must be sorted by expected_impact descending."""
        info = _make_info(num_episodes=10)
        coverage = _make_coverage(
            position=0.3, action=0.2, consistency=0.5, temporal=0.3,
        )
        quality = _make_quality(info, coverage)
        recs = generate_recommendations(info, coverage, quality)

        impacts = [r.expected_impact for r in recs]
        assert impacts == sorted(impacts, reverse=True)

    def test_max_six_recommendations(self):
        """No more than 6 recommendations should be returned."""
        info = _make_info(num_episodes=10)
        coverage = _make_coverage(
            position=0.2,
            action=0.1,
            consistency=0.5,
            temporal=0.2,
            outlier_episodes=[1, 2, 3],
        )
        quality = _make_quality(info, coverage)
        recs = generate_recommendations(info, coverage, quality)

        assert len(recs) <= MAX_RECOMMENDATIONS

    def test_recommendations_contain_concrete_numbers(self):
        """Recommendations should include specific numbers, not vague language."""
        info = _make_info(num_episodes=30)
        coverage = _make_coverage(position=0.4, action=0.3, consistency=0.8, temporal=0.8)
        quality = _make_quality(info, coverage)
        recs = generate_recommendations(info, coverage, quality)

        for rec in recs:
            # Each recommendation should have at least one digit
            assert any(c.isdigit() for c in rec.description), (
                f"Recommendation lacks concrete number: {rec.description}"
            )

    def test_recommendation_dataclass_fields(self):
        """Each recommendation should have all required fields."""
        info = _make_info(num_episodes=10)
        coverage = _make_coverage(position=0.3)
        quality = _make_quality(info, coverage)
        recs = generate_recommendations(info, coverage, quality)

        for rec in recs:
            assert isinstance(rec, Recommendation)
            assert rec.priority in ("HIGH", "MEDIUM", "LOW")
            assert len(rec.category) > 0
            assert len(rec.description) > 0
            assert rec.expected_impact > 0
            assert len(rec.effort) > 0

    def test_outlier_episodes_recommendation(self):
        """Outlier episodes should generate an episode_consistency recommendation."""
        info = _make_info(num_episodes=50)
        coverage = _make_coverage(
            consistency=0.8,
            outlier_episodes=[5, 12, 38],
        )
        quality = _make_quality(info, coverage)
        recs = generate_recommendations(info, coverage, quality)

        cons_recs = [r for r in recs if r.category == "episode_consistency"]
        assert len(cons_recs) == 1
        assert "5" in cons_recs[0].description
        assert "12" in cons_recs[0].description

    def test_temporal_coverage_recommendation(self):
        """Low temporal coverage should generate a recommendation."""
        info = _make_info(num_episodes=50)
        coverage = _make_coverage(
            temporal=0.3,
            dominant_path_pct=85.0,
        )
        quality = _make_quality(info, coverage)
        recs = generate_recommendations(info, coverage, quality)

        temp_recs = [r for r in recs if r.category == "temporal_coverage"]
        assert len(temp_recs) == 1
        assert "85" in temp_recs[0].description
