#!/usr/bin/env python
# coding=utf-8
"""
Gradio_UI.py — BRIANNA's interface.

Replaces the default smolagents template UI with:
  • BRIANNA branding & a dark neural-network aesthetic
  • 🎙️ Voice in / 🔊 voice out (speak to her, hear her reply) — open-source,
    no API keys (Whisper STT + gTTS TTS, see voice.py)
  • File upload (PDF, DOCX, TXT, images, .py files)
  • Image & audio rendering in chat
  • Clear-chat button, example prompts, per-step token/timing footer

The UI degrades gracefully: if voice models or optional Gradio features are
unavailable, the text chat keeps working.
"""

import mimetypes
import os
import re
import shutil
import tempfile
from typing import Optional

from smolagents.agent_types import AgentAudio, AgentImage, AgentText, handle_agent_output_types
from smolagents.agents import ActionStep, MultiStepAgent
from smolagents.memory import MemoryStep
from smolagents.utils import _is_package_available


# ─────────────────────────────────────────────────────────────────────────────
# Open-source voice helpers (inlined) — BRIANNA's ears and mouth.
#   • Listening : OpenAI Whisper via the transformers pipeline (local, no key).
#   • Speaking  : gTTS (free, no key). Both are lazy-loaded and fully graceful,
#     so a missing model/package degrades to text instead of crashing the app.
# ─────────────────────────────────────────────────────────────────────────────
_STT_PIPELINE = None
_STT_MODEL_ID = os.environ.get("BRIANNA_STT_MODEL", "openai/whisper-base")


def _get_stt_pipeline():
    """Lazily build (and cache) the Whisper speech-to-text pipeline."""
    global _STT_PIPELINE
    if _STT_PIPELINE is None:
        from transformers import pipeline  # imported lazily on first use

        _STT_PIPELINE = pipeline(
            task="automatic-speech-recognition",
            model=_STT_MODEL_ID,
            device=-1,  # CPU; set to 0 if a GPU Space is available
        )
    return _STT_PIPELINE


def _voice_transcribe(audio_path):
    """Speech → text: microphone filepath in, recognised text out."""
    if not audio_path:
        return ""
    try:
        pipe = _get_stt_pipeline()
        # return_timestamps=True lets Whisper handle clips longer than 30s.
        result = pipe(audio_path, return_timestamps=True)
        text = result["text"] if isinstance(result, dict) else str(result)
        return text.strip()
    except Exception as e:  # noqa: BLE001 — voice must never crash the app
        return f"[Speech-to-text unavailable: {e}]"


def _clean_for_speech(text, max_chars=1200):
    """Strip markdown / code so the spoken version sounds natural."""
    text = re.sub(r"```.*?```", " (code block omitted) ", text, flags=re.DOTALL)
    text = re.sub(r"[#*`_>|]", " ", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text[:max_chars]


def _voice_synthesize(text):
    """Text → spoken .mp3 path (autoplayed by Gradio), or None on failure."""
    spoken = _clean_for_speech(text or "")
    if not spoken:
        return None
    try:
        from gtts import gTTS  # imported lazily on first use

        tmp = tempfile.NamedTemporaryFile(suffix=".mp3", delete=False)
        tmp.close()
        gTTS(text=spoken, lang="en").save(tmp.name)
        return tmp.name
    except Exception:  # noqa: BLE001 — silently fall back to text-only
        return None


# ─────────────────────────────────────────────────────────────────────────────
# Message extraction from agent steps
# ─────────────────────────────────────────────────────────────────────────────
def pull_messages_from_step(step_log: MemoryStep):
    """Extract gr.ChatMessage objects from agent steps with proper nesting."""
    import gradio as gr

    if isinstance(step_log, ActionStep):
        step_number = f"Step {step_log.step_number}" if step_log.step_number is not None else ""
        yield gr.ChatMessage(role="assistant", content=f"**{step_number}**")

        # Thought / reasoning output
        if getattr(step_log, "model_output", None) is not None:
            model_output = step_log.model_output.strip()
            model_output = re.sub(r"```\s*<end_code>", "```", model_output)
            model_output = re.sub(r"<end_code>\s*```", "```", model_output)
            model_output = re.sub(r"```\s*\n\s*<end_code>", "```", model_output)
            model_output = model_output.strip()
            yield gr.ChatMessage(role="assistant", content=model_output)

        # Tool call block
        if getattr(step_log, "tool_calls", None):
            first_tool_call = step_log.tool_calls[0]
            used_code = first_tool_call.name == "python_interpreter"
            parent_id = f"call_{len(step_log.tool_calls)}"

            args = first_tool_call.arguments
            if isinstance(args, dict):
                content = str(args.get("answer", str(args)))
            else:
                content = str(args).strip()

            if used_code:
                content = re.sub(r"```.*?\n", "", content)
                content = re.sub(r"\s*<end_code>\s*", "", content)
                content = content.strip()
                if not content.startswith("```python"):
                    content = f"```python\n{content}\n```"

            parent_message_tool = gr.ChatMessage(
                role="assistant",
                content=content,
                metadata={
                    "title": f"🛠️ Used tool {first_tool_call.name}",
                    "id": parent_id,
                    "status": "pending",
                },
            )
            yield parent_message_tool

            # Execution logs nested under the tool call
            if getattr(step_log, "observations", None) and step_log.observations.strip():
                log_content = re.sub(r"^Execution logs:\s*", "", step_log.observations.strip())
                yield gr.ChatMessage(
                    role="assistant",
                    content=f"{log_content}",
                    metadata={"title": "📝 Execution Logs", "parent_id": parent_id, "status": "done"},
                )

            # Errors nested under the tool call
            if getattr(step_log, "error", None) is not None:
                yield gr.ChatMessage(
                    role="assistant",
                    content=str(step_log.error),
                    metadata={"title": "💥 Error", "parent_id": parent_id, "status": "done"},
                )

            parent_message_tool.metadata["status"] = "done"

        elif getattr(step_log, "error", None) is not None:
            yield gr.ChatMessage(
                role="assistant",
                content=str(step_log.error),
                metadata={"title": "💥 Error"},
            )

        # Step footer: token counts + duration
        step_footnote = f"{step_number}"
        if hasattr(step_log, "input_token_count") and hasattr(step_log, "output_token_count"):
            try:
                step_footnote += (
                    f" | In: {step_log.input_token_count:,} tokens"
                    f" | Out: {step_log.output_token_count:,} tokens"
                )
            except Exception:  # noqa: BLE001
                pass
        if getattr(step_log, "duration", None):
            step_footnote += f" | ⏱ {round(float(step_log.duration), 2)}s"

        yield gr.ChatMessage(
            role="assistant",
            content=f'<span style="color:#6ee7b7;font-size:11px;font-family:monospace;">{step_footnote}</span>',
        )
        yield gr.ChatMessage(role="assistant", content="---")


# ─────────────────────────────────────────────────────────────────────────────
# Streaming runner
# ─────────────────────────────────────────────────────────────────────────────
def stream_to_gradio(
    agent,
    task: str,
    reset_agent_memory: bool = False,
    additional_args: Optional[dict] = None,
):
    """Run the agent and stream messages as gradio ChatMessages."""
    if not _is_package_available("gradio"):
        raise ModuleNotFoundError(
            "Please install 'gradio': pip install 'smolagents[gradio]'"
        )
    import gradio as gr

    for step_log in agent.run(
        task, stream=True, reset=reset_agent_memory, additional_args=additional_args
    ):
        if hasattr(agent.model, "last_input_token_count") and isinstance(step_log, ActionStep):
            step_log.input_token_count = agent.model.last_input_token_count
            step_log.output_token_count = agent.model.last_output_token_count

        for message in pull_messages_from_step(step_log):
            yield message

    final_answer = handle_agent_output_types(step_log)

    if isinstance(final_answer, AgentText):
        yield gr.ChatMessage(
            role="assistant",
            content=f"**✅ Final Answer:**\n\n{final_answer.to_string()}",
        )
    elif isinstance(final_answer, AgentImage):
        yield gr.ChatMessage(
            role="assistant",
            content={"path": final_answer.to_string(), "mime_type": "image/png"},
        )
    elif isinstance(final_answer, AgentAudio):
        yield gr.ChatMessage(
            role="assistant",
            content={"path": final_answer.to_string(), "mime_type": "audio/wav"},
        )
    else:
        yield gr.ChatMessage(role="assistant", content=f"**✅ Final Answer:** {str(final_answer)}")


# ─────────────────────────────────────────────────────────────────────────────
# GradioUI class — BRIANNA's full interface
# ─────────────────────────────────────────────────────────────────────────────
class GradioUI:
    """BRIANNA's full-capability Gradio interface (chat + voice + files)."""

    ALLOWED_MIME_TYPES = [
        "application/pdf",
        "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        "text/plain",
        "text/x-python",
        "text/markdown",
        "image/png",
        "image/jpeg",
        "image/webp",
        "image/gif",
    ]

    def __init__(self, agent: MultiStepAgent, file_upload_folder: str = "uploads"):
        if not _is_package_available("gradio"):
            raise ModuleNotFoundError(
                "Please install 'gradio': pip install 'smolagents[gradio]'"
            )
        self.agent = agent
        self.file_upload_folder = file_upload_folder
        os.makedirs(self.file_upload_folder, exist_ok=True)

    # ── File handling ──────────────────────────────────────────────────────
    def upload_file(self, file, file_uploads_log):
        import gradio as gr

        if file is None:
            return gr.Textbox("No file uploaded", visible=True), file_uploads_log

        try:
            mime_type, _ = mimetypes.guess_type(file.name)
        except Exception as e:  # noqa: BLE001
            return gr.Textbox(f"Error: {e}", visible=True), file_uploads_log

        if mime_type not in self.ALLOWED_MIME_TYPES:
            return (
                gr.Textbox(
                    f"⚠️ File type '{mime_type}' not allowed. "
                    "Supported: PDF, DOCX, TXT, Python, Markdown, PNG/JPG/WEBP/GIF",
                    visible=True,
                ),
                file_uploads_log,
            )

        original_name = os.path.basename(file.name)
        sanitized_name = re.sub(r"[^\w\-.]", "_", original_name)
        file_path = os.path.join(self.file_upload_folder, sanitized_name)
        shutil.copy(file.name, file_path)

        return (
            gr.Textbox(f"✅ Uploaded: {sanitized_name}", visible=True),
            file_uploads_log + [file_path],
        )

    def log_user_message(self, text_input, file_uploads_log):
        """Attach uploaded file paths to the user message and clear the textbox."""
        full_message = text_input or ""
        if file_uploads_log:
            full_message += (
                f"\n\nYou have been provided with these files (use them as needed): "
                f"{file_uploads_log}"
            )
        return full_message, ""

    # ── Voice in ───────────────────────────────────────────────────────────
    def transcribe_voice(self, audio_path):
        """Microphone → text. Returns the recognised text for the input box."""
        if not audio_path:
            return ""
        return _voice_transcribe(audio_path)

    # ── Chat interaction (streams chat, then speaks the final answer) ───────
    def interact_with_agent(self, prompt, messages):
        import gradio as gr

        if not prompt or not str(prompt).strip():
            yield messages, None
            return

        messages.append(gr.ChatMessage(role="user", content=prompt))
        yield messages, None

        last_text = ""
        for msg in stream_to_gradio(self.agent, task=prompt, reset_agent_memory=False):
            messages.append(msg)
            if isinstance(getattr(msg, "content", None), str) and msg.content.strip() not in ("", "---"):
                last_text = msg.content
            yield messages, None

        # Speak BRIANNA's final answer aloud (autoplayed by the audio component).
        spoken_path = _voice_synthesize(last_text)
        yield messages, spoken_path

    def clear_chat(self):
        """Reset conversation history."""
        return [], [], None

    # ── Launch ─────────────────────────────────────────────────────────────
    def launch(self, **kwargs):
        import gradio as gr

        custom_css = """
        :root {
            --bg-deep:#020810; --bg-panel:#060d1a; --bg-card:#0a1628;
            --accent:#00e5a0; --accent-dim:#00a06e; --accent2:#3b82f6;
            --text:#e2eaf8; --text-muted:#6b7fa3; --border:#1a2d4a; --error:#f87171;
        }
        .gradio-container { background: var(--bg-deep) !important; color: var(--text) !important; }
        .brianna-header { text-align:center; padding:24px 0 6px 0; }
        .brianna-header h1 {
            font-weight:800; font-size:2.6rem; letter-spacing:0.12em; margin:0;
            background:linear-gradient(90deg,#00e5a0 0%,#3b82f6 60%,#a78bfa 100%);
            -webkit-background-clip:text; -webkit-text-fill-color:transparent; background-clip:text;
        }
        .brianna-header p {
            color:var(--text-muted); font-family:monospace; font-size:0.74rem;
            letter-spacing:0.16em; margin:6px 0 0 0;
        }
        .capability-pills { display:flex; flex-wrap:wrap; gap:6px; justify-content:center; margin:10px 0 16px 0; }
        .pill {
            background:var(--bg-card); border:1px solid var(--border); border-radius:20px;
            padding:3px 12px; font-family:monospace; font-size:0.7rem; color:var(--text-muted);
        }
        button.primary {
            background:linear-gradient(135deg,var(--accent) 0%,var(--accent-dim) 100%) !important;
            color:#020810 !important; font-weight:600 !important; border:none !important;
        }
        footer { display:none !important; }
        """

        EXAMPLES = [
            "Give me a PyTorch transformer template and explain each block",
            "Search the web for the latest papers on diffusion models and summarise the top 3",
            "Find the most downloaded text-classification model on HuggingFace",
            "Generate an image of a glowing neural network diagram",
            "What time is it right now in Tokyo and New York?",
            "Create a new tool that converts text to a word-frequency dictionary",
        ]

        with gr.Blocks(css=custom_css, title="BRIANNA — ML Code Agent", fill_height=True) as demo:
            stored_messages = gr.State([])
            file_uploads_log = gr.State([])

            gr.HTML(
                """
                <div class="brianna-header">
                    <h1>BRIANNA</h1>
                    <p>BRILLIANTLY RESPONSIVE INTELLIGENT ASSISTANT FOR NEURAL NETWORK APPLICATIONS</p>
                </div>
                <div class="capability-pills">
                    <span class="pill">🎙️ Voice</span>
                    <span class="pill">🔍 Web Search</span>
                    <span class="pill">🧠 PyTorch / ML</span>
                    <span class="pill">🤗 HuggingFace Hub</span>
                    <span class="pill">🎨 Image Gen</span>
                    <span class="pill">📁 Files</span>
                    <span class="pill">🔧 Self-Extension</span>
                    <span class="pill">🌐 Custom IP / API</span>
                </div>
                """
            )

            chatbot = gr.Chatbot(
                label="",
                type="messages",
                height=520,
                show_copy_button=True,
                avatar_images=(None, None),
            )

            # Spoken reply (autoplays BRIANNA's final answer).
            voice_output = gr.Audio(
                label="🔊 BRIANNA speaks",
                autoplay=True,
                interactive=False,
                visible=True,
            )

            with gr.Row(equal_height=True):
                text_input = gr.Textbox(
                    lines=1,
                    label="Message BRIANNA",
                    placeholder="Type, or use the mic below — ML, code, web research, or 'create a new tool that…'",
                    scale=5,
                )
                send_btn = gr.Button("Send ▶", variant="primary", scale=1, min_width=80)
                clear_btn = gr.Button("🗑 Clear", variant="secondary", scale=1, min_width=80)

            with gr.Accordion("🎙️ Talk to BRIANNA — record, and she replies out loud", open=True):
                voice_input = gr.Audio(
                    sources=["microphone"],
                    type="filepath",
                    label="Hold to record, release to send",
                )

            with gr.Accordion("📎 Attach Files (PDF, DOCX, TXT, Python, Images)", open=False):
                upload_file = gr.File(label="Upload file", file_count="single")
                upload_status = gr.Textbox(
                    label="Upload Status", interactive=False, visible=False, max_lines=1
                )

            with gr.Accordion("💡 Example Prompts — click to use", open=False):
                for ex in EXAMPLES:
                    ex_btn = gr.Button(f"› {ex}", size="sm")
                    ex_btn.click(fn=lambda msg=ex: msg, outputs=text_input)

            # ── Wiring ───────────────────────────────────────────────────
            upload_file.change(
                self.upload_file,
                inputs=[upload_file, file_uploads_log],
                outputs=[upload_status, file_uploads_log],
            )

            # Text submit (button or Enter): log message → run agent → speak.
            for trigger in [send_btn.click, text_input.submit]:
                trigger(
                    self.log_user_message,
                    inputs=[text_input, file_uploads_log],
                    outputs=[stored_messages, text_input],
                ).then(
                    self.interact_with_agent,
                    inputs=[stored_messages, chatbot],
                    outputs=[chatbot, voice_output],
                )

            # Voice submit: transcribe → fill textbox → run agent → speak (hands-free).
            voice_input.stop_recording(
                self.transcribe_voice,
                inputs=[voice_input],
                outputs=[text_input],
            ).then(
                self.log_user_message,
                inputs=[text_input, file_uploads_log],
                outputs=[stored_messages, text_input],
            ).then(
                self.interact_with_agent,
                inputs=[stored_messages, chatbot],
                outputs=[chatbot, voice_output],
            )

            clear_btn.click(self.clear_chat, outputs=[chatbot, stored_messages, voice_output])

        demo.launch(**kwargs)


__all__ = ["stream_to_gradio", "GradioUI"]