#    Copyright (C) 2025 James Brotosky, Brandon Wickline
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as published
#    by the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.
from typing import List, Optional

from textual.app import ComposeResult
from textual.binding import Binding
from textual.css.query import NoMatches
from textual.screen import Screen

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


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

    BINDINGS = [
        Binding("escape", "go_back", "Back"),
        Binding("q", "main_menu", "Main Menu"),
    ]

    def __init__(self, all_agents: Optional[List[Agent]]):
        super().__init__()
        self.all_agents = all_agents
        self.selected_agents = None
        self.workflow_stage = "select_agents"  # Track current stage

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

    def action_go_back(self) -> None:
        """Handle escape key to go back one step within the workflow."""
        if self.workflow_stage == "select_agents":
            # At first stage, go back to main menu
            self.app.pop_screen()
        elif self.workflow_stage == "operations":
            # Go back to agent selection
            try:
                ops_widget = self.query_one(AgentMoveOperations)
                ops_widget.remove()
            except NoMatches:
                pass
            self.mount(MultiAgentSelector(self.all_agents))
            self.workflow_stage = "select_agents"
        elif self.workflow_stage == "results":
            # Go back to operations
            try:
                results_widget = self.query_one(ResultsDisplay)
                results_widget.remove()
            except NoMatches:
                pass
            self.mount(AgentMoveOperations(self.selected_agents))
            self.workflow_stage = "operations"

    def action_main_menu(self) -> None:
        """Handle q key to go back to main menu."""
        while len(self.app.screen_stack) > 2:
            self.app.pop_screen()

    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))
        self.workflow_stage = "operations"

    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)
        )
        self.workflow_stage = "results"