agent_m/agentm/views/home.py

from textual.screen import Screen
from textual.widgets import Static, Button
from textual.containers import Vertical, Horizontal, HorizontalScroll
from textual.message import Message
from textual.reactive import reactive
from textual.widget import Widget
from rich.panel import Panel
from rich.console import Group
from PIL import Image, ImageEnhance, ImageFilter
import os

from rich_pixels import Pixels
from rich_pixels._renderer import HalfcellRenderer

from agentm.utils.logger import log_with_caller
from agentm.logic.roms import get_verified_roms, GAME_FILES


class GameSelected(Message):
    def __init__(self, sender: Widget, metadata: dict):
        self.metadata = metadata
        super().__init__(sender)


class ProgressWidget(Widget):
    message = reactive("Loading...")

    def render(self) -> str:
        return f"[bold cyan]{self.message}[/bold cyan]"


class GameAccordion(Static):
    def __init__(self, title: str, rom_file: str, metadata: dict, parent_view):
        self.title = title
        self.rom_file = rom_file
        self.metadata = metadata
        self.parent_view = parent_view
        self.safe_id = rom_file.replace(".", "_").replace("-", "_")

        image_path = os.path.abspath(self.metadata.get("image_path", ""))
        self.image_renderable = self.load_image_scaled(image_path)

        super().__init__(id=f"accordion_{self.safe_id}", classes="game_card")

    def load_image_scaled(self, path: str):
        try:
            with Image.open(path) as img:
                if img.mode != "RGBA":
                    img = img.convert("RGBA")
                scale_factor = 0.10
                target_width = int(img.width * scale_factor)
                target_height = int(img.height * scale_factor)
                resized = img.resize(
                    (target_width, target_height),
                    resample=Image.Resampling.LANCZOS
                )
                resized = resized.filter(ImageFilter.UnsharpMask(radius=1, percent=150, threshold=3))
                self._render_width = target_width
                self._render_height = target_height // 2
                return Pixels.from_image(
                    resized,
                    renderer=HalfcellRenderer(default_color="black"),
                )
        except Exception as e:
            self._render_width = 24
            self._render_height = 16
            return f"[red]Failed to load image[/red]\n[dim]{e}]"

    def render(self):
        return Panel(
            Group(
                self.image_renderable,
                "",  # spacer
                f"[center][bold orange1][u]{self.title.upper()}[/u][/bold orange1][/center]\n",
            ),
            border_style="bright_magenta",
            padding=(1, 2),
            expand=False
        )

    async def on_click(self):
        await self.display_info()
        self.parent_view.highlight_selected(self)

    async def display_info(self):
        meta = self.metadata
        log_with_caller("debug", f"Showing shared info for {self.rom_file}")
        info_lines = [
            f"[b]Title:[/b] {meta['title']}",
            f"[b]Game ID:[/b] {meta['game_id']}",
            f"[b]Difficulty:[/b] {meta.get('difficulty_min')} - {meta.get('difficulty_max')}",
            f"[b]Characters:[/b] {', '.join(meta.get('characters', []))}",
            f"[b]Keywords:[/b] {', '.join(meta.get('keywords', []))}",
            f"[b]SHA256:[/b] {meta['sha256']}",
        ]
        self.parent_view.shared_info_box.update("\n".join(info_lines))
        self.parent_view.shared_confirm_button.label = f"✅ Confirm {meta['title']}"
        self.parent_view.selected_game = meta


class HomeView(Screen):
    BINDINGS = [("escape", "app.quit", "Quit")]

    def highlight_selected(self, selected_widget: GameAccordion):
        for card in self.rom_scroll_row.children:
            if isinstance(card, GameAccordion):
                card.remove_class("game_card_clicked")
        selected_widget.add_class("game_card_clicked")

    def compose(self):
        self.logo = Static(
            "[b bright_magenta]\n\n"
            " █████╗   ██████╗  ███████╗ ███╗  ██╗ ████████╗   ███╗   ███╗\n"
            "██╔══██╗ ██╔════╝  ██╔════╝ ████╗ ██║ ╚══██╔══╝   ████╗ ████║\n"
            "███████║ ██║  ███╗ █████╗   ██╔██╗██║    ██║      ██╔████╔██║\n"
            "██╔══██║ ██║   ██║ ██╔══╝   ██║╚████║    ██║      ██║╚██╔╝██║\n"
            "██║  ██║ ╚██████╔╝ ███████╗ ██║ ╚███║    ██║      ██║     ██║\n"
            "╚═╝  ╚═╝  ╚═════╝  ╚══════╝ ╚═╝  ╚══╝    ╚═╝      ╚═╝     ╚═╝\n[/]",
            classes="header",
            expand=False,
        )
        self.progress_text = ProgressWidget()

        yield Vertical(
            self.logo,
            Static("[bold green]LOADING...[/bold green]", expand=True),
            self.progress_text,
            id="loading_container"
        )

    async def on_mount(self):
        log_with_caller("debug", "HomeView mounted. Starting ROM verification.")
        self.selected_game = None
        self.run_worker(self.run_verification, thread=True, exclusive=True, name="rom-verification")

    def run_verification(self):
        total = len(GAME_FILES)
        verified_roms = get_verified_roms()

        for idx, rom in enumerate(verified_roms, start=1):
            self.app.call_from_thread(
                lambda title=rom['title'], idx=idx: setattr(
                    self.progress_text, "message",
                    f"Processing {title} ({idx}/{total})"
                )
            )
            import time
            time.sleep(0.01)

        self.app.call_from_thread(lambda: self.display_verified_roms(verified_roms))

    async def display_verified_roms(self, verified_roms):
        log_with_caller("info", f"ROM verification complete. Total: {len(verified_roms)}")
        self.query_one("#loading_container").remove()

        self.shared_info_box = Static("", id="game_info_box", classes="game_info")
        self.shared_confirm_button = Button("✅ Confirm", id="confirm_button", classes="confirm_button")

        self.rom_scroll_row = HorizontalScroll(id="rom_scroll_row", classes="rom_row")

        await self.mount(
            Vertical(
                self.logo,
                Static("🎮 Welcome to Agent M", classes="header", expand=False),
                Static(
                    "This is an unofficial DIAMBRA launcher to help you easily train, evaluate, and submit RL agents for fighting games.\n\n"
                    "Verified ROMs are shown below. Click one to view game info and confirm selection.",
                    classes="body",
                    expand=False
                ),
                self.rom_scroll_row,
                Horizontal(
                    self.shared_info_box,
                    self.shared_confirm_button,
                    id="info_row",
                    classes="info_confirm_row",
                ),
                id="home_screen_container",
                classes="centered_layout"
            )
        )


        for rom in verified_roms:
            await self.rom_scroll_row.mount(GameAccordion(
                title=rom["title"],
                rom_file=rom["rom_file"],
                metadata=rom,
                parent_view=self
            ))

    async def on_button_pressed(self, event: Button.Pressed) -> None:
        if event.button.id == "confirm_button" and self.selected_game:
            await self.app.push_screen("training", self.selected_game)