"""Full pipeline test: analyze -> plan -> track with mocked data."""

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

_MOCK_INFO = DatasetInfo(
    repo_id="lerobot/test_dataset",
    num_episodes=40,
    num_frames=8000,
    fps=30.0,
    robot_type="so100",
    camera_names=["cam_top"],
    video_keys=["observation.images.cam_top"],
    shapes={"action": [6]},
    duration_seconds=8000 / 30.0,
)


class TestFullPipeline:
    """End-to-end pipeline test using mocked data."""

    @patch("orbit.analyzer.dataset_loader.load_dataset_info")
    def test_analyze_then_plan_then_track(self, mock_load, tmp_path):
        mock_load.return_value = _MOCK_INFO
        runner = CliRunner()

        # Step 1: Analyze
        result = runner.invoke(
            main,
            ["analyze", "lerobot/test_dataset", "--json", "--skip-embeddings"],
        )
        assert result.exit_code == 0
        analysis = json.loads(result.output)
        assert analysis["repo_id"] == "lerobot/test_dataset"
        assert 0 <= analysis["quality_score"] <= 1
        assert "recommendations" in analysis
        assert "estimated_success_rate" in analysis

        # Step 2: Plan
        plan_path = str(tmp_path / "plan.json")
        result = runner.invoke(
            main,
            ["plan", "pick up cup", "--robot", "so100", "--policy", "act", "--save", plan_path],
        )
        assert result.exit_code == 0

        # Verify plan file
        with open(plan_path) as f:
            plan_data = json.load(f)
        assert plan_data["task_type"] == "pick_and_place"
        assert plan_data["robot_type"] == "so100"
        assert plan_data["policy_type"] == "act"
        assert plan_data["total_target_episodes"] > 0
        assert len(plan_data["phases"]) >= 2

        # Step 3: Track
        result = runner.invoke(
            main,
            ["track", plan_path, "--dataset", "lerobot/test_dataset"],
        )
        assert result.exit_code == 0
        assert "Progress Tracker" in result.output
        assert "Coverage Checklist" in result.output
        assert "Progress Tracker" in result.output  # already checked above

    def test_plan_all_robots(self):
        """Verify plan command works with all supported robot types."""
        runner = CliRunner()
        for robot in ["so100", "so101", "koch", "aloha"]:
            result = runner.invoke(
                main,
                ["plan", "pick up cup", "--robot", robot, "--policy", "act", "--format", "json"],
            )
            assert result.exit_code == 0, f"Failed for robot={robot}"
            data = json.loads(result.output)
            assert data["robot_type"] == robot

    def test_plan_all_policies(self):
        """Verify plan command works with all supported policy types."""
        runner = CliRunner()
        for policy in ["act", "diffusion_policy", "smolvla"]:
            result = runner.invoke(
                main,
                ["plan", "pick up cup", "--robot", "so100", "--policy", policy, "--format", "json"],
            )
            assert result.exit_code == 0, f"Failed for policy={policy}"
            data = json.loads(result.output)
            assert data["policy_type"] == policy

    def test_plan_all_task_templates(self):
        """Verify each task template is matched correctly."""
        runner = CliRunner()
        template_map = {
            "pick up the cup from the table": "pick_and_place",
            "insert the peg into the socket": "insertion",
            "wipe the table surface clean": "wiping",
            "pour water from the bottle": "pouring",
            "stack blocks on top of each other": "stacking",
            "do something completely novel": "generic",
        }
        for task, expected_type in template_map.items():
            result = runner.invoke(
                main,
                ["plan", task, "--robot", "so100", "--policy", "act", "--format", "json"],
            )
            assert result.exit_code == 0, f"Failed for task='{task}'"
            data = json.loads(result.output)
            assert data["task_type"] == expected_type, (
                f"Task '{task}' matched '{data['task_type']}' instead of '{expected_type}'"
            )

    @patch("orbit.analyzer.dataset_loader.load_dataset_info")
    def test_analyze_validates_output_schema(self, mock_load):
        """Verify JSON output has all required fields."""
        mock_load.return_value = _MOCK_INFO
        runner = CliRunner()
        result = runner.invoke(
            main,
            ["analyze", "lerobot/test_dataset", "--json", "--skip-embeddings"],
        )
        assert result.exit_code == 0
        data = json.loads(result.output)

        required_fields = [
            "repo_id", "num_episodes", "num_frames", "fps", "robot_type",
            "camera_names", "coverage", "quality_score", "component_scores",
            "ready_to_train", "ready_to_train_verdict", "confidence",
            "estimated_success_rate", "recommendations",
        ]
        for field in required_fields:
            assert field in data, f"Missing field: {field}"

        # Verify coverage sub-fields
        for key in ["position_diversity", "action_diversity",
                     "episode_consistency", "temporal_coverage"]:
            assert key in data["coverage"], f"Missing coverage field: {key}"
            assert 0 <= data["coverage"][key] <= 1

    @patch("orbit.analyzer.dataset_loader.load_dataset_info")
    def test_track_plan_roundtrip(self, mock_load, tmp_path):
        """Verify plan save/load roundtrip produces valid tracking output."""
        mock_load.return_value = _MOCK_INFO
        runner = CliRunner()

        plan_path = str(tmp_path / "roundtrip.json")

        # Save plan
        runner.invoke(
            main,
            ["plan", "stack blocks", "--robot", "so100", "--policy", "act", "--save", plan_path],
        )

        # Load and verify saved JSON
        with open(plan_path) as f:
            saved = json.load(f)
        assert "phases" in saved
        assert "estimates" in saved
        assert "robot_config" in saved

        # Track against it
        result = runner.invoke(
            main, ["track", plan_path, "--dataset", "lerobot/test_dataset"]
        )
        assert result.exit_code == 0
        assert "Episode Progress" in result.output
