AirlockTools/widgets/multiagentselector.py

import difflib
import re
from typing import List

from textual.containers import Horizontal, Vertical
from textual.css.query import NoMatches
from textual.message import Message
from textual.widget import Widget
from textual.widgets import (
    Button,
    Footer,
    Header,
    SelectionList,
    Static,
    Switch,
    TextArea,
)

from models.agent import Agent


class MultiAgentSelector(Widget):
    class AgentsSelected(Message):
        def __init__(self, selected_agents: List[Agent]):
            super().__init__()
            self.selected_agents = selected_agents

    def __init__(self, all_agents: List[Agent]):
        super().__init__()
        self.all_agents = all_agents
        self._match_type = "exact"

    @property
    def match_type(self):
        return self._match_type

    @match_type.setter
    def match_type(self, value):
        self._match_type = value

    def compose(self):
        yield Header(show_clock=True, icon="⚙")
        title_text = Static("🖧 Multi-Agent Selector", id="selector_title")
        title_text.styles.margin = (0, 0, 0, 1)
        yield title_text

        with Horizontal() as main_layout:
            main_layout.styles.height = "auto"

            # Left side - Input and controls
            with Vertical() as left_pane:
                left_pane.styles.width = "1fr"
                left_pane.styles.height = "auto"

                text_area = TextArea(
                    id="device_input",
                    placeholder="Paste device names here (one per line). Supports wildcards: * and ?",
                )
                text_area.styles.height = 10
                text_area.styles.overflow_y = "auto"
                yield text_area

                with Horizontal(id="switch_search_container") as switch_search:
                    switch = Switch(value=False, id="match_switch")
                    switch.styles.width = "auto"
                    switch.styles.margin = (1, 0, 0, 0)
                    switch.styles.padding = (0, 0, 0, 0)
                    yield switch

                    switch_label = Static("Match: Exact", id="match_switch_label")
                    switch_label.styles.width = "auto"
                    switch_label.styles.margin = (2, 1, 0, 0)
                    yield switch_label

                    search = Button("🔍 Search", id="search_button")
                    search.styles.margin = (1, 0, 0, 0)
                    yield search

                with Horizontal() as select_buttons:
                    select_buttons.styles.margin = (0, 0, 0, 0)

                    select_all_button = Button("✅ Select All", id="select_all")
                    select_all_button.styles.margin = (1, 1, 0, 1)
                    yield select_all_button

                    select_none_button = Button("🚫 Select None", id="select_none")
                    select_none_button.styles.margin = (1, 0, 0, 1)
                    yield select_none_button

                with Horizontal() as button_row:
                    button_row.styles.height = "auto"
                    button_row.styles.margin = (1, 0, 0, 0)

                    back_button = Button("← Back", id="back_button")
                    back_button.styles.width = "1fr"
                    yield back_button

                    submit_button = Button(
                        "▶ Select & Continue", id="submit_selection", variant="primary"
                    )
                    submit_button.styles.margin = (0, 5, 2, 1)
                    submit_button.styles.padding = (0, 6, 0, 0)
                    yield submit_button

            # Right side - Results
            with Vertical() as right_pane:
                right_pane.styles.width = "2fr"
                yield SelectionList(id="match_results")
                yield Static(id="unmatched_label")
        yield Footer()

    def on_switch_changed(self, event: Switch.Changed):
        self.match_type = "fuzzy" if event.value else "exact"
        self.query_one("#match_switch_label", Static).update(
            f"Match: {self.match_type.capitalize()}"
        )

    def on_button_pressed(self, event: Button.Pressed):
        btn_id = event.button.id

        try:
            match_list = self.query_one("#match_results", SelectionList)
        except NoMatches:
            return
        if btn_id == "back_button":
            self.app.pop_screen()
            event.stop()
        elif btn_id == "select_all":
            match_list.select_all()
            event.stop()
        elif btn_id == "select_none":
            match_list.deselect_all()
            event.stop()
        elif btn_id == "submit_selection":
            # Get selected hostnames
            selected_hostnames = list(match_list.selected)
            # Convert back to Agent objects
            selected_agents = [
                agent
                for agent in self.all_agents
                if agent.hostname in selected_hostnames
            ]
            self.post_message(self.AgentsSelected(selected_agents))
            event.stop()
        elif btn_id == "search_button":
            self.update_matches()
            event.stop()

    def update_matches(self):
        raw_input = self.query_one("#device_input", TextArea).text.strip()
        device_names = [line.strip() for line in raw_input.split("\n") if line.strip()]
        matched, unmatched = self.match_devices(device_names)
        match_list = self.query_one("#match_results", SelectionList)
        match_list.clear_options()
        for name in matched:
            match_list.add_option((name, name))
        unmatched_label = self.query_one("#unmatched_label", Static)
        if unmatched:
            unmatched_label.update(f"⚠️ No matches for: {', '.join(unmatched)}")
        else:
            unmatched_label.update("")

    def match_devices(self, device_names: list[str]) -> tuple[list[str], list[str]]:
        if not self.all_agents or not device_names:
            return [], device_names
        agent_names = [agent.hostname for agent in self.all_agents]
        matched = set()
        unmatched = []

        for name in device_names:
            # Check if the name contains wildcards
            has_wildcards = "*" in name or "?" in name

            if has_wildcards:
                # Use regex for wildcard matching
                pattern = re.escape(name)
                pattern = pattern.replace(r"\*", ".*").replace(r"\?", ".")
                regex = re.compile(f"^{pattern}$", re.IGNORECASE)

                wildcard_matches = [
                    agent_name for agent_name in agent_names if regex.match(agent_name)
                ]
                if wildcard_matches:
                    matched.update(wildcard_matches)
                else:
                    unmatched.append(name)
            elif self.match_type == "exact":
                # Case-insensitive exact match
                name_lower = name.lower()
                exact_match = None
                for agent_name in agent_names:
                    if agent_name.lower() == name_lower:
                        exact_match = agent_name
                        break
                if exact_match:
                    matched.add(exact_match)
                else:
                    unmatched.append(name)
            else:
                # Fuzzy match
                matches = difflib.get_close_matches(name, agent_names, n=5, cutoff=0.5)
                if matches:
                    matched.update(matches)
                else:
                    unmatched.append(name)
        return sorted(matched), unmatched