AirlockTools/Loxide.py

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


# TODO Continue implementing logger
# TODO Add input sanitation and CSV injection prevention
# TODO Continue OTP and Local approval rewrites
# TODO Explore pywin32
# TODO Fix Requirements.txt
# TODO Create Generic system_config.json for gitea

from collections import Counter
from datetime import UTC, datetime, timedelta
import json
import logging
import os
from typing import Optional

from bson import ObjectId
import dotenv
import plotext as plt
from textual.app import App, ComposeResult
from textual.containers import Horizontal, Vertical, VerticalScroll
from textual.message import Message
from textual.screen import Screen
from textual.widgets import (
    Button,
    DirectoryTree,
    Footer,
    Header,
    Select,
    Static,
    Tab,
    Tabs,
)
import urllib3

import airlock_libs
from models.agent import Agent
from models.policy import Policy
from services.API import AirlockAPIWrapper
from services.security import getAPI
from TUI.Screens.executionhistoryscreen import ExecutionHistoryScreen
from TUI.Screens.moveagentworkflowscreen import MoveAgentWorkflowScreen
from TUI.Screens.otpactivityscreen import OTPActivitiesScreen
from TUI.Screens.otprevokescreen import OTPRevokeScreen
from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen
from TUI.Screens.policyprepworkflowscreen import PolicyPrepWorkflowScreen
from TUI.Screens.quietagentworkflowscreen import QuietAgentWorkflowScreen
from TUI.Widgets.agentmoveoperations import AgentMoveOperations
from TUI.Widgets.multiagentselector import MultiAgentSelector
from TUI.Widgets.policytreewidget import PolicyTreeWidget
from TUI.Widgets.resultsdisplay import ResultsDisplay
from TUI.Widgets.serverlogwidget import ServerLogWidget
from TUI.Widgets.settingswidget import SettingsWidget
from utils.configmanager import (
    get_system_value,
    get_user_value,
    load_env,
    save_user_config,
)
from utils.setup import get_base_directory, setup
from utils.utils import irtang, open_directory

dotenv.load_dotenv()
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

# ---------------------------------------------------------------------------
# GLOBAL STASH
# ---------------------------------------------------------------------------

_APP_RESTART_REASON = None

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Helper for plotext charts
# ---------------------------------------------------------------------------
def create_plotext_chart(
    chart_type: str,
    labels: list,
    values: list,
    height: int = 10,
    width: int = 60,
    title: str = "",
    color: str = None,
) -> Static:
    """Create a plotext chart and return it as a Static widget with the rendered output."""
    # Completely reset plotext state
    plt.clear_data()
    plt.clear_figure()
    plt.clear_color()
    plt.clear_terminal()
    plt.clf()

    # Set theme and canvas size
    plt.theme("clear")
    plt.plotsize(width, height)

    if chart_type == "simple_bar":
        # Simple horizontal bar - clean text-based look
        if color:
            plt.simple_bar(labels, values, width=width, color=color)
        else:
            plt.simple_bar(labels, values, width=width)
    elif chart_type == "bar_h":
        # Calculate integer ticks for bar charts
        max_val = max(values) if values else 0
        if max_val <= 5:
            ticks = list(range(0, max_val + 1))
        elif max_val <= 10:
            ticks = list(range(0, max_val + 1, 2))
        else:
            step = max(1, int(max_val / 5))
            ticks = list(range(0, max_val + step, step))
        plt.bar(labels, values, orientation="h", width=0.5)
        plt.xlabel("Count")
        plt.xticks(ticks)
        plt.grid(False, False)
    elif chart_type == "bar_v":
        max_val = max(values) if values else 0
        if max_val <= 5:
            ticks = list(range(0, max_val + 1))
        elif max_val <= 10:
            ticks = list(range(0, max_val + 1, 2))
        else:
            step = max(1, int(max_val / 5))
            ticks = list(range(0, max_val + step, step))
        plt.bar(labels, values, orientation="v", width=0.5)
        plt.ylabel("Count")
        plt.yticks(ticks)
        plt.grid(False, False)

    if title:
        plt.title(title)

    # Use build() to get the chart string
    chart_output = plt.build()

    # Clear state after building
    plt.clf()
    plt.clear_data()

    # Create Static widget with the chart - disable selection
    chart_widget = Static(chart_output, classes="chart_display")
    chart_widget.can_focus = False

    return chart_widget


# ---------------------------------------------------------------------------
# helper to persist TEXTUAL_THEME to *user* config and mirror to .env
# ---------------------------------------------------------------------------
def _persist_user_theme(theme_name: str) -> None:
    """
    Store the chosen Textual theme in the user's config using the config manager.
    No need to touch .env - config manager handles everything.
    """
    base_dir = get_base_directory()
    config_dir = base_dir / "config"

    try:
        save_user_config(config_dir, {"TEXTUAL_THEME": theme_name})
        logger.debug("Updated user config with TEXTUAL_THEME=%s", theme_name)
    except Exception as exc:
        logger.error("Failed to save TEXTUAL_THEME: %s", exc)


# ---------------------------------------------------------------------------
# 1) SCREEN
# ---------------------------------------------------------------------------
class MainMenuScreen(Screen):
    api: AirlockAPIWrapper

    BUTTON_DEFS = {
        "agent_actions": [
            {
                "label": "🔧 - Multi-Agent Operations",
                "id": "move_agent_workflow_button",
                "description": "Select agents to:\n  • Move policies\n  • Generate OTPs\n  • Toggle audit/enforcement\n  • View history\n  • Export data",
            },
            {
                "label": "🎫 - Review and approve OTP Activities",
                "id": "otp_activities_button",
            },
            {
                "label": "🛑 - Revoke Active OTP Session",
                "id": "otp_revoke_button",
            },
        ],
        "policy": [
            {
                "label": "⚙️ - Prepare Policy For Enforcement",
                "id": "policy_prep_button",
            },
            {
                "label": "🔍 - Find and Move Quiet Hosts to Enforcement",
                "id": "find_quiet_button",
            },
        ],
    }

    def __init__(self) -> None:
        super().__init__()
        self.extras = get_user_value("EXTRAS", str, "NOTTODAY")
        wd = load_env("WORKING_DIR") or os.getcwd()
        if not os.path.isdir(wd):
            wd = os.getcwd()
        self.working_dir = wd

        # Statistics cache: {days: {"data": stats_dict, "timestamp": datetime}}
        self.statistics_cache = {}
        self.current_stats_days = 1  # Default to 1 day

    def _make_buttons_for(self, tab_id: str) -> Vertical:
        defs = self.BUTTON_DEFS.get(tab_id, [])

        # Special layout for agent_actions tab
        if tab_id == "agent_actions":
            left_col = Vertical()
            left_col.styles.width = "1fr"
            right_col = Vertical()
            right_col.styles.width = "1fr"

            for item in defs:
                if isinstance(item, dict):
                    label = item["label"]
                    btn_id = item["id"]
                    description = item.get("description")
                else:
                    label, btn_id = item
                    description = None

                btn = Button(label, id=btn_id)
                btn.styles.width = "100%"
                btn.styles.margin = (0, 1, 1, 0)

                btn_container = Vertical()
                btn_container.styles.height = "auto"
                btn_container.compose_add_child(btn)

                if description:
                    desc_text = Static(description, classes="button_description")
                    desc_text.styles.width = "100%"
                    desc_text.styles.color = "ansi_bright_black"
                    desc_text.styles.text_align = "left"
                    desc_text.styles.margin = (0, 0, 1, 0)
                    btn_container.compose_add_child(desc_text)

                # Multi-Agent on left, OTP buttons on right
                if "otp" in btn_id:
                    right_col.compose_add_child(btn_container)
                else:
                    left_col.compose_add_child(btn_container)

            row = Horizontal()
            row.styles.width = "100%"
            row.styles.height = "auto"
            row.compose_add_child(left_col)
            row.compose_add_child(right_col)

            return Vertical(row)

        # Default single-column layout for other tabs
        widgets = []
        for item in defs:
            if isinstance(item, dict):
                label = item["label"]
                btn_id = item["id"]
                description = item.get("description")
            else:
                label, btn_id = item
                description = None

            btn = Button(label, id=btn_id)
            btn.styles.width = "100%"
            btn.styles.margin = (0, 1, 1, 0)
            widgets.append(btn)

            if description:
                desc_text = Static(description, classes="button_description")
                desc_text.styles.width = "100%"
                desc_text.styles.color = "ansi_bright_black"
                desc_text.styles.text_align = "left"
                desc_text.styles.margin = (0, 0, 1, 0)
                widgets.append(desc_text)

        return Vertical(*widgets)

    def compose(self) -> ComposeResult:
        yield Header(show_clock=True, icon="⚙")

        tabs = [
            Tab("Agents", id="agent_actions"),
            Tab("Tree View", id="p_tree"),
            Tab("Server Log", id="server_log"),
            Tab("Statistics", id="statistics"),
            Tab("Directory", id="dir"),
            Tab("Settings", id="settings"),
        ]

        if self.extras == "POLICYPREP":
            tabs.insert(2, Tab("Policy Prep", id="policy"))

        yield Tabs(*tabs, id="tabs")
        yield Vertical(id="content")
        yield Footer()

    def on_mount(self) -> None:
        self.switch_tab("agent_actions")

        # Delay version check to ensure UI is fully ready for notifications
        if hasattr(self.app, "_version_checker"):
            self.set_timer(1.0, self._do_version_check)

    def _do_version_check(self) -> None:
        """Perform version check and show notification if update available."""
        logger.info("_do_version_check called")

        if not hasattr(self.app, "_version_checker"):
            logger.warning(
                "No _version_checker on app - was create_update_notifier called?"
            )
            return

        checker = self.app._version_checker
        logger.info(f"Checker exists: {checker}")
        logger.info("Calling check_now(force=True)...")

        try:
            result = checker.check_now(force=True)
            logger.info(
                f"Version check result: update_available={result.update_available}, latest={result.latest_version}, current={result.current_version}, error={result.error}"
            )

            if result.error:
                logger.warning(f"Version check returned error: {result.error}")
                return

            if result.update_available:
                # Always nag on startup - don't check dismissed version here
                # User can dismiss from settings but we still want startup reminder
                logger.info(f"Showing toast for update {result.latest_version}")
                msg = f"🆕 Update available: {result.latest_version}\nGo to Settings to download"
                self.app.notify(
                    msg, title="Loxide Update Available", severity="warning", timeout=15
                )
                logger.info("Toast notification called")
            else:
                logger.info(
                    f"No update available - current {result.current_version} is latest"
                )
        except Exception as e:
            logger.error(f"Version check failed: {e}", exc_info=True)

    def on_key(self, event) -> None:
        """Handle up/down arrow keys for button navigation."""
        if event.key == "down":
            self._focus_nearby_button(1)
            event.prevent_default()
            event.stop()
        elif event.key == "up":
            self._focus_nearby_button(-1)
            event.prevent_default()
            event.stop()
        # left/right are handled by Textual's default tab navigation

    # focus helpers
    def _get_content_buttons(self) -> list[Button]:
        content = self.query_one("#content", Vertical)
        return list(content.query(Button))

    def _focus_first_button(self) -> None:
        buttons = self._get_content_buttons()
        if buttons:
            buttons[0].focus()

    def _focus_tabs(self) -> None:
        tabs = self.query_one("#tabs", Tabs)
        tabs.focus()

    def _focus_nearby_button(self, direction: int) -> None:
        buttons = self._get_content_buttons()
        if not buttons:
            return

        try:
            current = next(i for i, b in enumerate(buttons) if b.has_focus)
        except StopIteration:
            if direction > 0:
                buttons[0].focus()
            else:
                buttons[-1].focus()
            return

        if direction < 0 and current == 0:
            self._focus_tabs()
            return

        new_index = current + direction
        if 0 <= new_index < len(buttons):
            buttons[new_index].focus()

    def switch_tab(self, tab_id: str) -> None:
        content = self.query_one("#content", Vertical)
        content.remove_children()

        if tab_id in self.BUTTON_DEFS:
            content.mount(self._make_buttons_for(tab_id))
        elif tab_id == "server_log":
            content.mount(ServerLogWidget(self.app.api))
        elif tab_id == "dir":
            content.mount(DirectoryTree(self.working_dir, id="dir_tree"))
        elif tab_id == "p_tree":
            content.mount(PolicyTreeWidget(self.app.policies, self.app.devices))
        elif tab_id == "statistics":
            content.mount(self._create_statistics_widget())
        elif tab_id == "settings":
            content.mount(SettingsWidget())
        else:
            content.mount(Static(f"Unknown tab: {tab_id}"))

    def _create_statistics_widget(self) -> Vertical:
        """Create the enhanced statistics display widget with time period selection."""
        stats_container = Vertical(id="statistics_container")
        # Build all widgets first, then mount them all at once
        widgets_to_mount = []

        # Header
        header = Static("📊 System Statistics", classes="stats_header")
        header.styles.text_align = "center"
        header.styles.text_style = "bold"
        header.styles.margin = (1, 0, 1, 0)
        widgets_to_mount.append(header)

        # Controls row
        controls = Horizontal()
        controls.styles.height = "auto"
        controls.styles.margin = (0, 2, 1, 2)

        # Time period selector
        time_label = Static("Time Period: ")
        time_label.styles.width = "auto"
        time_label.styles.margin = (0, 1, 0, 0)

        time_select = Select(
            options=[
                ("1 Day", 1),
                ("3 Days", 3),
                ("7 Days", 7),
            ],
            value=self.current_stats_days,
            id="stats_time_select",
        )
        time_select.styles.width = 20

        # Fetch/Refresh button - starts as "Fetch"
        if self.current_stats_days in self.statistics_cache:
            button_label = "🔍„ Refresh"
        else:
            button_label = "📊 Fetch Statistics"

        fetch_btn = Button(button_label, id="stats_refresh_btn", variant="primary")
        fetch_btn.styles.margin = (0, 0, 0, 2)

        # Compose controls
        controls.compose_add_child(time_label)
        controls.compose_add_child(time_select)
        controls.compose_add_child(fetch_btn)
        widgets_to_mount.append(controls)

        # Timestamp and status area
        status_container = Vertical(id="stats_status_area")
        status_container.styles.margin = (0, 2, 0, 2)
        widgets_to_mount.append(status_container)

        # Data display area - scrollable vertical layout
        data_container = VerticalScroll(id="stats_data_area")
        data_container.styles.margin = (0, 2, 0, 2)
        widgets_to_mount.append(data_container)

        # Compose all widgets into the container
        for widget in widgets_to_mount:
            stats_container.compose_add_child(widget)

        # If cached, display it. Otherwise show prompt to fetch
        if self.current_stats_days in self.statistics_cache:
            self.call_after_refresh(self._update_statistics_display)
        else:
            self.call_after_refresh(self._show_fetch_prompt)

        return stats_container

    def _show_fetch_prompt(self) -> None:
        """Show a prompt to click Fetch to load statistics."""
        try:
            status_area = self.query_one("#stats_status_area", Vertical)
            status_area.remove_children()

            prompt = Static("Click 'Fetch Statistics' to load execution data.")
            prompt.styles.text_style = "italic"
            prompt.styles.text_align = "center"
            status_area.mount(prompt)
        except Exception as e:
            logger.warning(f"Could not show fetch prompt: {e}")

    def _skipback(self, days: int) -> ObjectId:
        """Generate a MongoDB ObjectId for a given number of days ago."""
        date_days_ago = datetime.now(UTC) - timedelta(days=days)
        timestamp = int(date_days_ago.timestamp())
        hex_timestamp = format(timestamp, "08x")
        objectid_hex = hex_timestamp + "0000000000000000"
        return ObjectId(objectid_hex)

    def _fetch_execution_statistics(self, days: int) -> None:
        """Fetch execution history and calculate statistics using airlock_libs."""
        try:
            status_area = self.query_one("#stats_status_area", Vertical)
        except Exception as e:
            logger.warning(f"Could not find status area: {e}")
            return

        status_area.remove_children()

        status = Static("⏳ Fetching execution data...")
        status.styles.text_style = "italic"
        status_area.mount(status)

        try:
            # Non-trusted execution types
            exec_types = [1, 2, 3, 6, 7, 12, 13, 14, 15, 16]

            logger.info(
                f"Fetching {days} days of execution history for all policies..."
            )
            logger.debug(f"Calling with exec_types={exec_types}, days={days}")

            # Use pull_policy_exec_histories - pass None for policy_name to get all policies
            execs = airlock_libs.pull_policy_exec_histories(
                self.app.api, None, str(exec_types), days  # None = all policies
            )

            logger.debug(f"Received response, length: {len(execs) if execs else 0}")

            if not execs:
                logger.warning("No execution data returned")
                status_area.remove_children()
                error = Static("No execution data found for the selected time period.")
                error.styles.color = "yellow"
                status_area.mount(error)
                return

            # Parse JSON response
            data = json.loads(execs)
            logger.debug(f"Parsed JSON, keys: {data.keys()}")

            all_executions = data.get("response", {}).get("exechistories", [])
            logger.debug(f"Found {len(all_executions)} executions in response")

            if not all_executions:
                logger.warning("No executions in response")
                status_area.remove_children()
                error = Static("No execution records found.")
                error.styles.color = "yellow"
                status_area.mount(error)
                return

            logger.info(f"Retrieved {len(all_executions)} execution records")

            # Process statistics
            stats = self._process_execution_data(all_executions)

            # Cache the results
            self.statistics_cache[days] = {
                "data": stats,
                "timestamp": datetime.now(UTC),
            }

            # Update display
            self._update_statistics_display()

        except Exception as e:
            logger.error(f"Failed to fetch execution statistics: {e}", exc_info=True)
            status_area.remove_children()
            error = Static(f"❌ Error fetching data: {str(e)}")
            error.styles.color = "red"
            status_area.mount(error)

    def _process_execution_data(self, executions: list) -> dict:
        """Process raw execution data into statistics."""
        # Execution type names
        type_names = {
            0: "Trusted Execution",
            1: "Blocked Execution",
            2: "Untrusted Execution [Audit]",
            3: "Untrusted Execution [OTP]",
            4: "Trusted Path Execution",
            5: "Trusted Publisher Execution",
            6: "Blocklist Execution",
            7: "Blocklist Execution [Audit]",
            8: "Trusted Process Execution",
            9: "Constrained Execution",
            10: "Trusted Metadata Execution",
            11: "Trusted Browser Execution",
            12: "Blocked Browser Execution",
            13: "Untrusted Browser Execution [Audit]",
            14: "Untrusted Browser Execution [OTP]",
            15: "Blocklist Browser Execution [Audit]",
            16: "Blocklist Browser Execution",
            17: "Trusted Installer Execution",
            18: "Trusted Browser Metadata Execution",
        }

        # Categorize execution types
        # "Untrusted" types: contain "Untrusted" in name
        untrusted_types = {2, 3, 13, 14}
        # "Block" types: contain "Block" in name (Blocked, Blocklist)
        block_types = {1, 6, 7, 12, 15, 16}

        # Count execution types
        type_counter = Counter()

        # Separate counters for untrusted vs block
        policy_counter_untrusted = Counter()
        policy_counter_block = Counter()
        hostname_counter_untrusted = Counter()
        hostname_counter_block = Counter()

        for exec_record in executions:
            exec_type = exec_record.get("type", -1)
            type_counter[exec_type] += 1

            policy_name = exec_record.get("policyname", "Unknown")
            hostname = exec_record.get("hostname", "Unknown")

            # Categorize by type
            if exec_type in untrusted_types:
                policy_counter_untrusted[policy_name] += 1
                hostname_counter_untrusted[hostname] += 1
            elif exec_type in block_types:
                policy_counter_block[policy_name] += 1
                hostname_counter_block[hostname] += 1

        # Get top 5 for each category
        top_policies_untrusted = policy_counter_untrusted.most_common(5)
        top_policies_block = policy_counter_block.most_common(5)
        top_hostnames_untrusted = hostname_counter_untrusted.most_common(5)
        top_hostnames_block = hostname_counter_block.most_common(5)

        # Format execution type counts
        type_counts = []
        for exec_type, count in sorted(type_counter.items()):
            type_name = type_names.get(exec_type, f"Unknown Type {exec_type}")
            type_counts.append((type_name, count))

        return {
            "total_executions": len(executions),
            "type_counts": type_counts,
            "top_policies_untrusted": top_policies_untrusted,
            "top_policies_block": top_policies_block,
            "top_hostnames_untrusted": top_hostnames_untrusted,
            "top_hostnames_block": top_hostnames_block,
        }

    def _update_statistics_display(self) -> None:
        """Update the statistics display with cached data using vertical scrollable layout."""
        try:
            data_area = self.query_one("#stats_data_area", VerticalScroll)
            status_area = self.query_one("#stats_status_area", Vertical)
        except Exception as e:
            logger.warning(f"Could not find statistics display areas: {e}")
            return

        # Clear any existing plotext state
        plt.clear_data()
        plt.clear_figure()

        data_area.remove_children()
        status_area.remove_children()

        # Get theme color for charts
        theme_name = (
            self.app._textual_theme
            if hasattr(self.app, "_textual_theme")
            else "textual-dark"
        )
        chart_color = None  # use plotext default

        # Get cached data
        cache_entry = self.statistics_cache.get(self.current_stats_days)
        if not cache_entry:
            status = Static("No data available. Click Refresh to fetch.")
            status.styles.text_style = "italic"
            status_area.mount(status)
            return

        stats = cache_entry["data"]
        timestamp = cache_entry["timestamp"]

        # Timestamp display
        timestamp_str = timestamp.strftime("%Y-%m-%d %H:%M:%S UTC")
        timestamp_label = Static(f"Last updated: {timestamp_str}")
        timestamp_label.styles.text_style = "dim"
        timestamp_label.styles.margin = (0, 0, 1, 0)
        status_area.mount(timestamp_label)

        # Create single scrollable column for better chart display
        main_col = Vertical()
        main_col.styles.width = "100%"
        main_col.styles.height = "auto"

        # === SYSTEM OVERVIEW (First) ===
        total_agents = len(self.app.devices) if self.app.devices else 0
        total_policies = len(self.app.policies) if self.app.policies else 0

        if self.app.devices:
            # status: 0=Offline, 1=Online, 2=Hidden, 3=Safemode
            online_agents = sum(
                1 for agent in self.app.devices if getattr(agent, "status", 0) == 1
            )
            offline_agents = sum(
                1 for agent in self.app.devices if getattr(agent, "status", 0) == 0
            )
            hidden_agents = sum(
                1 for agent in self.app.devices if getattr(agent, "status", 0) == 2
            )
            safemode_agents = sum(
                1 for agent in self.app.devices if getattr(agent, "status", 0) == 3
            )
        else:
            online_agents = offline_agents = hidden_agents = safemode_agents = 0

        system_header = Static("ℹ️  System Overview")
        system_header.styles.text_style = "bold"
        system_header.styles.margin = (0, 0, 1, 0)
        main_col.compose_add_child(system_header)

        # Total counts
        totals_line = Static(
            f"Total Agents: {total_agents}    Total Policies: {total_policies}"
        )
        totals_line.styles.margin = (0, 0, 1, 1)
        main_col.compose_add_child(totals_line)

        # Agent status with plotext bar chart and percentages
        status_data = [
            ("Online", online_agents),
            ("Offline", offline_agents),
            ("Hidden", hidden_agents),
            ("Safemode", safemode_agents),
        ]

        labels = []
        counts = []
        for status_name, count in status_data:
            percentage = int((count / total_agents) * 100) if total_agents > 0 else 0
            labels.append(f"{status_name} ({percentage}%)")
            counts.append(count)

        chart = create_plotext_chart(
            "simple_bar", labels, counts, height=15, width=60, color=chart_color
        )
        chart.styles.margin = (0, 0, 1, 0)
        main_col.compose_add_child(chart)

        # === EXECUTION SUMMARY ===
        total_execs = stats["total_executions"]
        total = Static(f"📊 Total Untrusted Executions: {total_execs}")
        total.styles.text_style = "bold"
        total.styles.margin = (1, 0, 1, 0)
        main_col.compose_add_child(total)

        # Execution type breakdown with percentages
        if stats["type_counts"]:
            type_header = Static("📋 Execution Types")
            type_header.styles.text_style = "bold"
            type_header.styles.margin = (1, 0, 0, 0)
            main_col.compose_add_child(type_header)

            labels = []
            counts = []
            for type_name, count in stats["type_counts"]:
                short_name = type_name.replace("Execution", "Exec").replace(
                    "[Audit]", "[Aud]"
                )
                if len(short_name) > 18:
                    short_name = short_name[:15] + "..."
                percentage = int((count / total_execs) * 100) if total_execs > 0 else 0
                labels.append(f"{short_name} ({percentage}%)")
                counts.append(count)

            chart = create_plotext_chart(
                "simple_bar", labels, counts, height=15, width=60, color=chart_color
            )
            chart.styles.margin = (0, 0, 1, 0)
            main_col.compose_add_child(chart)

        # === TOP POLICIES - Untrusted ===
        if stats["top_policies_untrusted"]:
            policy_header = Static("⚖️  Most Active Policies (Untrusted)")
            policy_header.styles.text_style = "bold"
            policy_header.styles.margin = (1, 0, 0, 0)
            main_col.compose_add_child(policy_header)

            policy_total = sum(count for _, count in stats["top_policies_untrusted"])
            labels = []
            values = []
            for policy_name, count in stats["top_policies_untrusted"]:
                short_name = policy_name[:18] if len(policy_name) > 18 else policy_name
                percentage = (
                    int((count / policy_total) * 100) if policy_total > 0 else 0
                )
                labels.append(f"{short_name} ({percentage}%)")
                values.append(count)
            chart = create_plotext_chart(
                "simple_bar", labels, values, height=12, width=60, color=chart_color
            )
            chart.styles.margin = (0, 0, 1, 0)
            main_col.compose_add_child(chart)

        # === TOP POLICIES - Blocked ===
        if stats["top_policies_block"]:
            policy_header = Static("⚖️  Most Active Policies (Blocked)")
            policy_header.styles.text_style = "bold"
            policy_header.styles.margin = (1, 0, 0, 0)
            main_col.compose_add_child(policy_header)

            policy_total = sum(count for _, count in stats["top_policies_block"])
            labels = []
            values = []
            for policy_name, count in stats["top_policies_block"]:
                short_name = policy_name[:18] if len(policy_name) > 18 else policy_name
                percentage = (
                    int((count / policy_total) * 100) if policy_total > 0 else 0
                )
                labels.append(f"{short_name} ({percentage}%)")
                values.append(count)
            chart = create_plotext_chart(
                "simple_bar", labels, values, height=12, width=60, color=chart_color
            )
            chart.styles.margin = (0, 0, 1, 0)
            main_col.compose_add_child(chart)

        # === TOP MACHINES - Untrusted ===
        if stats["top_hostnames_untrusted"]:
            host_header = Static("🖥️  Most Active Machines (Untrusted)")
            host_header.styles.text_style = "bold"
            host_header.styles.margin = (1, 0, 0, 0)
            main_col.compose_add_child(host_header)

            host_total = sum(count for _, count in stats["top_hostnames_untrusted"])
            labels = []
            values = []
            for hostname, count in stats["top_hostnames_untrusted"]:
                short_name = hostname[:18] if len(hostname) > 18 else hostname
                percentage = int((count / host_total) * 100) if host_total > 0 else 0
                labels.append(f"{short_name} ({percentage}%)")
                values.append(count)
            chart = create_plotext_chart(
                "simple_bar", labels, values, height=12, width=60, color=chart_color
            )
            chart.styles.margin = (0, 0, 1, 0)
            main_col.compose_add_child(chart)

        # === TOP MACHINES - Blocked ===
        if stats["top_hostnames_block"]:
            host_header = Static("🖥️  Most Active Machines (Blocked)")
            host_header.styles.text_style = "bold"
            host_header.styles.margin = (1, 0, 0, 0)
            main_col.compose_add_child(host_header)

            host_total = sum(count for _, count in stats["top_hostnames_block"])
            labels = []
            values = []
            for hostname, count in stats["top_hostnames_block"]:
                short_name = hostname[:18] if len(hostname) > 18 else hostname
                percentage = int((count / host_total) * 100) if host_total > 0 else 0
                labels.append(f"{short_name} ({percentage}%)")
                values.append(count)
            chart = create_plotext_chart(
                "simple_bar", labels, values, height=12, width=60, color=chart_color
            )
            chart.styles.margin = (0, 0, 1, 0)
            main_col.compose_add_child(chart)

        # Mount the main scrollable column
        data_area.mount(main_col)

    def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None:
        self.switch_tab(event.tab.id)

    def on_select_changed(self, event: Select.Changed) -> None:
        """Handle time period selection change."""
        if event.select.id == "stats_time_select":
            new_days = event.value
            if new_days != self.current_stats_days:
                self.current_stats_days = new_days

                # Update button label based on cache status
                try:
                    button = self.query_one("#stats_refresh_btn", Button)
                    if new_days in self.statistics_cache:
                        button.label = "🔍„ Refresh"
                        # Display cached data
                        self._update_statistics_display()
                    else:
                        button.label = "📊 Fetch Statistics"
                        # Show fetch prompt
                        self._show_fetch_prompt()
                except Exception as e:
                    logger.warning(f"Could not update button or display: {e}")

    def on_multi_agent_selector_agents_selected(
        self, message: MultiAgentSelector.AgentsSelected
    ) -> None:
        """Handle selected agents from AgentSelector."""
        global _APP_RESTART_REASON
        selected_agents = message.selected_agents
        logger.info("Selected agents: %s", selected_agents)
        # TODO: Implement actual handling of selected agents
        _APP_RESTART_REASON = ("multi_agent_action", selected_agents)
        self.app.exit()

    def on_settings_widget_theme_selected(
        self, message: SettingsWidget.ThemeSelected
    ) -> None:
        """Handle theme selection from SettingsWidget."""
        global _APP_RESTART_REASON
        _persist_user_theme(message.theme_name)
        _APP_RESTART_REASON = ("restart",)
        self.app.exit()

    def on_agent_move_operations_operation_complete(
        self, message: AgentMoveOperations.OperationComplete
    ) -> None:
        """Handle completion of agent move operation - show results."""
        logger.info(
            "Agent move operation completed: %s, %d successful, %d unsuccessful",
            message.operation,
            len(message.successful),
            len(message.unsuccessful),
        )

        # Format results for display
        successful_text = "\n".join(
            [f"{agent.hostname}" for agent, _ in message.successful]
        )
        unsuccessful_text = "\n".join(
            [f"{agent.hostname}: {error}" for agent, error in message.unsuccessful]
        )

        # Remove the operations widget
        try:
            ops_widget = self.query_one(AgentMoveOperations)
            ops_widget.remove()
        except Exception:
            pass

        # Show results
        self.query_one("#content", Vertical).mount(
            ResultsDisplay(message.operation, successful_text, unsuccessful_text)
        )

    def on_results_display_go_back(self, message: ResultsDisplay.GoBack) -> None:
        """Handle back button from results display."""
        try:
            results_widget = self.query_one(ResultsDisplay)
            results_widget.remove()
        except Exception:
            pass
        # Return to main menu
        self.app.pop_screen()

    def on_policy_tree_widget_view_execution_history(
        self, message: PolicyTreeWidget.ViewExecutionHistory
    ) -> None:
        """Handle request to view execution history for a device from tree view."""
        logger.info("Viewing execution history for device: %s", message.device.hostname)
        self.app.push_screen(ExecutionHistoryScreen([message.device]))
        message.stop()

    def on_policy_tree_widget_generate_otp(
        self, message: PolicyTreeWidget.GenerateOTP
    ) -> None:
        """Handle request to generate OTP for a device from tree view."""
        logger.info("Generating OTP for device: %s", message.device.hostname)
        self.app.push_screen(OTPWorkflowScreen([message.device]))
        message.stop()

    def on_policy_tree_widget_toggle_enforcement(
        self, message: PolicyTreeWidget.ToggleEnforcement
    ) -> None:
        """Handle request to toggle enforcement for a device from tree view."""
        logger.info("Toggling enforcement for device: %s", message.device.hostname)

        try:
            from TUI.Widgets.agentmoveoperations import moveAgentToRelatedPolicy
            from utils.configmanager import get_system_json

            policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")

            # Determine current mode and toggle
            if message.device.groupid in policy_relationship_map:
                # Currently in enforcement, move to audit
                result = moveAgentToRelatedPolicy(self.app.api, message.device, "audit")
                mode = "audit"
            else:
                # Currently in audit, move to enforcement
                result = moveAgentToRelatedPolicy(
                    self.app.api, message.device, "enforcement"
                )
                mode = "enforcement"

            logger.info(f"Successfully toggled {message.device.hostname} to {mode}")

            # Refresh data at the app level
            self.app.refresh_data()

            # Refresh the tree widget with new data
            try:
                tree_widget = self.query_one(PolicyTreeWidget)
                tree_widget.refresh_data(self.app.policies, self.app.devices)
            except:
                pass

        except Exception as e:
            logger.error(
                f"Failed to toggle enforcement for {message.device.hostname}: {e}"
            )
            self.app.bell()

        message.stop()

    def on_directory_tree_file_selected(
        self, event: DirectoryTree.FileSelected
    ) -> None:
        path = event.path
        logger.debug("Directory file selected: %s", path)
        try:
            open_directory(str(path))
        except Exception as exc:
            logger.error("Failed to open %s: %s", path, exc)
            self.app.bell()

    def on_button_pressed(self, event: Button.Pressed) -> None:
        button_id = event.button.id
        logger.debug("Button pressed: %s", button_id)

        # Handle statistics refresh button
        if button_id == "stats_refresh_btn":
            self._fetch_execution_statistics(self.current_stats_days)
            # Update button label to "Refresh" after first fetch
            try:
                button = event.button
                button.label = "🔍„ Refresh"
            except Exception as e:
                logger.warning(f"Could not update button label: {e}")
            event.stop()
            return

        match button_id:
            case "move_agent_workflow_button":
                self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices))
                event.stop()

            case "otp_generate_button":
                self.app.push_screen(OTPWorkflowScreen(self.app.devices))
                event.stop()

            case "find_quiet_button":
                self.app.push_screen(
                    QuietAgentWorkflowScreen(self.app.api, self.app.policies)
                )
                event.stop()
                return

            case "otp_activities_button":
                self.app.push_screen(OTPActivitiesScreen())
                event.stop()
                return

            case "otp_revoke_button":
                self.app.push_screen(OTPRevokeScreen())
                event.stop()
                return

            case "policy_prep_button":
                # Use the new TUI workflow screen instead of legacy
                self.app.push_screen(
                    PolicyPrepWorkflowScreen(self.app.api, self.app.policies)
                )
                event.stop()
                return

            case _:
                self.app.bell()
                logger.warning("Unknown button pressed: %s", button_id)
                return


# ---------------------------------------------------------------------------
# 2) APP
# ---------------------------------------------------------------------------
class Loxide(App[Message]):
    api: AirlockAPIWrapper
    working_dir: str
    policies: Optional[list[Policy]]
    devices: Optional[list[Agent]]

    CSS = """
    #logo {
        width: 100%;
        content-align: center middle;
        text-align: center;
    }
    .chart_display {
        overflow: hidden;
    }
    .chart_display:hover {
        background: transparent;
    }
    #statistics_container {
        height: 1fr;
    }
    #stats_status_area {
        height: auto;
    }
    #stats_data_area {
        height: 1fr;
        scrollbar-gutter: stable;
    }
    """
    BINDINGS = [
        ("q", "quit", "Quit"),
        ("f", "open_fe", "Launch Explorer"),
        ("r", "refresh", "Refresh"),
    ]

    def __init__(self, api: AirlockAPIWrapper):
        self._textual_theme = get_user_value("TEXTUAL_THEME", str, "textual-dark")
        super().__init__()
        self.api = api
        wd = load_env("WORKING_DIR") or os.getcwd()
        if not os.path.isdir(wd):
            wd = os.getcwd()
        self.working_dir = wd
        # Initial data load
        self.refresh_data()

    def refresh_data(self) -> None:
        """Public method to refresh policies and devices from the API."""
        try:
            self.policies = [
                Policy(**row.to_dict())
                for _, row in self.api.policy_find_all().iterrows()
            ]
            self.devices = [
                Agent(**row.to_dict())
                for _, row in self.api.agent_find_all().iterrows()
            ]
            if self.policies and self.devices:
                for agent in self.devices:
                    agent.enrich_with_policies(self.policies)
                logger.debug(
                    f"Enriched {len(self.devices)} agents with policy information"
                )
        except Exception as exc:
            logger.error("Failed to load policies/devices: %s", exc)
            self.policies = None
            self.devices = None

    def on_mount(self, api: AirlockAPIWrapper) -> None:
        self.theme = self._textual_theme
        self.push_screen(MainMenuScreen())

    def action_refresh(self) -> None:
        self.refresh_data()

    def action_quit(self) -> None:
        global _APP_RESTART_REASON
        _APP_RESTART_REASON = None
        self.exit()

    def action_open_fe(self) -> None:
        """Open the working directory in the OS file manager (footer binding)."""
        path_to_open = self.working_dir or os.getcwd()
        try:
            open_directory(path_to_open)
        except Exception as exc:
            logger.error("Failed to open directory %s: %s", path_to_open, exc)
            self.bell()  # optional feedback


# ---------------------------------------------------------------------------
# 3) PUBLIC ENTRYPOINT - Updated to accept attach_notification_handler
# ---------------------------------------------------------------------------
def run_Loxide(api: AirlockAPIWrapper, attach_notification_handler=None) -> None:
    global _APP_RESTART_REASON
    base_dir = get_base_directory()
    env_path = base_dir / ".env"
    dotenv.load_dotenv(dotenv_path=env_path, override=True)

    max_attempts = 5
    attempts = 0

    while attempts < max_attempts:
        attempts += 1
        logger.debug("Starting app loop iteration (attempt %d)", attempts)
        _APP_RESTART_REASON = None
        app = Loxide(api)

        # Attach the notification handler if provided
        if attach_notification_handler:
            attach_notification_handler(app)

        try:
            app.run()
        except SystemExit as exc:
            if exc.code != 0:
                logger.debug("Caught SystemExit from Textual: %s", exc)
                raise

        reason = _APP_RESTART_REASON
        logger.debug("After app.run(), _APP_RESTART_REASON = %r", reason)

        if not reason:
            logger.debug("No restart reason, exiting loop")
            break

        if reason[0] == "restart":
            logger.debug("Restarting app loop")
            continue

        if reason[0] == "multi_agent_action":
            logger.info("Multi-agent action with selected agents: %s", reason[1])
            continue

        logger.error("Unknown restart reason: %r", reason)
        break


# ---------------------------------------------------------------------------
# 4) MAIN FUNCTION - Updated to get and pass attach_notification_handler
# ---------------------------------------------------------------------------
def main():
    irtang()
    # Determine  working directory, setup directory, configure logging, sent env, get API and URL if not already stored
    # setup() now returns a function to attach the notification handler
    attach_notification_handler = setup()
    logger = logging.getLogger(__name__)

    try:
        url = get_system_value("URL")
        username = os.getenv("USERNAME")

        if not url:
            raise ValueError("Missing URL in environment variables.")
        if not username:
            raise ValueError("Missing USERNAME in environment variables.")

        logger.debug(f"Retrieved URL: {url}")
        logger.debug(f"Retrieved Username: {username}")

    except ValueError as e:
        logger.error(f"Configuration error: {e}", exc_info=True)
        raise

    api_key = getAPI(username, "Loxide")
    if api_key is None:
        raise ValueError("API key for Loxide is missing.")

    api = AirlockAPIWrapper(
        base_url=str(url),
        api_key=api_key,
    )
    run_Loxide(api, attach_notification_handler)


if __name__ == "__main__":
    main()