from textual.screen import Screen
from textual.widgets import Button, Static
from textual.containers import Vertical
from agentm.utils.logger import log_with_caller
from agentm.components.footer import AgentMFooter
from agentm.theme.palette import get_theme
from agentm.logic.db_functions import get_models_for_run
from agentm.views.model_select import ModelSelectView
from pyfiglet import Figlet

palette = get_theme()

class AgentHomeView(Screen):
    BINDINGS = [("escape", "app.pop_screen", "Back")]

    def __init__(self, agent: dict, run: dict):
        super().__init__()
        self.agent_metadata = agent
        self.run_metadata = run
        self.models_exist = False

    def compose(self):
        f = Figlet(font="ansi_shadow")
        ascii_header = f.renderText(self.agent_metadata["name"])

        yield Static(f"[{palette.ACCENT}]{ascii_header}[/{palette.ACCENT}]", classes="header", expand=False)
        yield Static(
            f"[b]Agent ID:[/] {self.agent_metadata['id']}  |  "
            f"[b]Run:[/] {self.run_metadata['name']}",
            classes="subheader",
            expand=False
        )

        # Buttons - only Eval and Submit may be disabled
        self.train_btn = Button("🚀 Train", id="train_btn", classes="confirm_button")
        self.eval_btn = Button("🎮 Eval", id="eval_btn", classes="confirm_button", disabled=True)
        self.submit_btn = Button("📦 Submit", id="submit_btn", classes="confirm_button", disabled=True)

        yield Vertical(
            self.train_btn,
            self.eval_btn,
            self.submit_btn,
            id="agent_home_menu",
            classes="centered_layout"
        )

        yield AgentMFooter(compact=True)

    async def on_mount(self):
        log_with_caller("info", f"Entered AgentHomeView for agent '{self.agent_metadata['name']}'")
        models = get_models_for_run(self.run_metadata["id"])
        self.models_exist = len(models) > 0
        self.eval_btn.disabled = not self.models_exist
        self.submit_btn.disabled = not self.models_exist
        log_with_caller("debug", f"Found {len(models)} model(s) for run '{self.run_metadata['name']}'")

    async def on_button_pressed(self, event: Button.Pressed) -> None:
        action = event.button.id.replace("_btn", "")  # 'train', 'eval', or 'submit'
        log_with_caller("info", f"{action.title()} selected for run '{self.run_metadata['name']}'")

        # Always go to model select before action
        await self.app.push_screen(
            ModelSelectView(
                agent_metadata=self.agent_metadata,
                mode=action,
                run=self.run_metadata
            )
        )