#    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/>.

"""
Quiet Agent Workflow Screen Module

Provides a TUI workflow for identifying quiet agents and moving them to target policies.
This screen replaces the legacy quietAgent.py with a comprehensive TUI interface that:
1. Allows selection of an initial policy to analyze
2. Categorizes devices into "Enforce Ready" and "Non-Enforce Ready" based on activity
3. Allows users to select target policies for each category
4. Uses the API to move devices to their target policies
"""

import datetime
import logging
import os
from typing import List, Optional

import pandas as pd
from textual.app import ComposeResult
from textual.containers import Horizontal, Vertical
from textual.reactive import reactive
from textual.screen import Screen
from textual.widgets import Button, DataTable, Footer, Header, Static

from models.policy import Policy
from services.API import AirlockAPIWrapper
from services.policyhandler import getPolicyInfo
from TUI.policyselector import PolicySelector
from utils.configmanager import load_env

logger = logging.getLogger(__name__)


class QuietAgentWorkflowScreen(Screen):
    """
    A Textual screen for the Quiet Agent analysis and migration workflow.

    This screen provides a multi-step workflow:
    1. Select initial policy to analyze
    2. View categorized agents (enforce ready vs. non-enforce ready)
    3. Select target policies for each category
    4. Execute agent migrations

    Attributes:
        api (AirlockAPIWrapper): API wrapper for Airlock operations
        policies (List[Policy]): List of all available policies
        selected_policy (Optional[Policy]): The initially selected policy to analyze
        history_days (int): Number of days of history to pull (default: 150)
        quiet_days (int): Number of days without execution to be considered quiet (default: 45)
        agents_df (Optional[pd.DataFrame]): DataFrame of all agents with analysis results
        enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents ready for enforcement
        non_enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents not ready for enforcement
        workflow_stage (str): Current stage of the workflow
    """

    BINDINGS = [
        ("escape", "go_back", "Back"),
    ]

    workflow_stage = reactive("select_policy")  # Tracks current workflow stage

    def __init__(self, api: AirlockAPIWrapper, policies: List[Policy]):
        """
        Initialize the QuietAgentWorkflowScreen.

        Args:
            api (AirlockAPIWrapper): API wrapper for Airlock operations
            policies (List[Policy]): List of all available policies
        """
        super().__init__()
        self.api = api
        self.policies = policies
        self.selected_policy: Optional[Policy] = None
        self.history_days = 150  # Fixed as per requirements
        self.quiet_days = 45  # Default value
        self.agents_df: Optional[pd.DataFrame] = None
        self.enforce_ready_df: Optional[pd.DataFrame] = None
        self.non_enforce_ready_df: Optional[pd.DataFrame] = None
        self.enforce_ready_target_policy: Optional[Policy] = None
        self.non_enforce_ready_target_policy: Optional[Policy] = None

    def compose(self) -> ComposeResult:
        """Build the UI layout for the workflow screen."""
        # Include Header and Footer like other standalone screens
        yield Header(show_clock=True, icon="⚙")

        # Title area
        title = Static("🔒 Quiet Agent Workflow", id="workflow_title")
        title.styles.margin = (0, 0, 0, 1)
        yield title

        # Status area
        status = Static("Step 1: Select Policy to Analyze", id="workflow_status")
        status.styles.margin = (0, 0, 1, 1)
        yield status

        # Content area - dynamically populated based on workflow stage
        yield Vertical(id="content_area")

        yield Footer()

    def on_mount(self) -> None:
        """Initialize the screen when mounted."""
        # Show initial policy selection
        self._show_policy_selection()

    def watch_workflow_stage(self, old_value: str, new_value: str) -> None:
        """React to workflow stage changes."""
        logger.debug(f"Workflow stage changed from {old_value} to {new_value}")
        self._update_status_message()

    def _update_status_message(self) -> None:
        """Update the status message based on current workflow stage."""
        status_widget = self.query_one("#workflow_status", Static)

        stage_messages = {
            "select_policy": "Step 1: Select Policy to Analyze",
            "select_quiet_days": "Step 2: Select Quiet Time Period",
            "analyzing": "📊 Analyzing agent activity...",
            "view_results": "Step 3: Review Categorized Agents",
            "select_enforce_target": "Step 4: Select Target Policy for Enforce Ready Agents",
            "select_non_enforce_target": "Step 5: Select Target Policy for Non-Enforce Ready Agents",
            "confirm_migration": "Step 6: Confirm and Execute Migration",
            "executing": "⏳ Executing agent migrations...",
            "complete": "✅ Migration Complete",
        }

        status_widget.update(stage_messages.get(self.workflow_stage, "Unknown Stage"))

    def _show_policy_selection(self) -> None:
        """Show the initial policy selection screen."""
        self.workflow_stage = "select_policy"
        content = self.query_one("#content_area", Vertical)
        content.remove_children()

        # Create policy selector widget
        policy_selector = PolicySelector(self.policies)
        content.mount(policy_selector)

    def on_policy_selector_policy_selected(
        self, message: PolicySelector.PolicySelected
    ) -> None:
        """Handle policy selection from PolicySelector widget."""
        # Handle based on current workflow stage
        if self.workflow_stage == "select_policy":
            # Initial policy selection for analysis
            self.selected_policy = message.policy
            logger.info(f"Selected policy for analysis: {self.selected_policy.name}")
            self._show_quiet_days_selection()
        elif self.workflow_stage == "select_enforce_target":
            # Target policy selection for enforce ready agents
            self.enforce_ready_target_policy = message.policy
            logger.info(
                f"Selected target policy for enforce ready: {message.policy.name}"
            )
            self._show_non_enforce_target_selection()
        elif self.workflow_stage == "select_non_enforce_target":
            # Target policy selection for non-enforce ready agents
            self.non_enforce_ready_target_policy = message.policy
            logger.info(
                f"Selected target policy for non-enforce ready: {message.policy.name}"
            )
            self._show_migration_confirmation()

    def _show_quiet_days_selection(self) -> None:
        """Show the quiet days selection screen."""
        self.workflow_stage = "select_quiet_days"
        content = self.query_one("#content_area", Vertical)
        content.remove_children()

        # Create info text
        info_widget = Static(
            f"Policy Selected: {self.selected_policy.name}\n\n"
            f"History Period: {self.history_days} days\n\n"
            "Select quiet time period (days without untrusted execution):",
            id="quiet_days_info",
        )
        info_widget.styles.margin = (0, 0, 2, 0)
        content.mount(info_widget)

        # Create button container and mount it first
        button_container = Vertical(id="quiet_days_buttons")
        button_container.styles.height = "auto"
        content.mount(button_container)

        # Now add buttons to the mounted container
        for days in [15, 30, 45, 60]:
            btn = Button(
                f"{days} days {'(Default)' if days == 45 else ''}",
                id=f"quiet_days_{days}",
                classes="quiet_day_btn",
            )
            btn.styles.width = "100%"
            btn.styles.margin = (0, 0, 1, 0)
            button_container.mount(btn)

        back_btn = Button("← Back", id="back_to_policy_selection")
        back_btn.styles.width = "100%"
        back_btn.styles.margin = (2, 0, 0, 0)
        button_container.mount(back_btn)

    def on_button_pressed(self, event: Button.Pressed) -> None:
        """Handle button press events."""
        button_id = event.button.id

        # Quiet days selection buttons
        if button_id and button_id.startswith("quiet_days_"):
            days = int(button_id.split("_")[-1])
            self.quiet_days = days
            logger.info(f"Selected quiet days: {days}")
            self._start_analysis()
            return

        # Navigation buttons
        if button_id == "back_to_policy_selection":
            self._show_policy_selection()
            return

        if button_id == "back_to_results":
            self._show_results()
            return

        if button_id == "select_enforce_target_btn":
            self._show_enforce_target_selection()
            return

        if button_id == "select_non_enforce_target_btn":
            self._show_non_enforce_target_selection()
            return

        if button_id == "skip_enforce_target_btn":
            # Skip enforce ready target selection
            self.enforce_ready_target_policy = None
            self._show_non_enforce_target_selection()
            return

        if button_id == "skip_non_enforce_target_btn":
            # Skip non-enforce ready target selection
            self.non_enforce_ready_target_policy = None
            self._show_migration_confirmation()
            return

        if button_id == "confirm_migration_btn":
            self._execute_migration()
            return

        if button_id == "cancel_migration_btn":
            self._show_results()
            return

        if button_id == "export_results_btn":
            self._export_results()
            return

        if button_id == "start_over_btn":
            self._show_policy_selection()
            return

    def _start_analysis(self) -> None:
        """Start the agent activity analysis."""
        self.workflow_stage = "analyzing"
        content = self.query_one("#content_area", Vertical)
        content.remove_children()

        # Show analyzing message with detailed steps
        analyzing_msg = Static(
            f"📊 Analyzing Agent Activity\n"
            f"{'=' * 50}\n\n"
            f"Policy: {self.selected_policy.name}\n"
            f"History Period: {self.history_days} days\n"
            f"Quiet Threshold: {self.quiet_days} days\n\n"
            f"Progress:\n"
            f"⏳ Step 1/4: Fetching agents from policy...\n"
            f"⏱️  Step 2/4: Pulling execution history (this may take a moment)...\n"
            f"⏱️  Step 3/4: Analyzing activity patterns...\n"
            f"⏱️  Step 4/4: Categorizing agents...\n\n"
            f"Please wait - this operation cannot be cancelled.",
            id="analyzing_message",
        )
        analyzing_msg.styles.margin = (2, 1)
        content.mount(analyzing_msg)

        # Show notification
        self.app.notify(
            "Starting analysis - this may take several minutes for large policies",
            severity="information",
            timeout=5,
        )

        # Perform the analysis asynchronously
        self.call_later(self._perform_analysis)

    def _perform_analysis(self) -> None:
        """Perform the actual agent activity analysis."""
        try:
            # Update status: Fetching agents
            self._update_analysis_status("Step 1/4: Fetching agents from policy...")

            # Get agents in the selected policy
            agents = self.api.agents_find_by_group(self.selected_policy.groupid)

            if agents.empty:
                self.app.notify(
                    f"No agents found in policy: {self.selected_policy.name}",
                    severity="warning",
                    timeout=5,
                )
                self._show_policy_selection()
                return

            agent_count = len(agents)
            self.app.notify(
                f"Found {agent_count} agents - fetching execution history...",
                severity="information",
                timeout=3,
            )

            # Update status: Pulling execution history
            self._update_analysis_status(
                f"Step 2/4: Pulling execution history for {agent_count} agents...\n"
                f"(This may take several minutes - progress shown in terminal)"
            )

            # Get execution history (this shows progress bars in terminal via airlock_libs)
            policy_exec_history = getPolicyInfo(
                self.api, self.selected_policy, [1, 2, 6, 7], self.history_days
            )

            # Update status: Analyzing patterns
            self._update_analysis_status("Step 3/4: Analyzing activity patterns...")
            self.app.notify(
                "History retrieved - analyzing patterns...",
                severity="information",
                timeout=2,
            )

            if policy_exec_history.empty:
                logger.info(
                    "No execution history found for the selected policy and time range."
                )
                # All agents are quiet (no executions)
                agents["execution_count"] = 0
                agents["days_since"] = None
                agents["required_quiet"] = self.quiet_days
                agents["enforce_ready"] = True
            else:
                # Convert datetime column
                policy_exec_history["datetime"] = pd.to_datetime(
                    policy_exec_history["datetime"],
                    format="%Y-%m-%dT%H:%M:%SZ",
                    utc=True,
                )

                # Calculate days ago
                now = datetime.datetime.now(datetime.timezone.utc)
                policy_exec_history["days_ago"] = policy_exec_history["datetime"].apply(
                    lambda dt: (now - dt).days
                )

                # Count total executions per hostname
                hostname_counts = policy_exec_history["hostname"].value_counts()
                agents["execution_count"] = (
                    agents["hostname"].map(hostname_counts).fillna(0).astype(int)
                )

                # Find most recent execution per hostname
                most_recent_exec = policy_exec_history.sort_values(
                    by="days_ago"
                ).drop_duplicates(subset="hostname", keep="first")

                # Map most recent execution age to agents
                agents["days_since"] = agents["hostname"].map(
                    most_recent_exec.set_index("hostname")["days_ago"]
                )

                # Check for enforcement readiness
                agents["required_quiet"] = self.quiet_days
                agents["enforce_ready"] = agents["days_since"].apply(
                    lambda x: True if pd.isna(x) or x > self.quiet_days else False
                )

            # Update status: Categorizing
            self._update_analysis_status("Step 4/4: Categorizing agents...")

            # Sort agents
            agents = agents.sort_values(
                by=["execution_count", "hostname"], ascending=[True, True]
            )

            # Store the results
            self.agents_df = agents

            # Categorize agents into DataFrames
            self.enforce_ready_df = agents[agents["enforce_ready"] == True].copy()
            self.non_enforce_ready_df = agents[agents["enforce_ready"] == False].copy()

            logger.info(
                f"Analysis complete: {len(self.enforce_ready_df)} enforce ready, "
                f"{len(self.non_enforce_ready_df)} non-enforce ready"
            )

            self.app.notify(
                f"Analysis complete! Found {len(self.enforce_ready_df)} enforce ready, "
                f"{len(self.non_enforce_ready_df)} not ready",
                severity="success",
                timeout=5,
            )

            # Show results
            self._show_results()

        except Exception as e:
            logger.error(f"Error during analysis: {e}", exc_info=True)
            self.app.notify(f"Analysis failed: {str(e)}", severity="error", timeout=5)
            self._show_policy_selection()

    def _update_analysis_status(self, status_text: str) -> None:
        """Update the analysis status message."""
        try:
            analyzing_msg = self.query_one("#analyzing_message", Static)

            # Build updated message
            updated_text = (
                f"📊 Analyzing Agent Activity\n"
                f"{'=' * 50}\n\n"
                f"Policy: {self.selected_policy.name}\n"
                f"History Period: {self.history_days} days\n"
                f"Quiet Threshold: {self.quiet_days} days\n\n"
                f"Progress:\n"
                f"✅ {status_text}\n\n"
                f"Please wait - this operation cannot be cancelled."
            )

            analyzing_msg.update(updated_text)
        except Exception as e:
            logger.debug(f"Could not update analysis status: {e}")

    def _show_results(self) -> None:
        """Show the categorized results."""
        self.workflow_stage = "view_results"
        content = self.query_one("#content_area", Vertical)
        content.remove_children()

        # Create results display container and mount it first
        results_container = Vertical(id="results_container")
        results_container.styles.height = "auto"
        results_container.styles.margin = (1, 1)
        content.mount(results_container)

        # Summary statistics
        total_agents = len(self.enforce_ready_df) + len(self.non_enforce_ready_df)
        ready_count = len(self.enforce_ready_df)
        not_ready_count = len(self.non_enforce_ready_df)
        ready_percentage = (ready_count / total_agents * 100) if total_agents > 0 else 0

        summary = Static(
            f"Analysis Results for: {self.selected_policy.name}\n\n"
            f"📊 Total Agents: {total_agents}\n"
            f"✅ Enforce Ready: {ready_count} ({ready_percentage:.1f}%)\n"
            f"❌ Not Ready: {not_ready_count} ({100 - ready_percentage:.1f}%)\n\n"
            f"Quiet Threshold: {self.quiet_days} days\n"
            f"History Period: {self.history_days} days",
            id="results_summary",
        )
        summary.styles.margin = (0, 0, 2, 0)
        results_container.mount(summary)

        # Action buttons
        button_container = Horizontal(id="results_buttons")
        button_container.styles.height = "auto"
        results_container.mount(button_container)

        if ready_count > 0:
            enforce_btn = Button(
                f"Select Target for Enforce Ready ({ready_count})",
                id="select_enforce_target_btn",
            )
            enforce_btn.styles.margin = (0, 1, 1, 0)
            button_container.mount(enforce_btn)

        if not_ready_count > 0:
            non_enforce_btn = Button(
                f"Select Target for Non-Enforce Ready ({not_ready_count})",
                id="select_non_enforce_target_btn",
            )
            non_enforce_btn.styles.margin = (0, 1, 1, 0)
            button_container.mount(non_enforce_btn)

        export_btn = Button("💾 Export Results", id="export_results_btn")
        export_btn.styles.margin = (0, 1, 1, 0)
        button_container.mount(export_btn)

        start_over_btn = Button("🔄 Start Over", id="start_over_btn")
        start_over_btn.styles.margin = (0, 0, 1, 0)
        button_container.mount(start_over_btn)

        # Tables showing agents
        tables_container = Horizontal()
        tables_container.styles.height = "1fr"
        results_container.mount(tables_container)

        # Enforce Ready table
        if ready_count > 0:
            enforce_col = Vertical()
            enforce_col.styles.width = "1fr"
            enforce_col.styles.margin = (1, 1, 0, 0)
            tables_container.mount(enforce_col)

            enforce_label = Static("✅ Enforce Ready Agents")
            enforce_label.styles.margin = (0, 0, 1, 0)
            enforce_col.mount(enforce_label)

            enforce_table = DataTable(id="enforce_ready_table")
            enforce_table.styles.height = "1fr"
            enforce_table.add_columns("Hostname", "Last Exec (days)")

            # Display first 50 agents
            for idx, row in self.enforce_ready_df.head(50).iterrows():
                days_since = row["days_since"]
                days_str = f"{int(days_since)}" if not pd.isna(days_since) else "Never"
                enforce_table.add_row(row["hostname"], days_str)

            if len(self.enforce_ready_df) > 50:
                enforce_table.add_row(
                    f"... and {len(self.enforce_ready_df) - 50} more", ""
                )

            enforce_col.mount(enforce_table)

        # Non-Enforce Ready table
        if not_ready_count > 0:
            non_enforce_col = Vertical()
            non_enforce_col.styles.width = "1fr"
            non_enforce_col.styles.margin = (1, 0, 0, 1)
            tables_container.mount(non_enforce_col)

            non_enforce_label = Static("❌ Non-Enforce Ready Agents")
            non_enforce_label.styles.margin = (0, 0, 1, 0)
            non_enforce_col.mount(non_enforce_label)

            non_enforce_table = DataTable(id="non_enforce_ready_table")
            non_enforce_table.styles.height = "1fr"
            non_enforce_table.add_columns("Hostname", "Last Exec (days)")

            # Display first 50 agents
            for idx, row in self.non_enforce_ready_df.head(50).iterrows():
                days_since = row["days_since"]
                days_str = f"{int(days_since)}" if not pd.isna(days_since) else "N/A"
                non_enforce_table.add_row(row["hostname"], days_str)

            if len(self.non_enforce_ready_df) > 50:
                non_enforce_table.add_row(
                    f"... and {len(self.non_enforce_ready_df) - 50} more", ""
                )

            non_enforce_col.mount(non_enforce_table)

    def _show_enforce_target_selection(self) -> None:
        """Show policy selection for enforce ready agents."""
        self.workflow_stage = "select_enforce_target"
        content = self.query_one("#content_area", Vertical)
        content.remove_children()

        # Info message
        info = Static(
            f"Select target policy for {len(self.enforce_ready_df)} Enforce Ready agents\n"
            f"Source Policy: {self.selected_policy.name}",
            id="enforce_target_info",
        )
        info.styles.margin = (0, 0, 2, 0)
        content.mount(info)

        # Policy selector
        policy_selector = PolicySelector(self.policies)
        content.mount(policy_selector)

        # Skip button
        skip_btn = Button("⭕️ Skip - No Migration", id="skip_enforce_target_btn")
        skip_btn.styles.width = "50%"
        skip_btn.styles.margin = (2, 0, 0, 0)
        content.mount(skip_btn)

    def _show_non_enforce_target_selection(self) -> None:
        """Show policy selection for non-enforce ready agents."""
        self.workflow_stage = "select_non_enforce_target"
        content = self.query_one("#content_area", Vertical)
        content.remove_children()

        # Info message
        info = Static(
            f"Select target policy for {len(self.non_enforce_ready_df)} Non-Enforce Ready agents\n"
            f"Source Policy: {self.selected_policy.name}",
            id="non_enforce_target_info",
        )
        info.styles.margin = (0, 0, 2, 0)
        content.mount(info)

        # Policy selector
        policy_selector = PolicySelector(self.policies)
        content.mount(policy_selector)

        # Skip button
        skip_btn = Button("⭕️ Skip - No Migration", id="skip_non_enforce_target_btn")
        skip_btn.styles.width = "50%"
        skip_btn.styles.margin = (2, 0, 0, 0)
        content.mount(skip_btn)

    def _show_migration_confirmation(self) -> None:
        """Show migration confirmation screen."""
        self.workflow_stage = "confirm_migration"
        content = self.query_one("#content_area", Vertical)
        content.remove_children()

        # Build confirmation message
        confirmation_lines = [
            "🔐 Migration Summary\n",
            f"Source Policy: {self.selected_policy.name}\n",
        ]

        if self.enforce_ready_target_policy:
            confirmation_lines.append(
                f"\n✅ Enforce Ready Migration:\n"
                f"  • Agents: {len(self.enforce_ready_df)}\n"
                f"  • Target: {self.enforce_ready_target_policy.name}\n"
            )

        if self.non_enforce_ready_target_policy:
            confirmation_lines.append(
                f"\n❌ Non-Enforce Ready Migration:\n"
                f"  • Agents: {len(self.non_enforce_ready_df)}\n"
                f"  • Target: {self.non_enforce_ready_target_policy.name}\n"
            )

        if (
            not self.enforce_ready_target_policy
            and not self.non_enforce_ready_target_policy
        ):
            confirmation_lines.append("\n⚠️ No migrations will be performed.")

        confirmation = Static("".join(confirmation_lines), id="migration_confirmation")
        confirmation.styles.margin = (1, 1, 2, 1)
        content.mount(confirmation)

        # Action buttons - mount container first, then add buttons
        button_container = Horizontal(id="confirmation_buttons")
        button_container.styles.height = "auto"
        button_container.styles.margin = (1, 1)
        content.mount(button_container)

        if self.enforce_ready_target_policy or self.non_enforce_ready_target_policy:
            confirm_btn = Button("✅ Confirm Migration", id="confirm_migration_btn")
            confirm_btn.styles.margin = (0, 1, 0, 0)
            button_container.mount(confirm_btn)

        cancel_btn = Button("❌ Cancel", id="cancel_migration_btn")
        button_container.mount(cancel_btn)

    def _execute_migration(self) -> None:
        """Execute the agent migrations."""
        self.workflow_stage = "executing"
        content = self.query_one("#content_area", Vertical)
        content.remove_children()

        # Show executing message
        executing_msg = Static(
            "⏳ Executing agent migrations...\nPlease wait...",
            id="executing_message",
        )
        executing_msg.styles.margin = (2, 1)
        content.mount(executing_msg)

        # Perform migrations asynchronously
        self.call_later(self._perform_migrations)

    def _perform_migrations(self) -> None:
        """Perform the actual agent migrations."""
        successful_migrations = []
        failed_migrations = []

        try:
            # Migrate enforce ready agents
            if self.enforce_ready_target_policy:
                for idx, row in self.enforce_ready_df.iterrows():
                    try:
                        result = self.api.agent_move(
                            row["agentid"], self.enforce_ready_target_policy.groupid
                        )
                        successful_migrations.append(
                            (row["hostname"], self.enforce_ready_target_policy.name)
                        )
                        logger.debug(
                            f"Moved {row['hostname']} to {self.enforce_ready_target_policy.name}"
                        )
                    except Exception as e:
                        failed_migrations.append((row["hostname"], str(e)))
                        logger.error(f"Failed to move {row['hostname']}: {e}")

            # Migrate non-enforce ready agents
            if self.non_enforce_ready_target_policy:
                for idx, row in self.non_enforce_ready_df.iterrows():
                    try:
                        result = self.api.agent_move(
                            row["agentid"], self.non_enforce_ready_target_policy.groupid
                        )
                        successful_migrations.append(
                            (row["hostname"], self.non_enforce_ready_target_policy.name)
                        )
                        logger.debug(
                            f"Moved {row['hostname']} to {self.non_enforce_ready_target_policy.name}"
                        )
                    except Exception as e:
                        failed_migrations.append((row["hostname"], str(e)))
                        logger.error(f"Failed to move {row['hostname']}: {e}")

            # Show completion results
            self._show_completion_results(successful_migrations, failed_migrations)

        except Exception as e:
            logger.error(f"Error during migration execution: {e}", exc_info=True)
            self.app.notify(f"Migration failed: {str(e)}", severity="error", timeout=5)
            self._show_results()

    def _show_completion_results(
        self, successful: List[tuple], failed: List[tuple]
    ) -> None:
        """Show migration completion results."""
        self.workflow_stage = "complete"
        content = self.query_one("#content_area", Vertical)
        content.remove_children()

        # Results summary
        total_attempted = len(successful) + len(failed)
        success_rate = (
            (len(successful) / total_attempted * 100) if total_attempted > 0 else 0
        )

        results = Static(
            f"✅ Migration Complete\n\n"
            f"Total Agents Migrated: {len(successful)}\n"
            f"Failed Migrations: {len(failed)}\n"
            f"Success Rate: {success_rate:.1f}%",
            id="completion_summary",
        )
        results.styles.margin = (1, 1, 2, 1)
        content.mount(results)

        # Details tables
        if successful:
            success_container = Vertical()
            success_container.styles.margin = (0, 1)
            content.mount(success_container)

            success_label = Static("✅ Successful Migrations")
            success_label.styles.margin = (0, 0, 1, 0)
            success_container.mount(success_label)

            success_table = DataTable(id="success_table")
            success_table.styles.height = "auto"
            success_table.add_columns("Hostname", "Target Policy")

            for hostname, target_policy in successful[:25]:  # Show first 25
                success_table.add_row(hostname, target_policy)

            if len(successful) > 25:
                success_table.add_row(f"... and {len(successful) - 25} more", "")

            success_container.mount(success_table)

        if failed:
            failed_container = Vertical()
            failed_container.styles.margin = (2, 1, 0, 1)
            content.mount(failed_container)

            failed_label = Static("❌ Failed Migrations")
            failed_label.styles.margin = (0, 0, 1, 0)
            failed_container.mount(failed_label)

            failed_table = DataTable(id="failed_table")
            failed_table.styles.height = "auto"
            failed_table.add_columns("Hostname", "Error")

            for hostname, error in failed[:25]:  # Show first 25
                failed_table.add_row(hostname, error[:50])  # Truncate error

            if len(failed) > 25:
                failed_table.add_row(f"... and {len(failed) - 25} more", "")

            failed_container.mount(failed_table)

        # Action button
        done_btn = Button("✔ Done", id="start_over_btn")
        done_btn.styles.width = "50%"
        done_btn.styles.margin = (2, 0, 0, 0)
        content.mount(done_btn)

    def _export_results(self) -> None:
        """Export analysis results to CSV."""
        try:
            working_dir = load_env("WORKING_DIR") or os.getcwd()
            filename = os.path.join(
                working_dir,
                f"{self.selected_policy.name}_quiet_analysis_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
            )

            self.agents_df.to_csv(filename, index=False)
            logger.info(f"Exported results to {filename}")
            self.app.notify(
                f"Results exported to:\n{filename}",
                severity="information",
                timeout=5,
            )

        except Exception as e:
            logger.error(f"Failed to export results: {e}")
            self.app.notify(f"Export failed: {str(e)}", severity="error", timeout=5)

    def action_go_back(self) -> None:
        """Handle back/escape action."""
        # Depending on stage, go back to previous stage or exit
        if self.workflow_stage in ["select_policy", "view_results", "complete"]:
            self.app.pop_screen()
        elif self.workflow_stage == "select_quiet_days":
            self._show_policy_selection()
        elif self.workflow_stage == "select_enforce_target":
            self._show_results()
        elif self.workflow_stage == "select_non_enforce_target":
            if self.enforce_ready_target_policy:
                self._show_enforce_target_selection()
            else:
                self._show_results()
        elif self.workflow_stage == "confirm_migration":
            self._show_non_enforce_target_selection()
        else:
            self.app.pop_screen()