"""Tests for orbit.cli."""

from __future__ import annotations

import json
from unittest.mock import patch

from click.testing import CliRunner

from orbit.analyzer.dataset_loader import DatasetInfo
from orbit.cli import main


def test_orbit_help():
    runner = CliRunner()
    result = runner.invoke(main, ["--help"])
    assert result.exit_code == 0
    assert "analyze" in result.output
    assert "plan" in result.output
    assert "track" in result.output


def test_analyze_help():
    runner = CliRunner()
    result = runner.invoke(main, ["analyze", "--help"])
    assert result.exit_code == 0
    assert "Analyze robot training data" in result.output


def test_plan_help():
    runner = CliRunner()
    result = runner.invoke(main, ["plan", "--help"])
    assert result.exit_code == 0
    assert "Plan data collection strategy" in result.output


def test_track_help():
    runner = CliRunner()
    result = runner.invoke(main, ["track", "--help"])
    assert result.exit_code == 0
    assert "Track training progress" in result.output


import pandas as pd


def _empty_episode_df():
    return pd.DataFrame(
        columns=["episode_index", "frame_index", "action", "state", "timestamp"]
    )


_FAKE_INFO = DatasetInfo(
    repo_id="lerobot/aloha_mobile_cabinet",
    num_episodes=50,
    num_frames=12847,
    fps=30.0,
    robot_type="aloha",
    camera_names=["cam_high", "cam_low"],
    video_keys=["observation.images.cam_high", "observation.images.cam_low"],
    shapes={
        "observation.images.cam_high": [480, 640, 3],
        "observation.images.cam_low": [480, 640, 3],
        "action": [14],
    },
    duration_seconds=12847 / 30.0,
)


class TestAnalyzeCommand:
    """Tests for the analyze command."""

    @patch("orbit.analyzer.dataset_loader.load_episode_data")
    @patch("orbit.analyzer.dataset_loader.load_dataset_info")
    def test_analyze_rich_output(self, mock_load, mock_ep):
        mock_load.return_value = _FAKE_INFO
        mock_ep.return_value = _empty_episode_df()
        runner = CliRunner()
        result = runner.invoke(
            main, ["analyze", "lerobot/aloha_mobile_cabinet", "--skip-embeddings"]
        )
        assert result.exit_code == 0
        assert "lerobot/aloha_mobile_cabinet" in result.output
        assert "50" in result.output
        # Core sections still present
        assert "ORBIT Analysis" in result.output or "aloha" in result.output.lower()

    @patch("orbit.analyzer.dataset_loader.load_episode_data")
    @patch("orbit.analyzer.dataset_loader.load_dataset_info")
    def test_analyze_json_output(self, mock_load, mock_ep):
        mock_load.return_value = _FAKE_INFO
        mock_ep.return_value = _empty_episode_df()
        runner = CliRunner()
        result = runner.invoke(
            main,
            ["analyze", "lerobot/aloha_mobile_cabinet", "--json", "--skip-embeddings"],
        )
        assert result.exit_code == 0
        data = json.loads(result.output)
        assert data["repo_id"] == "lerobot/aloha_mobile_cabinet"
        assert data["num_episodes"] == 50
        assert data["num_frames"] == 12847
        assert data["fps"] == 30.0
        assert data["robot_type"] == "aloha"
        assert data["camera_names"] == ["cam_high", "cam_low"]
        assert "duration_human" in data

    @patch("orbit.analyzer.dataset_loader.load_dataset_info")
    def test_analyze_not_found(self, mock_load):
        mock_load.side_effect = ValueError(
            "Dataset 'fake/nonexistent' not found on HuggingFace Hub"
        )
        runner = CliRunner()
        result = runner.invoke(main, ["analyze", "fake/nonexistent", "--skip-embeddings"])
        assert result.exit_code != 0
        assert "Error" in result.output
        assert "not found" in result.output

    @patch("orbit.analyzer.dataset_loader.load_dataset_info")
    def test_analyze_not_found_json(self, mock_load):
        mock_load.side_effect = ValueError(
            "Dataset 'fake/nonexistent' not found on HuggingFace Hub"
        )
        runner = CliRunner()
        result = runner.invoke(
            main, ["analyze", "fake/nonexistent", "--json", "--skip-embeddings"]
        )
        assert result.exit_code != 0
        data = json.loads(result.output)
        assert "error" in data

    @patch("orbit.analyzer.dataset_loader.load_dataset_info")
    def test_analyze_connection_error(self, mock_load):
        mock_load.side_effect = ConnectionError("Network error")
        runner = CliRunner()
        result = runner.invoke(
            main, ["analyze", "lerobot/aloha_mobile_cabinet", "--skip-embeddings"]
        )
        assert result.exit_code != 0
        assert "Error" in result.output

    @patch("orbit.analyzer.dataset_loader.load_episode_data")
    @patch("orbit.analyzer.dataset_loader.load_dataset_info")
    def test_analyze_episodes_flag(self, mock_load, mock_ep):
        mock_load.return_value = _FAKE_INFO
        mock_ep.return_value = _empty_episode_df()
        runner = CliRunner()
        result = runner.invoke(
            main,
            ["analyze", "lerobot/aloha_mobile_cabinet", "--episodes", "10", "--skip-embeddings"],
        )
        assert result.exit_code == 0

    def test_analyze_empty_repo_id(self):
        runner = CliRunner()
        result = runner.invoke(main, ["analyze", "", "--skip-embeddings"])
        assert result.exit_code != 0
        assert "empty" in result.output.lower()

    def test_analyze_incomplete_repo_id(self):
        runner = CliRunner()
        result = runner.invoke(main, ["analyze", "lerobot/", "--skip-embeddings"])
        assert result.exit_code != 0
        assert "error" in result.output.lower()

    @patch("orbit.analyzer.dataset_loader.load_dataset_info")
    def test_analyze_skip_coverage(self, mock_load):
        mock_load.return_value = _FAKE_INFO
        runner = CliRunner()
        result = runner.invoke(
            main,
            ["analyze", "lerobot/aloha_mobile_cabinet", "--skip-embeddings", "--skip-coverage"],
        )
        assert result.exit_code == 0
        assert "ORBIT Analysis" in result.output or "Action Quality" in result.output

    @patch("orbit.analyzer.dataset_loader.load_episode_data")
    @patch("orbit.analyzer.dataset_loader.load_dataset_info")
    def test_analyze_json_has_coverage(self, mock_load, mock_ep):
        mock_load.return_value = _FAKE_INFO
        mock_ep.return_value = _empty_episode_df()
        runner = CliRunner()
        result = runner.invoke(
            main,
            ["analyze", "lerobot/aloha_mobile_cabinet", "--json", "--skip-embeddings"],
        )
        assert result.exit_code == 0
        data = json.loads(result.output)
        assert "coverage" in data
        assert "position_diversity" in data["coverage"]
        assert "action_diversity" in data["coverage"]
        assert "quality_score" in data
        assert "ready_to_train" in data
        assert "recommendations" in data
        assert "estimated_success_rate" in data


class TestPlanCommand:
    """Tests for the plan command."""

    def test_plan_pick_and_place(self):
        runner = CliRunner()
        result = runner.invoke(
            main, ["plan", "pick up cups", "--robot", "so100", "--policy", "act"]
        )
        assert result.exit_code == 0
        assert "pick up cups" in result.output
        assert "pick_and_place" in result.output
        assert "so100" in result.output
        assert "act" in result.output

    def test_plan_insertion(self):
        runner = CliRunner()
        result = runner.invoke(
            main,
            ["plan", "insert peg into hole", "--robot", "koch", "--policy", "diffusion_policy"],
        )
        assert result.exit_code == 0
        assert "insertion" in result.output
        assert "koch" in result.output

    def test_plan_wiping(self):
        runner = CliRunner()
        result = runner.invoke(
            main, ["plan", "wipe table with cloth", "--robot", "aloha", "--policy", "act"]
        )
        assert result.exit_code == 0
        assert "wiping" in result.output
        assert "aloha" in result.output

    def test_plan_unknown_uses_generic(self):
        runner = CliRunner()
        result = runner.invoke(
            main, ["plan", "unknown task", "--robot", "so100", "--policy", "act"]
        )
        assert result.exit_code == 0
        assert "generic" in result.output

    def test_plan_save_json(self, tmp_path):
        save_path = str(tmp_path / "plan.json")
        runner = CliRunner()
        result = runner.invoke(
            main,
            ["plan", "stack blocks", "--robot", "so100", "--policy", "act", "--save", save_path],
        )
        assert result.exit_code == 0
        assert "Plan saved to" in result.output
        with open(save_path) as f:
            data = json.load(f)
        assert data["task_description"] == "stack blocks"
        assert data["task_type"] == "stacking"

    def test_plan_json_format(self):
        runner = CliRunner()
        result = runner.invoke(
            main,
            ["plan", "pick up cups", "--robot", "so100", "--policy", "act", "--format", "json"],
        )
        assert result.exit_code == 0
        data = json.loads(result.output)
        assert data["task_description"] == "pick up cups"
        assert "phases" in data


class TestTrackCommand:
    """Tests for the track command."""

    @patch("orbit.analyzer.dataset_loader.load_dataset_info")
    def test_track_with_plan(self, mock_load, tmp_path):
        mock_load.return_value = _FAKE_INFO
        # Save a plan first
        runner = CliRunner()
        save_path = str(tmp_path / "plan.json")
        runner.invoke(
            main,
            ["plan", "pick up cups", "--robot", "so100", "--policy", "act", "--save", save_path],
        )
        # Track against it
        result = runner.invoke(
            main, ["track", save_path, "--dataset", "lerobot/aloha_mobile_cabinet"]
        )
        assert result.exit_code == 0
        assert "Progress Tracker" in result.output
        assert "Episode Progress" in result.output
        assert "Coverage Checklist" in result.output

    def test_track_missing_plan(self):
        runner = CliRunner()
        result = runner.invoke(
            main, ["track", "/nonexistent/plan.json", "--dataset", "lerobot/pusht"]
        )
        assert result.exit_code != 0
        assert "not found" in result.output.lower()


class TestDisplayHelpers:
    """Tests for display utility functions."""

    def test_format_duration_seconds(self):
        from orbit.utils.display import format_duration

        assert format_duration(30.0) == "30.0 sec"

    def test_format_duration_minutes(self):
        from orbit.utils.display import format_duration

        assert format_duration(428.2) == "7.1 min"

    def test_format_duration_hours(self):
        from orbit.utils.display import format_duration

        assert format_duration(4320.0) == "1.2 hours"
