from typing import List, Optional

from textual.app import ComposeResult
from textual.screen import Screen

from models.agent import Agent
from TUI.agentmoveoperations import AgentMoveOperations
from TUI.multiagentselector import MultiAgentSelector
from TUI.resultsdisplay import ResultsDisplay


class MoveAgentWorkflowScreen(Screen):
    """Screen that handles the agent movement workflow."""

    def __init__(self, all_agents: Optional[List[Agent]]):
        super().__init__()
        self.all_agents = all_agents
        self.selected_agents = None

    def compose(self) -> ComposeResult:
        """Start with the multi-agent selector."""
        yield MultiAgentSelector(self.all_agents)

    def on_multi_agent_selector_agents_selected(
        self, message: MultiAgentSelector.AgentsSelected
    ) -> None:
        """Handle selected agents - switch to operations screen."""
        self.selected_agents = message.selected_agents

        # Remove the MultiAgentSelector
        selector = self.query_one(MultiAgentSelector)
        selector.remove()

        # Mount the AgentMoveOperations with the selected Agent objects
        self.mount(AgentMoveOperations(self.selected_agents))

    def on_agent_move_operations_operation_complete(
        self, message: AgentMoveOperations.OperationComplete
    ) -> None:
        """Handle completion of move operation - transition to results screen."""
        # Format successful results
        success_lines = []
        for agent, result in message.successful:
            success_lines.append(f"✓ {agent.hostname}")

        # Format unsuccessful results
        failure_lines = []
        for agent, error in message.unsuccessful:
            failure_lines.append(f"✗ {agent.hostname}: {error}")

        successful_text = "\n".join(success_lines) if success_lines else "(none)"
        unsuccessful_text = "\n".join(failure_lines) if failure_lines else "(none)"

        # Remove the operations widget
        ops_widget = self.query_one(AgentMoveOperations)
        ops_widget.remove()

        # Mount the results display
        self.mount(
            ResultsDisplay(message.operation, successful_text, unsuccessful_text)
        )
