import gradio as gr
import replicate
import os
from PIL import Image
import requests
from io import BytesIO
import time
import tempfile
import base64

# Set up Replicate API key from environment variable
os.environ['REPLICATE_API_TOKEN'] = os.getenv('REPLICATE_API_TOKEN')

def upload_image_to_hosting(image):
    """
    Upload image to multiple hosting services with fallback
    """
    # Method 1: Try imgbb.com (more reliable than Imgur)
    try:
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        buffered.seek(0)
        img_base64 = base64.b64encode(buffered.getvalue()).decode()
        
        # imgbb API (free tier)
        response = requests.post(
            "https://api.imgbb.com/1/upload",
            data={
                'key': '6d207e02198a847aa98d0a2a901485a5',  # Free API key
                'image': img_base64,
            }
        )
        
        if response.status_code == 200:
            data = response.json()
            if data.get('success'):
                return data['data']['url']
    except Exception as e:
        print(f"imgbb upload failed: {e}")
    
    # Method 2: Try freeimage.host
    try:
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        buffered.seek(0)
        
        files = {'source': buffered}
        response = requests.post(
            "https://freeimage.host/api/1/upload",
            files=files,
            data={'key': '6d207e02198a847aa98d0a2a901485a5'}
        )
        
        if response.status_code == 200:
            data = response.json()
            if 'image' in data and 'url' in data['image']:
                return data['image']['url']
    except Exception as e:
        print(f"freeimage.host upload failed: {e}")
    
    # Method 3: Try 0x0.st (simple and reliable)
    try:
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        buffered.seek(0)
        
        files = {'file': ('image.png', buffered, 'image/png')}
        response = requests.post(
            "https://0x0.st",
            files=files
        )
        
        if response.status_code == 200:
            return response.text.strip()
    except Exception as e:
        print(f"0x0.st upload failed: {e}")
    
    # Method 4: Original Imgur as last resort
    try:
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        buffered.seek(0)
        img_base64 = base64.b64encode(buffered.getvalue()).decode()
        
        headers = {
            'Authorization': 'Client-ID 0d90e8a3e7d8b4e'
        }
        
        response = requests.post(
            'https://api.imgur.com/3/image',
            headers=headers,
            data={'image': img_base64}
        )
        
        if response.status_code == 200:
            data = response.json()
            if data.get('success'):
                return data['data']['link']
    except Exception as e:
        print(f"Imgur upload failed: {e}")
    
    return None

def save_temp_image(image):
    """
    Save image temporarily and return file path
    """
    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
        image.save(tmp.name, 'PNG')
        return tmp.name

def process_images(prompt, image1, image2=None):
    """
    Process uploaded images with Replicate API
    """
    if not image1:
        return None, "❌ Please upload at least one image"
    
    # Check if API token is set
    if not os.getenv('REPLICATE_API_TOKEN'):
        return None, "⚠️ Please set REPLICATE_API_TOKEN environment variable"
    
    try:
        # Prepare image URLs list
        image_urls = []
        
        # Upload first image
        print("Uploading first image...")
        url1 = upload_image_to_hosting(image1)
        if url1:
            image_urls.append(url1)
            print(f"Image 1 uploaded: {url1}")
        else:
            # If all hosting services fail, use base64 data URI as fallback
            buffered = BytesIO()
            image1.save(buffered, format="PNG")
            buffered.seek(0)
            img_base64 = base64.b64encode(buffered.getvalue()).decode()
            data_uri = f"data:image/png;base64,{img_base64}"
            image_urls.append(data_uri)
            print("Using base64 data URI for image 1")
        
        # Upload second image if provided
        if image2:
            print("Uploading second image...")
            url2 = upload_image_to_hosting(image2)
            if url2:
                image_urls.append(url2)
                print(f"Image 2 uploaded: {url2}")
            else:
                buffered = BytesIO()
                image2.save(buffered, format="PNG")
                buffered.seek(0)
                img_base64 = base64.b64encode(buffered.getvalue()).decode()
                data_uri = f"data:image/png;base64,{img_base64}"
                image_urls.append(data_uri)
                print("Using base64 data URI for image 2")
        
        # Prepare input exactly as shown in the reference
        input_data = {
            "prompt": prompt,
            "image_input": image_urls
        }
        
        print(f"Sending to Replicate: prompt='{prompt}', images={len(image_urls)}")
        
        # Run the model
        output = replicate.run(
            "google/nano-banana",
            input=input_data
        )
        
        # Process output
        if output is None:
            return None, "❌ No output received from model"
        
        # Try different methods to get the image
        try:
            # Method 1: Using read() method
            if hasattr(output, 'read'):
                img_data = output.read()
                img = Image.open(BytesIO(img_data))
                return img, "✅ Free Nano Banana generated your image successfully! 🍌"
        except:
            pass
        
        try:
            # Method 2: Using url() method and downloading
            if hasattr(output, 'url'):
                output_url = output.url()
                response = requests.get(output_url, timeout=30)
                if response.status_code == 200:
                    img = Image.open(BytesIO(response.content))
                    return img, "✅ Free Nano Banana generated your image successfully! 🍌"
        except:
            pass
        
        # Method 3: If output is a URL string or list
        output_url = None
        if isinstance(output, str):
            output_url = output
        elif isinstance(output, list) and len(output) > 0:
            output_url = output[0]
        
        if output_url:
            response = requests.get(output_url, timeout=30)
            if response.status_code == 200:
                img = Image.open(BytesIO(response.content))
                return img, "✅ Free Nano Banana generated your image successfully! 🍌"
        
        return None, f"❌ Could not process output. Output type: {type(output)}"
        
    except replicate.exceptions.ReplicateError as e:
        error_msg = str(e)
        if "502" in error_msg:
            return None, "❌ Server error (502). The nano-banana model might be temporarily unavailable."
        elif "404" in error_msg:
            return None, "❌ Model 'google/nano-banana' not found. Please verify the model exists."
        elif "401" in error_msg or "403" in error_msg:
            return None, "❌ Authentication error. Please check your REPLICATE_API_TOKEN."
        elif "402" in error_msg:
            return None, "❌ Payment required. Please check your Replicate billing status."
        else:
            return None, f"❌ Replicate Error: {error_msg}"
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

# Create Gradio interface with gradient theme
css = """
.gradio-container {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    font-family: 'Inter', sans-serif;
}
.gr-button {
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
    border: none;
    color: white;
    font-weight: bold;
    transition: transform 0.2s;
}
.gr-button:hover {
    transform: scale(1.05);
    box-shadow: 0 10px 20px rgba(0,0,0,0.2);
}
.gr-input {
    border-radius: 10px;
    border: 2px solid rgba(255,255,255,0.3);
    background: rgba(255,255,255,0.9);
}
.header-text {
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    font-size: 2.5em;
    font-weight: bold;
    text-align: center;
    margin-bottom: 20px;
}
.description-text {
    color: white;
    text-align: center;
    font-size: 1.1em;
    margin-bottom: 30px;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.2);
}
"""

# Build the Gradio interface
with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    gr.HTML("""
        <div class="header-text">🍌 Free Nano Banana</div>
        <div class="description-text">
            Upload 1-2 images and describe how you want them styled.
            The AI will create a beautiful transformation using nano-banana model!
        </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Input Section")
            
            prompt = gr.Textbox(
                label="✏️ Style Prompt",
                placeholder="Describe how you want to style your images...",
                lines=3,
                value="Make the sheets in the style of the logo. Make the scene natural."
            )
            
            with gr.Row():
                image1 = gr.Image(
                    label="Image 1 (Required)",
                    type="pil",
                    height=200
                )
                image2 = gr.Image(
                    label="Image 2 (Optional)",
                    type="pil",
                    height=200
                )
            
            generate_btn = gr.Button(
                "🚀 Generate Styled Image",
                variant="primary",
                size="lg"
            )
            
            gr.Markdown("""
            #### 💡 Tips:
            - Upload high-quality images for best results
            - Be specific in your style description
            - Model: google/nano-banana 🍌
            - Free image hosting included!
            """)
        
        with gr.Column(scale=1):
            gr.Markdown("### 🎯 Output Section")
            
            output_image = gr.Image(
                label="Generated Image",
                type="pil",
                height=400
            )
            
            status = gr.Textbox(
                label="Status",
                interactive=False,
                lines=2
            )
    
    # Examples section
    with gr.Row():
        gr.Examples(
            examples=[
                ["Make the sheets in the style of the logo. Make the scene natural.", None, None],
                ["Transform into watercolor painting style", None, None],
                ["Make it look like a vintage photograph", None, None],
                ["Apply cyberpunk neon style", None, None],
            ],
            inputs=[prompt, image1, image2],
            label="Example Prompts"
        )
    
    # Event handlers
    generate_btn.click(
        fn=process_images,
        inputs=[prompt, image1, image2],
        outputs=[output_image, status],
        api_name="generate"
    )
    
    # Additional information
    gr.Markdown("""
    ---
    ### ⚙️ Setup Instructions:
    
    1. **Set Environment Variable:**
       ```bash
       export REPLICATE_API_TOKEN="r8_your_token_here"
       ```
       Get your token from: https://replicate.com/account/api-tokens
    
    2. **Install Required Packages:**
       ```bash
       pip install gradio replicate pillow requests
       ```
    
    3. **Model Information:**
       - Model: `google/nano-banana`
       - Input format: 
         - `prompt`: Text description
         - `image_input`: List of image URLs
       - Output: Generated/styled image
    
    4. **Important Notes:**
       - Images must be hosted with public URLs
       - Currently using Imgur for free image hosting
       - If model returns 502 error, it may be temporarily unavailable
    
    ### 🔒 Security:
    - API keys are managed through environment variables
    - Never commit API keys to version control
    """)

# Launch the app
if __name__ == "__main__":
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )