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

# --- 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 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="🎨 준비 중...")
        
        # Prepare input for Replicate API
        input_data = {
            "prompt": prompt
        }
        
        # If there's an input image, add it to the input
        if image_path:
            # Upload the local image file and get a URL
            # For Replicate, we need to provide URLs, not local paths
            # We'll read the file and create a temporary URL or use file directly
            input_data["image_input"] = [image_path]
        
        progress(0.5, desc="✨ 생성 중...")
        
        # Run the model on Replicate
        output = replicate.run(
            "google/nano-banana",
            input=input_data
        )
        
        progress(0.8, desc="🖼️ 마무리 중...")
        
        # Handle the output
        if output:
            # If output is a FileOutput object, get the URL
            if hasattr(output, 'url'):
                image_url = output.url()
            elif isinstance(output, str):
                image_url = output
            elif isinstance(output, list) and len(output) > 0:
                image_url = output[0]
            else:
                raise ValueError("Unexpected output format from Replicate")
            
            # 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="✅ 완료!")
                return tmpfile.name
        else:
            raise ValueError("No output received from Replicate API")

    except Exception as e:
        print(f"Error details: {e}")
        print(f"Error type: {type(e)}")
        raise gr.Error(f"이미지 생성 실패: {e}")

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.
    """
    if not images:
        raise gr.Error("'여러 이미지' 탭에서 최소 한 개의 이미지를 업로드해주세요.")
    
    try:
        progress(0.2, desc="🎨 이미지 준비 중...")
        
        # Prepare input for Replicate API with multiple images
        input_data = {
            "prompt": prompt,
            "image_input": images  # Pass the list of image paths directly
        }

        progress(0.5, desc="✨ 생성 중...")
        
        # Run the model on Replicate
        output = replicate.run(
            "google/nano-banana",
            input=input_data
        )
        
        progress(0.8, desc="🖼️ 마무리 중...")
        
        # Handle the output
        if output:
            # If output is a FileOutput object, get the URL
            if hasattr(output, 'url'):
                image_url = output.url()
            elif isinstance(output, str):
                image_url = output
            elif isinstance(output, list) and len(output) > 0:
                image_url = output[0]
            else:
                raise ValueError("Unexpected output format from Replicate")
            
            # 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="✅ 완료!")
                return tmpfile.name
        else:
            raise ValueError("No output received from Replicate API")
            
    except Exception as e:
        print(f"Multi-image error details: {e}")
        raise gr.Error(f"이미지 생성 실패: {e}")

# --- 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">
            Google Nano Banana로 구동되는 AI 이미지 생성기
        </p>
    </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;">
            🔐 이 서비스를 사용하려면 Hugging Face 계정으로 로그인해주세요.
        </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;">📸 모드 선택</h3>')
                active_tab_state = gr.State(value="single")
                
                with gr.Tabs(elem_classes="tabs") as tabs:
                    with gr.TabItem("🖼️ 단일 이미지", id="single") as single_tab:
                        image_input = gr.Image(
                            type="filepath",
                            label="입력 이미지 (선택사항)",
                            elem_classes="image-input"
                        )
                        gr.HTML('''
                        <p style="text-align: center; color: #6b7280; font-size: 0.9rem; margin-top: 0.5rem;">
                            💡 텍스트→이미지 생성은 비워두세요
                        </p>
                        ''')
                        
                    with gr.TabItem("🎨 다중 이미지", id="multiple") as multi_tab:
                        gallery_input = gr.Gallery(
                            label="입력 이미지들", 
                            file_types=["image"],
                            elem_classes="gallery-input"
                        )
                        gr.HTML('''
                        <p style="text-align: center; color: #6b7280; font-size: 0.9rem; margin-top: 0.5rem;">
                            💡 여러 이미지를 드래그 앤 드롭하세요
                        </p>
                        ''')
                
                # Prompt Input
                gr.HTML('<h3>✍️ 프롬프트</h3>')
                prompt_input = gr.Textbox(
                    label="",
                    info="AI에게 원하는 이미지를 설명하세요",
                    placeholder="예: 맛있어 보이는 피자, 우주를 배경으로 한 고양이, 미래적인 도시 풍경...",
                    lines=3,
                    elem_classes="prompt-input"
                )
                
                # Generate Button
                generate_button = gr.Button(
                    "🚀 생성하기", 
                    variant="primary",
                    elem_classes="generate-btn"
                )
                
                # Examples
                with gr.Accordion("💡 예제 프롬프트", 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;">🎨 생성 결과</h3>')
                
                output_image = gr.Image(
                    label="", 
                    interactive=False, 
                    elem_id="output"
                )
                
                use_image_button = gr.Button(
                    "♻️ 이 이미지를 다음 편집에 사용", 
                    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;">💡 팁</h4>
                    <ul style="margin: 0; padding-left: 1.5rem; color: #0c4a6e;">
                        <li>구체적이고 상세한 프롬프트를 사용하세요</li>
                        <li>생성된 이미지를 재사용하여 반복적으로 개선할 수 있습니다</li>
                        <li>다중 이미지 모드로 여러 참조 이미지를 결합할 수 있습니다</li>
                        <li>영어 프롬프트가 더 좋은 결과를 생성합니다</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("로그인이 필요합니다. 상단의 'Sign in with Hugging Face' 버튼을 클릭해주세요.")
        if not prompt:
            raise gr.Error("프롬프트를 입력해주세요.")
        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>🔐 로그인이 필요합니다</h2>
                <p style="font-size: 1.1rem; margin: 1rem 0;">
                    이 AI 이미지 생성 도구를 사용하려면 Hugging Face 계정으로 로그인해주세요.
                </p>
                <p style="margin: 1rem 0;">
                    로그인하면 다음 기능을 사용할 수 있습니다:
                </p>
                <ul style="text-align: left; display: inline-block; margin: 1rem 0;">
                    <li>🚀 Google Nano Banana를 통한 고품질 이미지 생성</li>
                    <li>⚡ 빠른 이미지 생성 및 편집</li>
                    <li>🎨 텍스트를 이미지로 변환</li>
                    <li>🔧 다중 이미지 편집 및 결합</li>
                </ul>
                <p style="margin-top: 1.5rem; font-weight: bold;">
                    상단의 "Sign in with Hugging Face" 버튼을 클릭하여 시작하세요!
                </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()