ui/contentagentui.py

import os
import gradio as gr

from status_check import is_endpoint_healthy
from endpoint_utils import wake_endpoint


class ContentAgentUI:
    """
    Gradio UI that:
      - shows a minimal control panel first (status + Start button),
      - auto-initializes the agent on load if the endpoint is already healthy,
      - otherwise lets the user 'Start Agent' (wake -> health -> init),
      - reveals the main chat panel (with header, guidance, examples, footer) after init.
    """

    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 in 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
        self.examples_radio = None

    # ---------- 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:
        # neutral; on_load will set real status and maybe auto-init
        return "Checking endpoint status…"

    def _load_examples(self) -> list[str]:
        ex_dir = os.path.join(os.path.dirname(__file__), "examples")
        out: list[str] = []
        if os.path.isdir(ex_dir):
            for name in sorted(os.listdir(ex_dir)):
                if name.lower().endswith(".txt"):
                    p = os.path.join(ex_dir, name)
                    try:
                        with open(p, "r", encoding="utf-8", errors="ignore") as f:
                            out.append(f.read())
                    except Exception:
                        pass
        return out

    # ---------- 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()
        examples = self._load_examples()

        with gr.Blocks(css=css) as demo:
            # global header (always visible)
            gr.Markdown("# Content Agent")

            # Control panel (shown first; may auto-hide on load)
            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:
                # Guidance / about
                gr.Markdown(
                    "### What this does\n"
                    "Enter text below; the agent will classify tone (polite/neutral/impolite) and can use tools.\n\n"
                    "**Tech:**\n"
                    "- `deepseek-ai/DeepSeek-R1-Distill-Qwen-32B` (endpoint)\n"
                    "- Intel Polite Guard tool\n"
                    "- Runs via your Inference Endpoint\n"
                )

                # Chat controls
                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)

                # Examples (optional)
                gr.Markdown("### Try one of these examples")
                if examples:
                    self.examples_radio = gr.Radio(choices=examples, label="Examples")
                    # fill the prompt when an example is picked
                    self.examples_radio.change(lambda ex: ex, inputs=self.examples_radio, outputs=self.prompt)
                else:
                    gr.Markdown("*No examples found.*")

                # Footer
                gr.Markdown("<div id='footer'>Thanks for trying it out!</div>")

            # --- AUTO INIT ON LOAD IF HEALTHY ---
            def on_load():
                healthy, msg = is_endpoint_healthy(self.endpoint_uri)
                if healthy:
                    try:
                        agent = self.agent_initializer(self.endpoint_uri)
                        return (
                            f"Endpoint healthy ✅ — {msg}. Agent initialized.",
                            gr.update(visible=False),   # hide control panel
                            gr.update(visible=True),    # show main panel
                            agent,
                        )
                    except Exception as e:
                        return (
                            f"Agent init failed: {e}",
                            gr.update(visible=True),
                            gr.update(visible=False),
                            None,
                        )
                # not healthy → keep Start button path
                return (
                    f"Endpoint not healthy: {msg}\nClick 'Start Agent' to wake and initialize.",
                    gr.update(visible=True),
                    gr.update(visible=False),
                    None,
                )

            demo.load(
                on_load,
                inputs=None,
                outputs=[self.status_box, self.control_panel, self.main_panel, self.agent_state],
            )

            # --- MANUAL START (wake → health → init) ---
            def on_start():
                lines: list[str] = []

                def push(s: str):
                    lines.append(s)
                    return ("\n".join(lines), gr.update(), gr.update(), None)

                # Wake with progress
                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=lines.append)
                yield ("\n".join(lines), gr.update(), gr.update(), None)  # flush all logs

                if not ok:
                    yield push(f"[Server message] {err or 'wake failed'}")
                    return

                # Health → init
                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)

            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)