ui/contentagentui.py

import os
import gradio as gr

from status_check import is_endpoint_healthy      # your tiny checker
from endpoint_utils import wake_endpoint          # your wake+poll helper


class ContentAgentUI:
    """
    Gradio UI that:
      - shows a minimal control panel first (status + Start button),
      - wakes & inits the agent when requested,
      - reveals the main chat panel only after the agent is ready.
    """

    def __init__(self, endpoint_uri: str, is_healthy: bool, health_message: str, agent_initializer):
        self.endpoint_uri = endpoint_uri
        self.is_healthy = bool(is_healthy)
        self.health_message = health_message or ""
        self.agent_initializer = agent_initializer  # callable: (uri) -> CodeAgent

        # set at build()
        self.app: gr.Blocks | None = None
        self.status_box = None
        self.control_panel = None
        self.main_panel = None
        self.prompt = None
        self.reply = None
        self.agent_state = None  # where we store the agent object
        # NOTE: we don't keep a separate self.agent; we read from agent_state in callbacks

    # ---------- helpers ----------

    def _read_css(self) -> str | None:
        css_path = os.path.join(os.getcwd(), "ui", "styles.css")
        if os.path.exists(css_path):
            try:
                with open(css_path, "r", encoding="utf-8") as f:
                    return f.read()
            except Exception:
                return None
        return None

    def _initial_status_text(self) -> str:
        if not self.endpoint_uri:
            return "No endpoint URI configured."
        if self.is_healthy:
            return f"Endpoint healthy ✅ — {self.health_message or 'OK'}.\nClick 'Start Agent' to initialize."
        return (
            f"Endpoint not healthy: {self.health_message or 'unknown'}.\n"
            "Click 'Start Agent' to wake the endpoint and initialize."
        )

    # ---------- agent call ----------

    @staticmethod
    def _call_agent(text: str, agent) -> str:
        try:
            if agent is None:
                return "Agent not initialized yet. Click 'Start Agent'."
            return str(agent.run(text))  # smolagents.CodeAgent API
        except Exception as e:
            return f"Error: {e}"

    # ---------- UI build ----------

    def build(self) -> gr.Blocks:
        if self.app is not None:
            return self.app

        css = self._read_css()
        with gr.Blocks(css=css) as demo:
            gr.Markdown("# Content Agent")

            # Control panel (shown first)
            with gr.Group(visible=True) as self.control_panel:
                self.status_box = gr.Textbox(
                    label="Status",
                    value=self._initial_status_text(),
                    lines=8,
                    interactive=False,
                )
                start_btn = gr.Button("Start Agent (wake if needed)")

            # Main panel (hidden until agent is initialized)
            with gr.Group(visible=False) as self.main_panel:
                self.agent_state = gr.State(None)
                self.prompt = gr.Textbox(label="Your Input", placeholder="Enter something here…")
                self.reply = gr.Textbox(label="Content feedback", interactive=False, lines=12)
                submit_btn = gr.Button("Submit")

                submit_btn.click(self._call_agent, inputs=[self.prompt, self.agent_state], outputs=self.reply)
                self.prompt.submit(self._call_agent, inputs=[self.prompt, self.agent_state], outputs=self.reply)

            # Event: Start Agent (wake → health-check → init → reveal main panel)
            
            def on_start():
                lines = []
                def push(s): 
                    lines.append(s)
                    return ("\n".join(lines), gr.update(), gr.update(), None)
            
                yield push("Waking endpoint… (this can take several minutes for cold starts)")
                ok, err = wake_endpoint(self.endpoint_uri, max_wait=600, poll_every=5.0, log=lambda m: lines.append(m))
                yield push(lines[-1] if lines else "…")  # flush last log line
            
                if not ok:
                    yield push(f"[Server message] {err or 'wake failed'}")
                    return
            
                yield push("Endpoint awake ✅. Checking health…")
                healthy, msg = is_endpoint_healthy(self.endpoint_uri)
                if not healthy:
                    yield push(f"[Server message] {msg}")
                    return
            
                yield push("Initializing agent…")
                try:
                    agent = self.agent_initializer(self.endpoint_uri)
                except Exception as e:
                    yield push(f"Agent init failed: {e}")
                    return
            
                yield ("Agent initialized ✅", gr.update(visible=False), gr.update(visible=True), agent)

            # wire the button: outputs are (status_text, control_panel_update, main_panel_update, agent_state)
            start_btn.click(on_start, inputs=None, outputs=[self.status_box, self.control_panel, self.main_panel, self.agent_state])

            self.app = demo

        return self.app

    # ---------- public API ----------

    def launch(self, **kwargs):
        return self.build().launch(**kwargs)