agent_m/agentm/views/home.py

from textual.screen import Screen
from textual.widgets import Static, Button
from textual.containers import Vertical, Horizontal
from textual.message import Message
from textual.reactive import reactive
from textual.widget import Widget
import asyncio
from math import ceil

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(Vertical):
    def __init__(self, title: str, rom_file: str, metadata: dict, parent_view):
        safe_id = rom_file.replace(".", "_").replace("-", "_")
        super().__init__(id=f"accordion_{safe_id}")
        self.title = title
        self.rom_file = rom_file
        self.safe_id = safe_id
        self.metadata = metadata
        self.parent_view = parent_view

    def compose(self):
        yield Button(f"{self.title}", id=f"btn_{self.safe_id}", classes="game_button")

    async def on_button_pressed(self, event: Button.Pressed) -> None:
        if event.button.id == f"btn_{self.safe_id}":
            await self.display_info()

    async def display_info(self):
        log_with_caller("debug", f"Showing shared info for {self.rom_file}")
        meta = self.metadata

        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 compose(self):
        self.logo = Static(
            "[b bright_magenta]\n\n"
            " █████╗   ██████╗  ███████╗ ███╗  ██╗ ████████╗   ███╗   ███╗\n"
            "██╔══██╗ ██╔════╝  ██╔════╝ ████╗ ██║ ╚══██╔══╝   ████╗ ████║\n"
            "███████║ ██║  ███╗ █████╗   ██╔██╗██║    ██║      ██╔████╔██║\n"
            "██╔══██║ ██║   ██║ ██╔══╝   ██║╚████║    ██║      ██║╚██╔╝██║\n"
            "██║  ██║ ╚██████╔╝ ███████╗ ██║ ╚███║    ██║      ██║     ██║  ██║\n"
            "╚═╝  ╚═╝  ╚═════╝  ╚══════╝ ╚═╝  ╚══╝    ╚═╝      ╚═╝     ╚═╝  ╚═╝\n[/]",
            classes="header",
            expand=True,
        )
        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")

        await self.mount(
            Vertical(
                Static("🎮 Welcome to Agent M", classes="header"),
                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"
                ),
                Static(f"✅ Found {len(verified_roms)} valid ROM(s).", id="status_message"),
                Vertical(
                    Horizontal(id="rom_row_1", classes="rom_row"),
                    Horizontal(id="rom_row_2", classes="rom_row"),
                    id="rom_rows",
                    classes="rom_rows_container"
                ),
                self.shared_info_box,
                self.shared_confirm_button,
                id="home_screen_container",
                classes="centered_layout"
            )
        )

        per_row = ceil(len(verified_roms) / 2)
        rows = [verified_roms[:per_row], verified_roms[per_row:]]

        for i, row in enumerate(rows, start=1):
            rom_row = self.query_one(f"#rom_row_{i}", Horizontal)
            for rom in row:
                await rom_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)