import os
import re
import gradio as gr

class GradioUI:
    """
    Gradio UI for a smolagents agent:
      - unwraps FinalAnswerStep to plain text,
      - extracts 'IMAGE:<path>' lines and shows them in a gallery,
      - keeps chat and gallery state across messages (Spaces safe).
    """

    def __init__(self, agent):
        self.agent = agent
        self.images_dir = os.path.abspath("generated_images")
        os.makedirs(self.images_dir, exist_ok=True)

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

    def _pretty_text(self, out):
        # smolagents FinalAnswerTool result object
        if hasattr(out, "final_answer"):
            v = out.final_answer
            return v if isinstance(v, str) else str(v)
        # dict with "final_answer"
        if isinstance(out, dict) and "final_answer" in out:
            v = out["final_answer"]
            return v if isinstance(v, str) else str(v)
        # plain string or anything else
        return out if isinstance(out, str) else str(out)

    IMAGE_LINE_RE = re.compile(
        r"^\s*(?:IMAGE|IMG_PATH)\s*:\s*(.+\.(?:png|jpg|jpeg|webp|bmp|gif))\s*$",
        re.IGNORECASE | re.MULTILINE
    )

    def _extract_image_paths(self, text):
        paths = []
        for m in self.IMAGE_LINE_RE.finditer(text or ""):
            p = m.group(1).strip().strip('"').strip("'")
            p_abs = os.path.abspath(p)
            if os.path.isfile(p_abs):
                paths.append(p_abs)
        return paths

    # --- callbacks ----------------------------------------------------------

    def _chat(self, history_state, message, gallery_state):
        out = self.agent.run(message)
        clean = self._pretty_text(out)
        img_paths = self._extract_image_paths(clean)

        # update conversation history
        history = (history_state or []) + [(message, clean)]
        # update gallery items
        gallery_list = list(gallery_state or [])
        for p in img_paths:
            if p not in gallery_list:
                gallery_list.append(p)

        # Return updated states and component values
        return history, gallery_list, history, gallery_list

    def _clear(self):
        # clear both the UI components and the states
        return [], [], [], []

    # --- layout -------------------------------------------------------------

    def launch(self, **kwargs):
        with gr.Blocks(fill_height=True) as demo:
            gr.Markdown("## 🧠 Tools Agent (text + images) — Hugging Face Space")

            with gr.Row():
                chatbot = gr.Chatbot(label="Chat", height=420, type="messages")
                gallery = gr.Gallery(label="Generated images", height=420, columns=[2], preview=True)

            msg = gr.Textbox(placeholder="Ask: 'What's the time in America/New_York?' or 'Generate an image of a cat astronaut'", lines=2)
            with gr.Row():
                send = gr.Button("Send", variant="primary")
                clear = gr.Button("Clear")

            # persistent states
            chat_state = gr.State([])
            gallery_state = gr.State([])

            send.click(
                self._chat,
                inputs=[chat_state, msg, gallery_state],
                outputs=[chat_state, gallery_state, chatbot, gallery],
            ).then(
                lambda: "",
                inputs=None,
                outputs=msg
            )

            clear.click(
                self._clear,
                inputs=None,
                outputs=[chat_state, gallery_state, chatbot, gallery]
            )

        # queue() is recommended on Spaces
        demo.queue().launch(**kwargs)