"""Tests for orbit.planner.playbook."""

from __future__ import annotations

import json
import os
import tempfile

from click.testing import CliRunner

from orbit.cli import main
from orbit.planner.playbook import (
    generate_playbook,
)


class TestKeywordMatching:
    """Test that task descriptions match the correct template."""

    def test_pick_and_place_keywords(self):
        playbook = generate_playbook("pick up mugs", "so100", "act")
        assert playbook.task_type == "pick_and_place"

    def test_pick_and_place_grab(self):
        playbook = generate_playbook("grab objects from table", "so100", "act")
        assert playbook.task_type == "pick_and_place"

    def test_pick_and_place_lift(self):
        playbook = generate_playbook("lift blocks and place them", "so100", "act")
        assert playbook.task_type == "pick_and_place"

    def test_insertion_keywords(self):
        playbook = generate_playbook("insert peg into hole", "so100", "act")
        assert playbook.task_type == "insertion"

    def test_insertion_plug(self):
        playbook = generate_playbook("plug USB cable into socket", "so100", "act")
        assert playbook.task_type == "insertion"

    def test_wiping_keywords(self):
        playbook = generate_playbook("wipe the table clean", "so100", "act")
        assert playbook.task_type == "wiping"

    def test_wiping_scrub(self):
        playbook = generate_playbook("scrub the surface", "so100", "act")
        assert playbook.task_type == "wiping"

    def test_pouring_keywords(self):
        playbook = generate_playbook("pour water into cup", "so100", "act")
        assert playbook.task_type == "pouring"

    def test_pouring_fill(self):
        playbook = generate_playbook("fill the bowl with liquid", "so100", "act")
        assert playbook.task_type == "pouring"

    def test_stacking_keywords(self):
        playbook = generate_playbook("stack blocks into tower", "so100", "act")
        assert playbook.task_type == "stacking"

    def test_stacking_pile(self):
        playbook = generate_playbook("pile up cubes", "so100", "act")
        assert playbook.task_type == "stacking"

    def test_generic_fallback(self):
        playbook = generate_playbook("do something completely random", "so100", "act")
        assert playbook.task_type == "generic"

    def test_generic_no_keywords(self):
        playbook = generate_playbook("calibrate the sensors", "so100", "act")
        assert playbook.task_type == "generic"


class TestRobotCustomization:
    """Test robot-specific adjustments."""

    def test_so100_config_loaded(self):
        playbook = generate_playbook("pick up mugs", "so100", "act")
        assert playbook.robot_config is not None
        assert playbook.robot_config.robot_type == "so100"
        assert playbook.robot_config.arm_count == 1

    def test_aloha_bimanual_episodes_scaled(self):
        single = generate_playbook("pick up mugs", "so100", "act")
        bimanual = generate_playbook("pick up mugs", "aloha", "act")
        # Aloha (bimanual) should have more target episodes
        assert bimanual.total_target_episodes > single.total_target_episodes

    def test_aloha_bimanual_requirement_added(self):
        playbook = generate_playbook("pick up mugs", "aloha", "act")
        # Should have bimanual requirement added
        for phase in playbook.phases:
            assert "bimanual" in phase.requirements

    def test_unknown_robot_no_config(self):
        playbook = generate_playbook("pick up mugs", "unknown_bot", "act")
        assert playbook.robot_config is None
        assert playbook.robot_type == "unknown_bot"

    def test_koch_config(self):
        playbook = generate_playbook("pick up mugs", "koch", "act")
        assert playbook.robot_config is not None
        assert playbook.robot_config.robot_type == "koch"


class TestPolicyAdjustments:
    """Test policy-specific adjustments."""

    def test_act_minimum_episodes(self):
        playbook = generate_playbook("pick up mugs", "so100", "act")
        for phase in playbook.phases:
            assert phase.target_episodes >= 50
        assert any("ACT" in n for n in playbook.policy_notes)

    def test_diffusion_policy_more_episodes(self):
        act = generate_playbook("pick up mugs", "so100", "act")
        diff = generate_playbook("pick up mugs", "so100", "diffusion_policy")
        assert diff.total_target_episodes > act.total_target_episodes

    def test_diffusion_policy_notes(self):
        playbook = generate_playbook("pick up mugs", "so100", "diffusion_policy")
        assert any("Diffusion" in n for n in playbook.policy_notes)

    def test_smolvla_notes(self):
        playbook = generate_playbook("pick up mugs", "so100", "smolvla")
        assert any("SmolVLA" in n for n in playbook.policy_notes)

    def test_unknown_policy(self):
        playbook = generate_playbook("pick up mugs", "so100", "my_custom_policy")
        assert any("my_custom_policy" in n for n in playbook.policy_notes)


class TestPlaybookStructure:
    """Test overall playbook structure."""

    def test_has_phases(self):
        playbook = generate_playbook("pick up mugs", "so100", "act")
        assert len(playbook.phases) >= 2

    def test_has_estimates(self):
        playbook = generate_playbook("pick up mugs", "so100", "act")
        assert playbook.estimates.eighty_percent_success
        assert playbook.estimates.ninety_five_percent_success

    def test_total_episodes_sum(self):
        playbook = generate_playbook("pick up mugs", "so100", "act")
        phase_sum = sum(p.target_episodes for p in playbook.phases)
        assert playbook.total_target_episodes == phase_sum

    def test_to_dict(self):
        playbook = generate_playbook("pick up mugs", "so100", "act")
        d = playbook.to_dict()
        assert d["task_description"] == "pick up mugs"
        assert d["task_type"] == "pick_and_place"
        assert d["robot_type"] == "so100"
        assert d["policy_type"] == "act"
        assert len(d["phases"]) >= 2

    def test_save_creates_valid_json(self):
        playbook = generate_playbook("pick up mugs", "so100", "act")
        with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f:
            path = f.name
        try:
            playbook.save(path)
            with open(path) as f:
                data = json.load(f)
            assert data["task_type"] == "pick_and_place"
            assert data["total_target_episodes"] > 0
            assert len(data["phases"]) >= 2
        finally:
            os.unlink(path)


class TestPlanCLI:
    """Test the orbit plan CLI command."""

    def test_plan_basic_output(self):
        runner = CliRunner()
        result = runner.invoke(main, [
            "plan", "pick up mugs", "--robot", "so100", "--policy", "act",
        ])
        assert result.exit_code == 0
        assert "ORBIT Data Collection Playbook" in result.output
        assert "pick up mugs" in result.output
        assert "pick_and_place" in result.output

    def test_plan_save_flag(self):
        runner = CliRunner()
        with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f:
            path = f.name
        try:
            result = runner.invoke(main, [
                "plan", "pick up mugs", "--robot", "so100", "--policy", "act",
                "--save", path,
            ])
            assert result.exit_code == 0
            assert "Plan saved" in result.output
            with open(path) as f:
                data = json.load(f)
            assert data["task_type"] == "pick_and_place"
        finally:
            os.unlink(path)

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

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