"""Tests for the orbit cleaner module."""

import json
from unittest.mock import MagicMock, patch

import numpy as np
import pandas as pd
import pytest

from orbit.cleaner import CleaningPlan, build_cleaning_plan, print_dry_run


def _make_episode_data(n_episodes=10, ep_length=100, action_dim=6, add_bad=True):
    """Create synthetic episode data for testing."""
    rows = []
    for ep_idx in range(n_episodes):
        for frame_idx in range(ep_length):
            if add_bad and ep_idx == 0:
                # Episode 0: dead frames (no motion)
                action = [0.0] * action_dim
            elif add_bad and ep_idx == 1:
                # Episode 1: very short (will have only 10 frames)
                if frame_idx >= 10:
                    continue
                action = list(np.random.randn(action_dim) * 0.5)
            else:
                action = list(np.random.randn(action_dim) * 0.5)
            state = list(np.random.randn(action_dim) * 0.3)
            rows.append({
                "episode_index": ep_idx,
                "frame_index": frame_idx,
                "action": action,
                "state": state,
            })
    return pd.DataFrame(rows)


def _make_info(n_episodes=10, fps=30.0, action_dim=6):
    """Create a mock DatasetInfo."""
    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 = []
    info.video_keys = []
    info.shapes = {"action": [action_dim], "state": [action_dim]}
    info.duration_seconds = n_episodes * 100 / fps
    return info


class TestBuildCleaningPlan:
    """Test the cleaning plan builder."""

    def test_remove_episodes_classified_as_remove(self):
        """Episodes the ranker classifies as 'remove' should be in the removal list."""
        ep_data = _make_episode_data(n_episodes=10, add_bad=True)
        info = _make_info(n_episodes=10)
        plan = build_cleaning_plan("test/dataset", ep_data, info)
        # Episode 0 has all dead frames, should be flagged
        remove_indices = [ep["index"] for ep in plan.episodes_to_remove]
        assert 0 in remove_indices  # dead frames episode

    def test_aggressive_includes_review_episodes(self):
        """With --aggressive, 'review' episodes should also be removed."""
        ep_data = _make_episode_data(n_episodes=10, add_bad=True)
        info = _make_info(n_episodes=10)
        plan_standard = build_cleaning_plan("test/dataset", ep_data, info, aggressive=False)
        plan_aggressive = build_cleaning_plan("test/dataset", ep_data, info, aggressive=True)
        # Aggressive should remove at least as many as standard
        assert len(plan_aggressive.episodes_to_remove) >= len(plan_standard.episodes_to_remove)

    def test_dry_run_produces_plan_without_modification(self):
        """Building a plan should not modify anything — it's just analysis."""
        ep_data = _make_episode_data(n_episodes=5, add_bad=False)
        info = _make_info(n_episodes=5)
        plan = build_cleaning_plan("test/dataset", ep_data, info)
        assert isinstance(plan, CleaningPlan)
        assert plan.original_episodes == 5

    def test_cleaning_report_json_structure(self):
        """The cleaning plan should have the expected structure."""
        ep_data = _make_episode_data(n_episodes=10, add_bad=True)
        info = _make_info(n_episodes=10)
        plan = build_cleaning_plan("test/dataset", ep_data, info)
        assert plan.original_dataset == "test/dataset"
        assert plan.original_episodes == 10
        assert isinstance(plan.episodes_to_remove, list)
        assert isinstance(plan.dead_joints, list)
        assert plan.episodes_to_keep == plan.original_episodes - len(plan.episodes_to_remove)
        assert plan.cleaning_mode in ("standard", "aggressive")
        assert isinstance(plan.quality_before, float)
        assert isinstance(plan.quality_after, float)

    def test_episode_removal_has_reason(self):
        """Each removed episode should have a reason."""
        ep_data = _make_episode_data(n_episodes=10, add_bad=True)
        info = _make_info(n_episodes=10)
        plan = build_cleaning_plan("test/dataset", ep_data, info)
        for ep in plan.episodes_to_remove:
            assert "index" in ep
            assert "reason" in ep
            assert len(ep["reason"]) > 0

    def test_empty_episode_data(self):
        """Should handle empty data gracefully."""
        ep_data = pd.DataFrame()
        info = _make_info(n_episodes=0)
        plan = build_cleaning_plan("test/dataset", ep_data, info)
        assert plan.episodes_to_keep == 0
        assert len(plan.episodes_to_remove) == 0

    def test_clean_dataset_no_removals(self):
        """A clean dataset should have no episodes to remove."""
        # Create uniform, well-behaved data
        rows = []
        for ep_idx in range(20):
            for frame_idx in range(100):
                t = frame_idx / 100.0
                action = [np.sin(t + i * 0.5) * 0.5 for i in range(6)]
                state = [np.cos(t + i * 0.3) * 0.3 for i in range(6)]
                rows.append({
                    "episode_index": ep_idx,
                    "frame_index": frame_idx,
                    "action": action,
                    "state": state,
                })
        ep_data = pd.DataFrame(rows)
        info = _make_info(n_episodes=20)
        plan = build_cleaning_plan("test/dataset", ep_data, info)
        # With well-behaved data, removals should be minimal
        assert plan.episodes_to_keep >= 15  # at least 75% kept
