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, 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):
        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")

    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