"""Test custom prompts for parallel agents in CAI CLI."""

import pytest
from unittest.mock import MagicMock, patch
from cai.repl.commands.parallel import ParallelCommand, PARALLEL_CONFIGS, ParallelConfig
from rich.console import Console


class TestParallelCustomPrompts:
    """Test suite for parallel agent custom prompts."""

    def setup_method(self):
        """Set up test environment before each test."""
        # Clear any existing configurations
        PARALLEL_CONFIGS.clear()
        self.console = Console()
        self.command = ParallelCommand()

    def teardown_method(self):
        """Clean up after each test."""
        PARALLEL_CONFIGS.clear()

    def test_prompt_subcommand_adds_prompt_to_config(self):
        """Test that the prompt subcommand correctly adds a custom prompt to a config."""
        # Add an agent first
        with patch("cai.repl.commands._parallel_monolith.console"):
            self.command.handle_add(["redteam_agent"])

        # Verify agent was added
        assert len(PARALLEL_CONFIGS) == 1
        assert PARALLEL_CONFIGS[0].prompt is None

        # Set a custom prompt
        with patch("cai.repl.commands._parallel_monolith.console") as mock_console:
            result = self.command.handle_prompt(["P1", "Focus on SQL injection vulnerabilities"])

        assert result is True
        assert PARALLEL_CONFIGS[0].prompt == "Focus on SQL injection vulnerabilities"

        # Verify success message was printed
        mock_console.print.assert_any_call(
            "[green]Updated prompt for Red Team Agent (ID: P1)[/green]"
        )

    def test_prompt_subcommand_with_index(self):
        """Test that the prompt subcommand works with numeric index."""
        # Add an agent
        with patch("cai.repl.commands._parallel_monolith.console"):
            self.command.handle_add(["bug_bounter_agent"])

        # Set prompt using index
        with patch("cai.repl.commands._parallel_monolith.console"):
            result = self.command.handle_prompt(["1", "Test for XSS vulnerabilities"])

        assert result is True
        assert PARALLEL_CONFIGS[0].prompt == "Test for XSS vulnerabilities"

    def test_prompt_subcommand_error_handling(self):
        """Test error handling for invalid prompt commands."""
        # Test with no arguments
        with patch("cai.repl.commands._parallel_monolith.console") as mock_console:
            result = self.command.handle_prompt([])

        assert result is False
        mock_console.print.assert_any_call("[red]Error: Agent ID/index and prompt required[/red]")

        # Test with invalid ID
        with patch("cai.repl.commands._parallel_monolith.console") as mock_console:
            result = self.command.handle_prompt(["P99", "Some prompt"])

        assert result is False
        mock_console.print.assert_any_call("[red]Error: No agent found with ID/index 'P99'[/red]")

    def test_custom_prompt_displayed_in_list(self):
        """Test that custom prompts are displayed in the list command."""
        # Add agents with prompts
        config1 = ParallelConfig("redteam_agent", prompt="Focus on authentication bypass")
        config1.id = "P1"
        config2 = ParallelConfig(
            "bug_bounter_agent", prompt="Look for IDOR vulnerabilities in the API endpoints"
        )
        config2.id = "P2"
        PARALLEL_CONFIGS.extend([config1, config2])

        # Mock the table print to capture output
        with patch("cai.repl.commands._parallel_monolith.Table") as mock_table:
            with patch("cai.repl.commands._parallel_monolith.console"):
                self.command.handle_list()

            # Verify table was created with correct columns
            mock_table.assert_called_once()
            table_instance = mock_table.return_value

            # Verify add_row was called for each config
            assert table_instance.add_row.call_count == 2

            # Check first row
            first_call = table_instance.add_row.call_args_list[0]
            args = first_call[0]
            assert args[6] == "Focus on authentication bypass"  # Custom prompt column

            # Check second row (should be truncated)
            second_call = table_instance.add_row.call_args_list[1]
            args = second_call[0]
            assert args[6] == "Look for IDOR vulnerabilities in the ..."  # Truncated prompt

    def test_custom_prompt_in_status_display(self):
        """Test that custom prompts are shown in the status display."""
        # Add agent with prompt
        config = ParallelConfig("dfir_agent", prompt="Analyze memory dumps for malware artifacts")
        config.id = "P1"
        PARALLEL_CONFIGS.append(config)

        with patch("cai.repl.commands._parallel_monolith.console") as mock_console:
            self.command.handle_no_args()

        # Verify that prompt info is included in status
        # We need to look through all the print calls to find the Panel
        panel_found = False
        for call in mock_console.print.call_args_list:
            if call[0]:  # Check if arguments exist
                arg = call[0][0]
                # Check if it's a Panel object
                if hasattr(arg, "__class__") and arg.__class__.__name__ == "Panel":
                    # Check the renderable content
                    if hasattr(arg, "renderable"):
                        content = str(arg.renderable)
                        if "Prompt: Analyze memory dumps for malware artifacts" in content:
                            panel_found = True
                            break

        assert panel_found, "Prompt not found in status display"

    def test_parallel_execution_uses_custom_prompts(self):
        """Test that parallel execution correctly uses custom prompts instead of user input."""
        # This test would require mocking the actual parallel execution in cli.py
        # For now, we just verify the configuration is set up correctly

        config1 = ParallelConfig("redteam_agent", prompt="Custom prompt 1")
        config1.id = "P1"
        config2 = ParallelConfig("bug_bounter_agent", prompt="Custom prompt 2")
        config2.id = "P2"
        config3 = ParallelConfig("dfir_agent")  # No custom prompt
        config3.id = "P3"

        PARALLEL_CONFIGS.extend([config1, config2, config3])

        # Verify each config has the correct prompt
        assert PARALLEL_CONFIGS[0].prompt == "Custom prompt 1"
        assert PARALLEL_CONFIGS[1].prompt == "Custom prompt 2"
        assert PARALLEL_CONFIGS[2].prompt is None

    def test_parallel_history_persistence_on_interrupt(self):
        """Test that parallel agents' histories are saved when interrupted."""
        # This test verifies the configuration for history persistence
        from cai.sdk.agents.parallel_isolation import PARALLEL_ISOLATION

        # Setup parallel configs
        config1 = ParallelConfig("redteam_agent")
        config1.id = "P1"
        config2 = ParallelConfig("bug_bounter_agent")
        config2.id = "P2"

        PARALLEL_CONFIGS.extend([config1, config2])

        # Simulate parallel mode
        PARALLEL_ISOLATION._parallel_mode = True

        # Add some test history
        test_history1 = [{"role": "user", "content": "Test message 1"}]
        test_history2 = [{"role": "user", "content": "Test message 2"}]

        PARALLEL_ISOLATION.replace_isolated_history("P1", test_history1)
        PARALLEL_ISOLATION.replace_isolated_history("P2", test_history2)

        # Verify histories are stored
        assert PARALLEL_ISOLATION.get_isolated_history("P1") == test_history1
        assert PARALLEL_ISOLATION.get_isolated_history("P2") == test_history2

        # Clean up
        PARALLEL_ISOLATION.clear_all_histories()
        PARALLEL_ISOLATION._parallel_mode = False

    def test_prompt_update_overwrites_existing(self):
        """Test that updating a prompt overwrites the existing one."""
        # Add agent with initial prompt
        config = ParallelConfig("redteam_agent", prompt="Initial prompt")
        config.id = "P1"
        PARALLEL_CONFIGS.append(config)

        # Update the prompt
        with patch("cai.repl.commands._parallel_monolith.console") as mock_console:
            self.command.handle_prompt(["P1", "Updated prompt with new instructions"])

        assert PARALLEL_CONFIGS[0].prompt == "Updated prompt with new instructions"

        # Verify old prompt was shown
        old_prompt_found = False
        for call in mock_console.print.call_args_list:
            if call[0] and "[dim]Old prompt: Initial prompt[/dim]" in str(call[0][0]):
                old_prompt_found = True
                break

        assert old_prompt_found, "Old prompt message not found"
