#!/usr/bin/env python
# coding=utf-8
"""
Gradio UI for Travel Catalogue Creator
Production-ready interface with streaming agent responses
"""
import os
import re
from typing import Optional, List
from smolagents.agent_types import AgentText, handle_agent_output_types
from smolagents.agents import ActionStep, MultiStepAgent
from smolagents.memory import MemoryStep
from smolagents.utils import _is_package_available


def pull_messages_from_step(step_log: MemoryStep):
    """Extract ChatMessage objects from agent steps with proper nesting"""
    if not _is_package_available("gradio"):
        raise ModuleNotFoundError("Install gradio: `pip install 'smolagents[gradio]'`")
    
    import gradio as gr
    
    if isinstance(step_log, ActionStep):
        # Step header
        step_number = f"Step {step_log.step_number}" if step_log.step_number is not None else "Processing"
        yield gr.ChatMessage(role="assistant", content=f"**{step_number}**")
        
        # Show LLM reasoning/thinking
        if hasattr(step_log, "model_output") and step_log.model_output:
            model_output = step_log.model_output.strip()
            # Clean up code blocks
            model_output = re.sub(r"```\s*<end_code>.*", "```", model_output)
            model_output = re.sub(r"<end_code>\s*```", "```", model_output)
            
            if model_output:
                yield gr.ChatMessage(role="assistant", content=model_output)
        
        parent_id = None
        
        # Handle tool calls
        if hasattr(step_log, "tool_calls") and step_log.tool_calls:
            tool_call = step_log.tool_calls[0]
            parent_id = f"tool_{step_log.step_number}"
            
            # Format tool arguments
            args = tool_call.arguments
            if isinstance(args, dict):
                content = "\n".join(f"• {k}: {v}" for k, v in args.items() if v and k != 'self')
            else:
                content = str(args).strip()
            
            metadata = {
                "title": f"🛠️ Using: {tool_call.name}",
                "id": parent_id,
                "status": "pending",  # Gradio requires "pending", not "running"
            }
            yield gr.ChatMessage(role="assistant", content=content, metadata=metadata)
        
        # Show observations/results
        if hasattr(step_log, "observations") and step_log.observations:
            obs = step_log.observations.strip()
            if obs and not obs.startswith("Execution logs:"):
                metadata = {
                    "title": "✅ Result",
                    "status": "done"
                }
                if parent_id is not None:
                    metadata["parent_id"] = parent_id
                yield gr.ChatMessage(role="assistant", content=obs, metadata=metadata)
        
        # Show errors
        if hasattr(step_log, "error") and step_log.error:
            metadata = {
                "title": "⚠️ Warning",
                "status": "done"
            }
            if parent_id is not None:
                metadata["parent_id"] = parent_id
            yield gr.ChatMessage(role="assistant", content=str(step_log.error), metadata=metadata)
        
        # Step footer with timing and token info
        footer_parts = [step_number]
        if hasattr(step_log, "duration") and step_log.duration:
            footer_parts.append(f"⏱️ {float(step_log.duration):.1f}s")
        if hasattr(step_log, "input_token_count") and hasattr(step_log, "output_token_count"):
            footer_parts.append(f"💬 {step_log.input_token_count + step_log.output_token_count:,} tokens")
        
        yield gr.ChatMessage(
            role="assistant",
            content=f'<span style="color: #888; font-size: 0.85em;">{" | ".join(footer_parts)}</span>',
        )
        
        # Divider between steps
        yield gr.ChatMessage(
            role="assistant",
            content='<hr style="margin: 8px 0; border: 0; border-top: 1px solid #eee">'
        )


def stream_to_gradio(
    agent: MultiStepAgent,
    task: str,
    reset_agent_memory: bool = False,
    additional_args: Optional[dict] = None,
):
    """
    Runs agent and streams messages as gradio ChatMessages.
    
    Args:
        agent: The MultiStepAgent instance to run
        task: User's task/query string
        reset_agent_memory: Whether to clear agent memory before running
        additional_args: Optional additional arguments for the agent
    
    Yields:
        Gradio ChatMessage objects for UI display
    """
    if not _is_package_available("gradio"):
        raise ModuleNotFoundError("Install gradio: `pip install 'smolagents[gradio]'`")
    
    import gradio as gr
    
    try:
        # Run agent and stream steps
        for step_log in agent.run(task, stream=True, reset=reset_agent_memory, additional_args=additional_args):
            if isinstance(step_log, ActionStep):
                # Track token usage if available
                if hasattr(agent.model, "last_input_token_count") and hasattr(agent.model, "last_output_token_count"):
                    step_log.input_token_count = agent.model.last_input_token_count
                    step_log.output_token_count = agent.model.last_output_token_count
                
                # Yield messages from this step
                for message in pull_messages_from_step(step_log):
                    yield message
        
        # Handle final output
        final = handle_agent_output_types(step_log)
        if isinstance(final, AgentText):
            content = final.to_string()
        else:
            # Extract the actual content from the final answer
            content = str(final)
            # Remove the wrapper if it exists (e.g., "FinalAnswerStep(final_answer='...')")
            if "final_answer=" in content or "FinalAnswerStep" in content:
                import re
                match = re.search(r"final_answer=['\"](.+?)['\"](?:\)|$)", content, re.DOTALL)
                if match:
                    content = match.group(1)
                elif "final_answer='" in content:
                    # Alternative extraction if regex fails
                    content = content.split("final_answer='", 1)[1].rsplit("')", 1)[0]
                    # Unescape newlines and other escape sequences
                    content = content.encode().decode('unicode_escape')
        
        yield gr.ChatMessage(
            role="assistant",
            content=content,
            metadata={"status": "done"}
        )
    
    except Exception as e:
        # Handle errors gracefully
        error_msg = str(e)
        
        # Provide helpful error messages
        if "500 Internal Server Error" in error_msg or "Bad Request" in error_msg:
            helpful_msg = (
                "⚠️ **API Error:** The model service encountered an issue.\n\n"
                "**Possible fixes:**\n"
                "• Try rephrasing your request more clearly\n"
                "• Include all required details: destination, dates, origin city, budget\n"
                "• Example: *'5-day Barcelona trip from NYC, Oct 15-19, budget $1500 USD'*\n\n"
                f"Technical details: {error_msg}"
            )
        elif "No results found" in error_msg:
            helpful_msg = (
                "⚠️ **Search Issue:** Couldn't find information about the destination.\n\n"
                "Please try again with a different destination or check the spelling."
            )
        else:
            helpful_msg = (
                f"⚠️ **Error:** {error_msg}\n\n"
                "Please try again with a clearer trip description including:\n"
                "• Destination city\n"
                "• Travel dates\n"
                "• Origin city\n"
                "• Budget amount + currency"
            )
        
        yield gr.ChatMessage(
            role="assistant",
            content=helpful_msg,
            metadata={"status": "done"}
        )


class GradioUI:
    """Production-ready Gradio interface for travel agent"""
    
    def __init__(self, agent: MultiStepAgent, file_upload_folder: Optional[str] = None):
        """
        Initialize Gradio UI wrapper.
        
        Args:
            agent: The MultiStepAgent instance to use
            file_upload_folder: Optional folder path for file uploads
        """
        if not _is_package_available("gradio"):
            raise ModuleNotFoundError("Install gradio: `pip install 'smolagents[gradio]'`")
        
        self.agent = agent
        self.file_upload_folder = file_upload_folder
        
        # Create upload folder if specified
        if self.file_upload_folder and not os.path.exists(file_upload_folder):
            os.makedirs(file_upload_folder, exist_ok=True)
    
    def interact_with_agent(self, prompt: str, history: List):
        """
        Handle user interaction with the agent.
        
        Args:
            prompt: User's input text
            history: Conversation history
        
        Yields:
            Updated conversation history
        """
        import gradio as gr
        
        # Add user message to history
        history.append(gr.ChatMessage(role="user", content=prompt))
        yield history
        
        # Stream agent responses
        for msg in stream_to_gradio(self.agent, task=prompt, reset_agent_memory=False):
            history.append(msg)
            yield history
    
    def launch(self, **kwargs):
        """
        Launch the Gradio interface.
        
        Args:
            **kwargs: Additional arguments passed to demo.launch()
        """
        import gradio as gr
        
        # Build UI
        with gr.Blocks(
            title="✈️ Travel Catalogue Creator",
            fill_height=True,
            theme=gr.themes.Soft()
        ) as demo:
            # Header
            gr.Markdown("# ✈️ Smart Travel Catalogue Creator")
            gr.Markdown(
                "Plan your perfect trip with AI: weather forecasts, custom itineraries, "
                "packing lists & visual inspiration"
            )
            
            # Chat interface
            chatbot = gr.Chatbot(
                label="Your Travel Assistant",
                type="messages",
                avatar_images=(
                    None,  # User avatar (default)
                    "https://huggingface.co/datasets/agents-course/course-images/resolve/main/en/communication/Alfred.png",  # Agent avatar
                ),
                height=600,
                show_copy_button=True,
            )
            
            # Input area
            with gr.Row():
                text_input = gr.Textbox(
                    label="Describe your trip",
                    placeholder="e.g., '4-day Barcelona trip from NYC, Oct 15-19, budget $1500 USD'",
                    lines=2,
                )
                submit_btn = gr.Button("Plan My Trip", variant="primary")
            
            # Help text
            gr.Markdown("""
            ### 💡 Tips for best results:
            • **Include:** destination, dates, origin city, budget amount + currency
            • **Example:** *"5-day Lisbon trip from London, Sep 20-24, budget £800 GBP"*
            • **Example:** *"Weekend Paris getaway from Amsterdam, March 10-12, €600 budget"*
            """)
            
            # Connect interactions
            submit_btn.click(
                self.interact_with_agent,
                inputs=[text_input, chatbot],
                outputs=[chatbot],
                show_progress="full",
            )
            
            text_input.submit(
                self.interact_with_agent,
                inputs=[text_input, chatbot],
                outputs=[chatbot],
                show_progress="full",
            )
        
        # Launch
        demo.launch(**kwargs)


# Export public API
__all__ = ["stream_to_gradio", "GradioUI"]