from smolagents import CodeAgent, HfApiModel, tool
import yaml
import requests
import base64
import io
import os
import hashlib
from datetime import datetime
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Create a custom image generation tool that returns a string instead of AgentImage
@tool
def text_to_image_generator(prompt: str) -> str:
    """Generate an image from text and return the result as a string with image details.
    Args:
        prompt: A detailed description of the image to generate
    """
    try:
        print(f"🎨 Starting image generation for: {prompt}")
        
        # Create output directory if it doesn't exist
        output_dir = "generated_images"
        os.makedirs(output_dir, exist_ok=True)
        
        # Method 1: Try Hugging Face Inference API
        api_urls = [
            "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2-1",
            "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5",
            "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
        ]
        
        # Get HF token from environment variable if available
        hf_token = os.getenv('HUGGINGFACE_TOKEN')
        headers = {}
        if hf_token:
            headers['Authorization'] = f'Bearer {hf_token}'
            print("🔐 Using Hugging Face authentication token")
        else:
            print("⚠️ No HF token found, using public access")
        
        for api_url in api_urls:
            try:
                print(f"🔄 Trying API: {api_url}")
                
                # Prepare request payload
                payload = {
                    "inputs": prompt,
                    "parameters": {
                        "num_inference_steps": 30,
                        "guidance_scale": 7.5,
                        "width": 512,
                        "height": 512
                    }
                }
                
                # Make request
                response = requests.post(
                    api_url,
                    json=payload,
                    headers=headers,
                    timeout=60  # Increased timeout for image generation
                )
                
                print(f"📡 API Response: {response.status_code}")
                
                if response.status_code == 200:
                    print("✅ Image generated successfully!")
                    
                    # Generate unique filename
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                    prompt_hash = hashlib.md5(prompt.encode()).hexdigest()[:8]
                    image_filename = f"{output_dir}/generated_image_{timestamp}_{prompt_hash}.png"
                    
                    # Save the image
                    with open(image_filename, 'wb') as f:
                        f.write(response.content)
                    
                    image_size = len(response.content)
                    print(f"💾 Image saved as: {image_filename} (Size: {image_size} bytes)")
                    
                    return f"""✅ IMAGE GENERATED SUCCESSFULLY!

📝 Prompt: {prompt}
🖼️ File: {image_filename}
📦 Size: {image_size} bytes
🌐 API: {api_url.split('/')[-1]}
⏰ Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

The image has been generated and saved successfully. You can view it at: {os.path.abspath(image_filename)}"""
                    
                elif response.status_code == 503:
                    print("⏳ Model is loading, trying next option...")
                    continue
                elif response.status_code == 429:
                    print("⚠️ Rate limit exceeded, trying next option...")
                    continue
                elif response.status_code == 401:
                    print("🔐 Authentication required, trying next option...")
                    continue
                else:
                    print(f"❌ API Error: {response.status_code}")
                    if response.text:
                        print(f"Error details: {response.text[:200]}...")
                    continue
                    
            except requests.exceptions.Timeout:
                print("⏰ Request timeout, trying next option...")
                continue
            except requests.exceptions.ConnectionError:
                print("🔗 Connection error, trying next option...")
                continue
            except requests.exceptions.RequestException as e:
                print(f"🔗 Request error: {e}")
                continue
            except Exception as e:
                print(f"🚫 Unexpected error: {e}")
                continue
        
        # Method 2: If all APIs fail, create a detailed description
        print("🎭 All APIs failed, generating detailed description...")
        
        description = f"""🎨 IMAGE CONCEPT GENERATED

📝 Original Prompt: "{prompt}"

🖼️ Detailed Visual Description:
This image would feature {prompt.lower()}, rendered in high quality with:

🎯 Main Elements:
- Primary subject matter as described in the prompt
- Professional composition and framing
- Balanced lighting and shadows
- Rich, vibrant colors

🎨 Style Characteristics:
- Digital art quality
- Realistic rendering
- High resolution output
- Professional photography aesthetic

📐 Technical Specs:
- Dimensions: 512x512 pixels (standard)
- Format: PNG with transparency support
- Quality: Professional grade
- Style: Photorealistic

💡 Note: This is a detailed concept description. To generate actual images, please:
1. Set up a Hugging Face token: export HUGGINGFACE_TOKEN=your_token_here
2. Ensure stable internet connection
3. Try again with specific, detailed prompts

⚠️ Status: Image generation APIs unavailable - concept description provided instead."""

        print("📋 Description generated successfully!")
        return description
        
    except Exception as e:
        error_msg = f"❌ Image generation failed: {str(e)}"
        print(error_msg)
        return error_msg

def initialize_agent():
    """Initialize the agent with proper error handling"""
    try:
        print("🚀 Initializing Image Generation Agent...")
        
        # Initialize final answer tool
        final_answer = FinalAnswerTool()
        
        # Create model with optimized settings
        model = HfApiModel(
            max_tokens=1024,
            temperature=0.7,
            model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
            custom_role_conversions=None,
        )
        
        # Load prompts with better error handling
        try:
            with open("prompts.yaml", 'r', encoding='utf-8') as stream:
                prompt_templates = yaml.safe_load(stream)
            print("✅ Loaded custom prompts from prompts.yaml")
        except FileNotFoundError:
            print("⚠️ prompts.yaml not found, using default prompts")
            prompt_templates = {
                "system": """You are an AI assistant specialized in image generation. 
When users request images, use the text_to_image_generator tool with detailed, descriptive prompts.
Always provide clear feedback about the image generation process.
Be creative and add artistic details to make prompts more vivid.""",
                "user": "Request: {input}"
            }
        except yaml.YAMLError as e:
            print(f"⚠️ Error parsing prompts.yaml: {e}, using default prompts")
            prompt_templates = {
                "system": """You are an AI assistant specialized in image generation. 
When users request images, use the text_to_image_generator tool with detailed, descriptive prompts.
Always provide clear feedback about the image generation process.""",
                "user": "Request: {input}"
            }
        
        # Create agent with simplified tool set
        agent = CodeAgent(
            model=model,
            tools=[final_answer, text_to_image_generator],
            max_steps=3,
            verbosity_level=1,
            grammar=None,
            planning_interval=None,
            name="StringBasedImageAgent",
            description="AI agent that generates images and returns detailed text descriptions",
            prompt_templates=prompt_templates
        )
        
        print("✅ Agent initialized successfully")
        print("🎯 Available tools: FinalAnswer, TextToImageGenerator")
        return agent
        
    except Exception as e:
        print(f"❌ Agent initialization failed: {e}")
        raise

def launch_agent():
    """Launch the agent with proper error handling"""
    try:
        agent = initialize_agent()
        print("🌐 Starting Gradio interface...")
        
        # Check if GradioUI is available
        try:
            ui = GradioUI(agent)
            ui.launch(share=False, debug=True)
        except Exception as e:
            print(f"❌ Gradio launch failed: {e}")
            print("💡 Possible solutions:")
            print("1. Install Gradio: pip install gradio")
            print("2. Check GradioUI import path")
            print("3. Ensure all dependencies are installed")
            
            # Fallback: simple command line interface
            print("\n🔄 Falling back to command line interface...")
            run_cli_interface(agent)
            
    except Exception as e:
        print(f"❌ Launch failed: {e}")
        print("💡 Please check your dependencies and configuration")

def run_cli_interface(agent):
    """Simple command line interface as fallback"""
    print("\n" + "=" * 50)
    print("🎨 IMAGE GENERATION AGENT - CLI MODE")
    print("=" * 50)
    print("Type 'quit' to exit")
    
    while True:
        try:
            user_input = input("\n🎯 Enter your image prompt: ").strip()
            if user_input.lower() in ['quit', 'exit', 'q']:
                print("👋 Goodbye!")
                break
            
            if user_input:
                print(f"\n🔄 Processing: {user_input}")
                result = agent.run(user_input)
                print(f"\n📋 Result:\n{result}")
            else:
                print("⚠️ Please enter a valid prompt")
                
        except KeyboardInterrupt:
            print("\n👋 Goodbye!")
            break
        except Exception as e:
            print(f"❌ Error: {e}")

if __name__ == "__main__":
    print("=" * 50)
    print("🎨 STRING-BASED IMAGE GENERATION AGENT")
    print("=" * 50)
    print("💡 Setup tips:")
    print("1. Set HF token: export HUGGINGFACE_TOKEN=your_token")
    print("2. Install dependencies: pip install gradio pyyaml requests")
    print("3. Ensure tools/ directory exists with required modules")
    print("=" * 50)
    
    launch_agent()