import gradio as gr
import replicate
import os
from typing import Optional, List
from huggingface_hub import whoami
from PIL import Image
import requests
from io import BytesIO
import tempfile
import base64

# --- Replicate API Configuration ---
REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN")

if not REPLICATE_API_TOKEN:
    raise ValueError("REPLICATE_API_TOKEN environment variable is not set.")

# Initialize Replicate client
os.environ["REPLICATE_API_TOKEN"] = REPLICATE_API_TOKEN

def verify_login_status(token: Optional[gr.OAuthToken]) -> bool:
    """Verifies if the user is logged in to Hugging Face."""
    if not token:
        return False
    try:
        user_info = whoami(token=token.token)
        return True if user_info else False
    except Exception as e:
        print(f"Could not verify user's login status: {e}")
        return False

def upload_image_to_hosting(image_path: str) -> str:
    """
    Upload image to hosting service and return URL.
    Using multiple fallback methods for reliability.
    """
    # Open the image
    img = Image.open(image_path)
    
    # Method 1: Try imgbb.com (most reliable)
    try:
        buffered = BytesIO()
        img.save(buffered, format="PNG")
        buffered.seek(0)
        img_base64 = base64.b64encode(buffered.getvalue()).decode()
        
        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 0x0.st (simple and reliable)
    try:
        buffered = BytesIO()
        img.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:
            url = response.text.strip()
            if url.startswith('http'):
                return url
    except Exception as e:
        print(f"0x0.st upload failed: {e}")
    
    # Method 3: Fallback to data URI (last resort)
    buffered = BytesIO()
    img.save(buffered, format="PNG")
    buffered.seek(0)
    img_base64 = base64.b64encode(buffered.getvalue()).decode()
    return f"data:image/png;base64,{img_base64}"

def image_to_data_uri(image_path: str) -> str:
    """Convert local image file to data URI format (kept for backwards compatibility)."""
    with open(image_path, "rb") as img_file:
        img_data = img_file.read()
        img_base64 = base64.b64encode(img_data).decode('utf-8')
        
    # Get the image format
    img = Image.open(image_path)
    img_format = img.format.lower() if img.format else 'png'
    
    # Create data URI
    data_uri = f"data:image/{img_format};base64,{img_base64}"
    return data_uri

def process_output(output, progress=gr.Progress()) -> str:
    """Process the output from Replicate API and return a local file path."""
    try:
        # Check if output has a url attribute (FileObject)
        if hasattr(output, 'url'):
            # If url is a method, call it; if it's a property, just access it
            image_url = output.url() if callable(output.url) else output.url
        # If output is already a string URL
        elif isinstance(output, str):
            image_url = output
        # If output is a list of URLs
        elif isinstance(output, list) and len(output) > 0:
            # Check first item in list
            first_item = output[0]
            if hasattr(first_item, 'url'):
                image_url = first_item.url() if callable(first_item.url) else first_item.url
            else:
                image_url = first_item
        else:
            raise ValueError(f"Unexpected output format from Replicate: {type(output)}")
        
        # Download the image from URL
        response = requests.get(image_url)
        response.raise_for_status()
        
        # Save to temporary file
        img = Image.open(BytesIO(response.content))
        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile:
            img.save(tmpfile.name)
            progress(1.0, desc="✅ Complete!")
            return tmpfile.name
            
    except Exception as e:
        print(f"Error processing output: {e}")
        raise ValueError(f"Failed to process output: {str(e)}")

def run_single_image_logic(prompt: str, image_path: Optional[str] = None, progress=gr.Progress()) -> str:
    """Handles text-to-image or single image-to-image using Replicate's Nano Banana."""
    try:
        progress(0.2, desc="🎨 Preparing...")
        
        # Prepare input for Replicate API
        input_data = {
            "prompt": prompt
        }
        
        # If there's an input image, upload it to get a proper URL
        if image_path:
            progress(0.3, desc="📤 Uploading image...")
            # Upload to hosting service for proper URL
            image_url = upload_image_to_hosting(image_path)
            
            if image_url.startswith('http'):
                print(f"Image uploaded successfully: {image_url[:50]}...")
            else:
                print("Using data URI fallback")
            
            input_data["image_input"] = [image_url]
        
        progress(0.5, desc="✨ Generating...")
        
        # Run the model on Replicate
        # Note: Replace "google/nano-banana" with actual model name if it doesn't exist
        # Examples of real models: "stability-ai/stable-diffusion", "tencentarc/photomaker", etc.
        output = replicate.run(
            "google/nano-banana",  # This might need to be changed to a real model
            input=input_data
        )
        
        progress(0.8, desc="🖼️ Finalizing...")
        
        # Handle the output - output is already a URL string or FileObject
        if output:
            return process_output(output, progress)
        else:
            raise ValueError("No output received from Replicate API")

    except replicate.exceptions.ModelError as e:
        print(f"Replicate Model Error: {e}")
        error_msg = str(e)
        if "does not exist" in error_msg.lower() or "not found" in error_msg.lower():
            raise gr.Error("The specified model 'google/nano-banana' was not found. Please check the model name and ensure your Replicate API token has access.")
        else:
            raise gr.Error(f"Model error: {error_msg[:200]}")
    except Exception as e:
        print(f"Error details: {e}")
        print(f"Error type: {type(e)}")
        if 'output' in locals():
            print(f"Output value: {output}")
            print(f"Output type: {type(output)}")
        raise gr.Error(f"Image generation failed: {str(e)[:200]}")

def run_multi_image_logic(prompt: str, images: List[str], progress=gr.Progress()) -> str:
    """
    Handles multi-image editing by sending a list of images and a prompt.
    Note: Since the actual model might not support multiple images, 
    we'll process only the first image or combine them.
    """
    if not images:
        raise gr.Error("Please upload at least one image in the 'Multiple Images' tab.")
    
    try:
        progress(0.2, desc="🎨 Preparing images...")
        
        # For now, we'll use only the first image since the model might not support multiple
        # You can modify this based on the actual model's capabilities
        image_path = images[0]
        if isinstance(image_path, (list, tuple)):
            image_path = image_path[0]
        
        progress(0.3, desc="📤 Uploading image...")
        image_url = upload_image_to_hosting(image_path)
        
        if image_url.startswith('http'):
            print(f"Image uploaded successfully: {image_url[:50]}...")
        else:
            print("Using data URI fallback")
        
        # Prepare input for Replicate API
        # Using single image format since model might not support multiple
        input_data = {
            "prompt": prompt,
            "image_input": [image_url]  # Send as array with single image
        }

        progress(0.5, desc="✨ Generating...")
        
        # Run the model on Replicate
        # Note: Replace "google/nano-banana" with actual model name
        # Examples of real models: "stability-ai/stable-diffusion", "tencentarc/photomaker", etc.
        output = replicate.run(
            "google/nano-banana",  # This might need to be changed to a real model
            input=input_data
        )
        
        progress(0.8, desc="🖼️ Finalizing...")
        
        # Handle the output using the process_output function
        if output:
            return process_output(output, progress)
        else:
            raise ValueError("No output received from Replicate API")
            
    except replicate.exceptions.ModelError as e:
        print(f"Replicate Model Error: {e}")
        error_msg = str(e)
        if "does not exist" in error_msg.lower() or "not found" in error_msg.lower():
            raise gr.Error("The specified model 'google/nano-banana' was not found. Please check the model name.")
        elif "no image content" in error_msg.lower():
            raise gr.Error("Failed to process images. The model may not support the provided image format or multiple images.")
        else:
            raise gr.Error(f"Model error: {error_msg[:200]}")
    except Exception as e:
        print(f"Multi-image error details: {e}")
        print(f"Input data sent: {input_data if 'input_data' in locals() else 'Not set'}")
        print(f"Output value: {output if 'output' in locals() else 'Not set'}")
        raise gr.Error(f"Image generation failed: {str(e)[:200]}")

# --- Gradio App UI ---
css = '''
/* Header Styling */
.main-header {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    padding: 2rem;
    border-radius: 1rem;
    margin-bottom: 2rem;
    box-shadow: 0 10px 30px rgba(0,0,0,0.1);
}
.header-title {
    font-size: 2.5rem !important;
    font-weight: bold;
    color: white;
    text-align: center;
    margin: 0 !important;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.2);
}
.header-subtitle {
    color: rgba(255,255,255,0.9);
    text-align: center;
    margin-top: 0.5rem !important;
    font-size: 1.1rem;
}
/* Card Styling */
.card {
    background: white;
    border-radius: 1rem;
    padding: 1.5rem;
    box-shadow: 0 4px 6px rgba(0,0,0,0.1);
    border: 1px solid rgba(0,0,0,0.05);
}
.dark .card {
    background: #1f2937;
    border: 1px solid #374151;
}
/* Tab Styling */
.tabs {
    border-radius: 0.5rem;
    overflow: hidden;
    margin-bottom: 1rem;
}
.tabitem {
    padding: 1rem !important;
}
button.selected {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    color: white !important;
}
/* Button Styling */
.generate-btn {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    color: white !important;
    font-size: 1.1rem !important;
    font-weight: 600 !important;
    padding: 0.8rem 2rem !important;
    border-radius: 0.5rem !important;
    cursor: pointer !important;
    transition: all 0.3s ease !important;
    width: 100% !important;
    margin-top: 1rem !important;
}
.generate-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 10px 20px rgba(102, 126, 234, 0.4) !important;
}
.use-btn {
    background: linear-gradient(135deg, #10b981 0%, #059669 100%) !important;
    border: none !important;
    color: white !important;
    font-weight: 600 !important;
    padding: 0.6rem 1.5rem !important;
    border-radius: 0.5rem !important;
    cursor: pointer !important;
    transition: all 0.3s ease !important;
    width: 100% !important;
}
.use-btn:hover {
    transform: translateY(-1px) !important;
    box-shadow: 0 5px 15px rgba(16, 185, 129, 0.4) !important;
}
/* Input Styling */
.prompt-input textarea {
    border-radius: 0.5rem !important;
    border: 2px solid #e5e7eb !important;
    padding: 0.8rem !important;
    font-size: 1rem !important;
    transition: border-color 0.3s ease !important;
}
.prompt-input textarea:focus {
    border-color: #667eea !important;
    outline: none !important;
}
.dark .prompt-input textarea {
    border-color: #374151 !important;
    background: #1f2937 !important;
}
/* Image Output Styling */
#output {
    border-radius: 0.5rem !important;
    overflow: hidden !important;
    box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important;
}
/* Progress Bar Styling */
.progress-bar {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
}
/* Examples Styling */
.examples {
    background: #f9fafb;
    border-radius: 0.5rem;
    padding: 1rem;
    margin-top: 1rem;
}
.dark .examples {
    background: #1f2937;
}
/* Login Message Styling */
.login-message {
    background: linear-gradient(135deg, #fef3c7 0%, #fde68a 100%);
    border-radius: 1rem;
    padding: 2rem;
    text-align: center;
    border: 2px solid #f59e0b;
}
.dark .login-message {
    background: linear-gradient(135deg, #7c2d12 0%, #92400e 100%);
    border-color: #f59e0b;
}
/* Emoji Animations */
@keyframes bounce {
    0%, 100% { transform: translateY(0); }
    50% { transform: translateY(-10px); }
}
.emoji-icon {
    display: inline-block;
    animation: bounce 2s infinite;
}
/* Responsive Design */
@media (max-width: 768px) {
    .header-title {
        font-size: 2rem !important;
    }
    
    .main-container {
        padding: 1rem !important;
    }
}
'''

with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
    # Header
    gr.HTML('''
    <div class="main-header">
        <h1 class="header-title">
            🍌 Real Nano Banana
        </h1>
        <p class="header-subtitle">
            AI Image Generator powered by Google Nano Banana
        </p>
        <div style="display: flex; justify-content: center; align-items: center; gap: 10px; margin-top: 20px;">
            <a href="https://huggingface.co/spaces/ginigen/Nano-Banana-PRO" target="_blank">
                <img src="https://img.shields.io/static/v1?label=NANO%20BANANA&message=PRO&color=%230000ff&labelColor=%23800080&logo=HUGGINGFACE&logoColor=white&style=for-the-badge" alt="badge">
            </a>           
            <a href="https://huggingface.co/spaces/openfree/Nano-Banana-Upscale" target="_blank">
                <img src="https://img.shields.io/static/v1?label=NANO%20BANANA&message=UPSCALE&color=%230000ff&labelColor=%23800080&logo=GOOGLE&logoColor=white&style=for-the-badge" alt="Nano Banana Upscale">
            </a>
            <a href="https://huggingface.co/spaces/openfree/Free-Nano-Banana" target="_blank">
                <img src="https://img.shields.io/static/v1?label=NANO%20BANANA&message=FREE&color=%230000ff&labelColor=%23800080&logo=GOOGLE&logoColor=white&style=for-the-badge" alt="Free Nano Banana">
            </a>
 
            <a href="https://huggingface.co/spaces/ginigen/Nano-Banana-Video" target="_blank">
                <img src="https://img.shields.io/static/v1?label=NANO%20BANANA&message=VIDEO&color=%230000ff&labelColor=%23800080&logo=GOOGLE&logoColor=white&style=for-the-badge" alt="Nano Banana VIDEO">
            </a>                    
            <a href="https://discord.gg/openfreeai" target="_blank">
                <img src="https://img.shields.io/static/v1?label=Discord&message=Openfree%20AI&color=%230000ff&labelColor=%23800080&logo=discord&logoColor=white&style=for-the-badge" alt="Discord Openfree AI">
            </a>
        </div>
    </div>
    ''')  # 여기에 닫는 따옴표 추가
    
    # Login Notice
    gr.HTML('''
    <div style="background: linear-gradient(135deg, #e0f2fe 0%, #bae6fd 100%); 
                border-radius: 0.5rem; padding: 1rem; margin-bottom: 1.5rem; 
                border-left: 4px solid #0284c7;">
        <p style="margin: 0; color: #075985; font-weight: 600;">
            🔐 Please sign in with your Hugging Face account to use this service.
        </p>
    </div>
    ''')
    
    login_message = gr.Markdown(visible=False)
    main_interface = gr.Column(visible=False, elem_classes="main-container")

    with main_interface:
        with gr.Row():
            with gr.Column(scale=1):
                gr.HTML('<div class="card">')
                
                # Mode Selection
                gr.HTML('<h3 style="margin-top: 0;">📸 Select Mode</h3>')
                active_tab_state = gr.State(value="single")
                
                with gr.Tabs(elem_classes="tabs") as tabs:
                    with gr.TabItem("🖼️ Single Image", id="single") as single_tab:
                        image_input = gr.Image(
                            type="filepath",
                            label="Input Image (Optional)",
                            elem_classes="image-input"
                        )
                        gr.HTML('''
                        <p style="text-align: center; color: #6b7280; font-size: 0.9rem; margin-top: 0.5rem;">
                            💡 Leave empty for text-to-image generation
                        </p>
                        ''')
                        
                    with gr.TabItem("🎨 Multiple Images", id="multiple") as multi_tab:
                        gallery_input = gr.Gallery(
                            label="Input Images (Max 2 images)", 
                            file_types=["image"],
                            elem_classes="gallery-input"
                        )
                        gr.HTML('''
                        <p style="text-align: center; color: #6b7280; font-size: 0.9rem; margin-top: 0.5rem;">
                            💡 Upload up to 2 images for combination/editing
                        </p>
                        ''')
                
                # Prompt Input
                gr.HTML('<h3>✍️ Prompt</h3>')
                prompt_input = gr.Textbox(
                    label="",
                    info="Describe what you want the AI to generate",
                    placeholder="e.g., A delicious pizza, a cat in space, futuristic cityscape...",
                    lines=3,
                    elem_classes="prompt-input"
                )
                
                # Generate Button
                generate_button = gr.Button(
                    "🚀 Generate", 
                    variant="primary",
                    elem_classes="generate-btn"
                )
                
                # Examples
                with gr.Accordion("💡 Example Prompts", open=False):
                    gr.Examples(
                        examples=[
                            ["A delicious looking pizza with melting cheese"],
                            ["A cat in a spacesuit walking on the moon surface"],
                            ["Cyberpunk city at night with neon lights"],
                            ["Japanese garden with cherry blossoms in spring"],
                            ["Fantasy wizard tower in a magical world"],
                            ["Make the scene more dramatic and cinematic"],
                            ["Transform this into a watercolor painting style"],
                        ],
                        inputs=prompt_input
                    )
                
                gr.HTML('</div>')

            with gr.Column(scale=1):
                gr.HTML('<div class="card">')
                gr.HTML('<h3 style="margin-top: 0;">🎨 Generated Result</h3>')
                
                output_image = gr.Image(
                    label="", 
                    interactive=False, 
                    elem_id="output"
                )
                
                use_image_button = gr.Button(
                    "♻️ Use this image for next edit", 
                    elem_classes="use-btn",
                    visible=False
                )
                
                # Tips
                gr.HTML('''
                <div style="background: #f0f9ff; border-radius: 0.5rem; padding: 1rem; margin-top: 1rem;">
                    <h4 style="margin-top: 0; color: #0369a1;">💡 Tips</h4>
                    <ul style="margin: 0; padding-left: 1.5rem; color: #0c4a6e;">
                        <li>Use specific and detailed prompts for better results</li>
                        <li>You can reuse generated images for iterative improvements</li>
                        <li>Multiple image mode supports up to 2 images for combination</li>
                        <li>English prompts tend to produce better results</li>
                    </ul>
                </div>
                ''')
                
                gr.HTML('</div>')
        
        # Footer
        gr.HTML('''
        <div style="text-align: center; margin-top: 2rem; padding: 1rem; 
                    border-top: 1px solid #e5e7eb;">
            <p style="color: #6b7280;">
                Made with 💜 using Replicate API | Powered by Google Nano Banana
            </p>
        </div>
        ''')
    
    login_button = gr.LoginButton()
    
    # --- Event Handlers ---
    def unified_generator(
        prompt: str,
        single_image: Optional[str],
        multi_images: Optional[List[str]],
        active_tab: str,
        oauth_token: Optional[gr.OAuthToken] = None,
    ):
        if not verify_login_status(oauth_token):
            raise gr.Error("Login required. Please click the 'Sign in with Hugging Face' button at the top.")
        if not prompt:
            raise gr.Error("Please enter a prompt.")
        if active_tab == "multiple" and multi_images:
            result = run_multi_image_logic(prompt, multi_images)
        else:
            result = run_single_image_logic(prompt, single_image)
        return result, gr.update(visible=True)

    single_tab.select(lambda: "single", None, active_tab_state)
    multi_tab.select(lambda: "multiple", None, active_tab_state)

    generate_button.click(
        unified_generator,
        inputs=[prompt_input, image_input, gallery_input, active_tab_state],
        outputs=[output_image, use_image_button],
    )

    use_image_button.click(
        lambda img: (img, gr.update(visible=False)), 
        inputs=[output_image],
        outputs=[image_input, use_image_button]
    )

    # --- Access Control Logic ---
    def control_access(
        profile: Optional[gr.OAuthProfile] = None,
        oauth_token: Optional[gr.OAuthToken] = None
    ):
        if not profile:
            return gr.update(visible=False), gr.update(visible=False)
        if verify_login_status(oauth_token):
            return gr.update(visible=True), gr.update(visible=False)
        else:
            message = '''
            <div class="login-message">
                <h2>🔐 Login Required</h2>
                <p style="font-size: 1.1rem; margin: 1rem 0;">
                    Please sign in with your Hugging Face account to use this AI image generation tool.
                </p>
                <p style="margin: 1rem 0;">
                    After logging in, you can access:
                </p>
                <ul style="text-align: left; display: inline-block; margin: 1rem 0;">
                    <li>🚀 High-quality image generation via Google Nano Banana</li>
                    <li>⚡ Fast image generation and editing</li>
                    <li>🎨 Text-to-image conversion</li>
                    <li>🔧 Multiple image editing and combining</li>
                </ul>
                <p style="margin-top: 1.5rem; font-weight: bold;">
                    Click the "Sign in with Hugging Face" button at the top to get started!
                </p>
            </div>
            '''
            return gr.update(visible=False), gr.update(visible=True, value=message)

    demo.load(control_access, inputs=None, outputs=[main_interface, login_message])

if __name__ == "__main__":
    demo.queue(max_size=None, default_concurrency_limit=None)
    demo.launch()