from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
import re
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Below is an example of a tool that does nothing. Amaze us with your creativity!
@tool
def create_new_vday_card_tool(theme: str) -> str:
    # it's import to specify the return type
    # Keep this format for the description / args / args description but feel free to modify the tool
    """A tool that generates a custom Valentine's Day card based on a given theme.
    Args:
        theme: A string representing the theme for the Valentine's Day card.
    """
    try:
        # Import tool from Hub
        '''
        image_generation_tool = Tool.from_space(
            space_id="black-forest-labs/FLUX.1-schnell",
            name="generate-image",
            description="Generate image from prompt"
        )
        '''
        image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
        image = image_generation_tool(theme)
        return final_answer(image)
    except Exception as e:
        return f"Error creating card for theme '{theme}': {str(e)}"

# image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
final_answer = FinalAnswerTool()

# Using Mistral-7B-Instruct-v0.1 instead of Qwen2.5-Coder-32B-Instruct
model = HfApiModel(
    max_tokens=2096,
    temperature=0.7,  # Slightly higher temperature for Mistral as it's smaller
    model_id='mistralai/Mistral-7B-Instruct-v0.1',  # Changed to Mistral model
    # custom_role_conversions can be adjusted for Mistral
    custom_role_conversions={
        "system": "<s>[INST] <<SYS>>\n{message}\n<</SYS>>\n\n",
        "human": "{message} [/INST]",
        "assistant": "{message} </s><s>[INST] "
    },
)

# Let's add a fix for the Gradio issue by patching the launch method
def modified_launch():
    try:
        # Original code
        with open("prompts.yaml", 'r') as stream:
            prompt_templates = yaml.safe_load(stream)
        
        agent = CodeAgent(
            model=model,
            tools=[final_answer, create_new_vday_card_tool],  # add your tools here (don't remove final answer)
            max_steps=6,
            verbosity_level=1,
            grammar=None,
            planning_interval=None,
            name=None,
            description=None,
            prompt_templates=prompt_templates
        )
        
        # Get the Gradio UI instance but don't launch yet
        ui = GradioUI(agent)
        
        # Patch the launch method to avoid API issues
        original_launch = ui.launch
        
        def patched_launch(*args, **kwargs):
            # Try to patch the problematic API processing in routes.py
            try:
                import types
                from gradio import routes
                
                # Create a dummy method that returns an empty dict to avoid the error
                def dummy_api_info(self):
                    return {"endpoints": {}, "components": [], "dependencies": []}
                
                # Try to patch FastAPI.api_info
                try:
                    if hasattr(routes, 'FastAPI'):
                        routes.FastAPI.api_info = types.MethodType(dummy_api_info, routes.FastAPI)
                except:
                    pass
                    
            except:
                pass
                
            # Call the original launch method
            return original_launch(*args, **kwargs)
        
        # Replace the launch method with our patched version
        ui.launch = patched_launch
        
        # Now launch the UI with our patched method
        return ui.launch()
    
    except Exception as e:
        # If there's an error with Gradio, fall back to a Flask interface
        print(f"Error with Gradio: {str(e)}")
        print("Falling back to Flask interface...")
        
        # Import Flask and create a simple interface
        from flask import Flask, request, render_template_string, jsonify
        import base64
        from PIL import Image
        from io import BytesIO
        
        app = Flask(__name__)
        
        # Simple HTML template
        html_template = """
        <!DOCTYPE html>
        <html>
        <head>
            <title>Valentine's Day Card Generator</title>
            <style>
                body {
                    font-family: Arial, sans-serif;
                    max-width: 800px;
                    margin: 0 auto;
                    padding: 20px;
                }
                h1 {
                    color: #e91e63;
                }
                .form-container {
                    margin: 20px 0;
                    padding: 15px;
                    border: 1px solid #ffcccc;
                    border-radius: 5px;
                    background-color: #fff0f5;
                }
                label {
                    display: block;
                    margin-bottom: 5px;
                    font-weight: bold;
                }
                input {
                    width: 100%;
                    padding: 8px;
                    margin-bottom: 15px;
                    border: 1px solid #ccc;
                    border-radius: 4px;
                }
                button {
                    background-color: #e91e63;
                    color: white;
                    border: none;
                    padding: 10px 15px;
                    cursor: pointer;
                    border-radius: 4px;
                }
                button:hover {
                    background-color: #c2185b;
                }
                .result-container {
                    margin-top: 20px;
                    text-align: center;
                }
                .card-image {
                    max-width: 100%;
                    border: 10px solid #fff;
                    box-shadow: 0 0 10px rgba(0,0,0,0.2);
                }
            </style>
        </head>
        <body>
            <h1>Valentine's Day Card Generator</h1>
            <p>Enter a theme to generate a custom Valentine's Day card.</p>
            
            <div class="form-container">
                <form method="POST">
                    <label for="theme">Card Theme:</label>
                    <input type="text" id="theme" name="theme" placeholder="e.g., romantic sunset, cute animals, etc." required>
                    <button type="submit">Generate Card</button>
                </form>
            </div>
            
            {% if error %}
            <div style="color: red; margin-top: 20px;">
                <p>{{ error }}</p>
            </div>
            {% endif %}
            
            {% if image_data %}
            <div class="result-container">
                <h2>Your Valentine's Day Card</h2>
                <p>Theme: {{ theme }}</p>
                <img src="data:image/png;base64,{{ image_data }}" alt="Valentine's Day Card" class="card-image">
            </div>
            {% endif %}
        </body>
        </html>
        """
        
        @app.route('/', methods=['GET', 'POST'])
        def home():
            if request.method == 'POST':
                theme = request.form.get('theme', '')
                if theme:
                    try:
                        # Add Valentine's context to the theme
                        valentine_theme = f"Valentine's Day card with theme: {theme}, romantic, love, hearts"
                        
                        # Use the agent to generate the card
                        image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
                        image = image_generation_tool(valentine_theme)
                        
                        # Convert PIL Image to base64 string
                        buffered = BytesIO()
                        image.save(buffered, format="PNG")
                        img_str = base64.b64encode(buffered.getvalue()).decode()
                        
                        return render_template_string(html_template, image_data=img_str, theme=theme)
                    except Exception as e:
                        return render_template_string(html_template, error=f"Error: {str(e)}")
                
            return render_template_string(html_template)
        
        # Run the Flask app
        return app.run(debug=True, host='0.0.0.0', port=8080)

# Run the modified launch function
if __name__ == "__main__":
    modified_launch()