from textual.screen import Screen
from textual.widgets import Static, Input, Button
from textual.containers import Vertical, Horizontal
from agentm.theme.palette import get_theme
from agentm.utils.logger import log_with_caller
from agentm.logic.db_functions import insert_model, update_run_pending
from datetime import datetime

palette = get_theme()

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

    def __init__(self, agent: dict, model: dict | None, run: dict):
        super().__init__()
        self.agent_metadata = agent
        self.model_metadata = model
        self.run_metadata = run
        self.is_pending_run = run.get("pending", True)

    def compose(self):
        log_with_caller("info", f"Opening TrainingView for agent='{self.agent_metadata['name']}', run='{self.run_metadata['name']}', pending={self.is_pending_run}")

        yield Static(f"[{palette.ACCENT} bold]Training: {self.agent_metadata['name']}[/]", classes="header")

        if self.is_pending_run:
            yield self.render_pending_setup()
        else:
            yield self.render_resume_options()

    def render_pending_setup(self):
        yield Static("[b]Initial Model Setup[/b]", classes="subheader")

        self.name_input = Input(placeholder="Model Name", id="model_name")
        self.steps_input = Input(placeholder="Total Training Steps", id="total_steps")
        self.lr_input = Input(placeholder="Initial Learning Rate", id="learning_rate")
        self.clip_input = Input(placeholder="Initial Clip Range", id="clip_range")
        self.notes_input = Input(placeholder="Notes", id="notes")

        self.confirm_button = Button("✅ Create & Start Training", id="confirm_model_btn", classes="confirm_button")

        return Vertical(
            self.name_input,
            self.steps_input,
            self.lr_input,
            self.clip_input,
            self.notes_input,
            self.confirm_button,
            classes="centered_layout"
        )

    def render_resume_options(self):
        model = self.model_metadata
        from rich.table import Table
        from rich.panel import Panel

        table = Table.grid(padding=(0, 1))
        table.add_column("Key", style="bold underline")
        table.add_column("Value", style=palette.ACCENT, overflow="fold")

        table.add_row("Name", model["name"])
        table.add_row("Steps", f"{model['total_steps_completed']} / {model['total_steps_planned']}")
        table.add_row("Reward", str(model.get("average_reward", "—")))
        table.add_row("Learning Rate", str(model.get("current_learning_rate", "—")))
        table.add_row("Clip Range", str(model.get("current_clip_range", "—")))
        table.add_row("Created", model.get("created_at", "—"))

        return Vertical(
            Static(Panel(table, title="Model Info", border_style=palette.BORDER), classes="agent_info_box"),
            Button("⏯️ Resume Training", id="resume_training_btn", classes="confirm_button"),
            classes="centered_layout"
        )

    async def on_button_pressed(self, event: Button.Pressed) -> None:
        if event.button.id == "confirm_model_btn":
            try:
                name = self.name_input.value.strip()
                total_steps = int(self.steps_input.value)
                learning_rate = float(self.lr_input.value)
                clip_range = float(self.clip_input.value)
                notes = self.notes_input.value.strip()

                log_with_caller("info", f"Creating new model: {name} for agent_id={self.agent_metadata['id']}")

                insert_model(
                    agent_id=self.agent_metadata["id"],
                    name=name,
                    total_steps_planned=total_steps,
                    current_learning_rate=learning_rate,
                    current_clip_range=clip_range,
                    notes=notes
                )

                update_run_pending(self.run_metadata["id"], False)
                log_with_caller("info", f"Model '{name}' created and run marked as not pending")

                await self.app.pop_screen()

            except Exception as e:
                log_with_caller("error", f"Failed to create model: {e}")

        elif event.button.id == "resume_training_btn":
            log_with_caller("info", f"Resuming training for model: {self.model_metadata['name']}")
            await self.app.pop_screen()