import gradio as gr
from google import genai 
from google.genai import types 
import os
from typing import Optional, List
from huggingface_hub import whoami
from PIL import Image
from io import BytesIO
import tempfile

# --- Google Gemini API Configuration ---
# Use GEMINI_API_KEY if available, otherwise fall back to GOOGLE_API_KEY
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

API_KEY = GEMINI_API_KEY or GOOGLE_API_KEY

if not API_KEY:
    raise ValueError("Neither GEMINI_API_KEY nor GOOGLE_API_KEY environment variable is set.")

client = genai.Client(
    api_key=API_KEY,
)

# Note: Gemini models don't directly generate images - they analyze/describe them
# For image generation, you'd need to use a different API like Imagen
# This code is updated to work with text generation about images
GEMINI_MODEL_NAME = 'gemini-2.0-flash-exp'  # Updated model name

def verify_pro_status(token: Optional[gr.OAuthToken]) -> bool:
    """Verifies if the user is a Hugging Face PRO user or part of an enterprise org."""
    if not token:
        return False
    try:
        user_info = whoami(token=token.token)
        if user_info.get("isPro", False):
            return True
        orgs = user_info.get("orgs", [])
        if any(org.get("isEnterprise", False) for org in orgs):
            return True
        return False
    except Exception as e:
        print(f"Could not verify user's PRO/Enterprise status: {e}")
        return False

def _extract_image_data_from_response(response) -> Optional[bytes]:
    """Helper to extract image data from the model's response."""
    # Debug: Print response structure
    print(f"Response type: {type(response)}")
    
    # Note: Gemini doesn't generate images directly
    # This would need to be replaced with actual image generation API
    # For now, return None to indicate no image was generated
    
    if hasattr(response, 'text'):
        print(f"Response text: {response.text[:200] if response.text else 'Empty'}")
    
    return None

def run_single_image_logic(prompt: str, image_path: Optional[str] = None, progress=gr.Progress()) -> str:
    """Handles text or image analysis using Google Gemini."""
    try:
        progress(0.2, desc="🎨 준비 중...")
        
        contents = []
        if image_path:
            # Image analysis
            input_image = Image.open(image_path)
            contents.append(input_image)
            contents.append(f"Describe what you see and how it could be modified based on: {prompt}")
        else:
            # Text-only prompt
            contents.append(f"Describe an image concept for: {prompt}")

        progress(0.5, desc="✨ 생성 중...")
        
        # Remove the generation_config parameter that's causing the error
        # Use the simpler API call format
        try:
            response = client.models.generate_content( 
                model=GEMINI_MODEL_NAME,
                contents=contents,
            )
        except Exception as api_error:
            # Fallback: try with just the contents as a simple string/list
            print(f"First attempt failed: {api_error}")
            if image_path:
                # For image input, we need to handle it differently
                # The API might expect a different format
                raise gr.Error("이미지 입력은 현재 지원되지 않습니다. Gemini API는 이미지 생성이 아닌 분석용입니다.")
            else:
                # For text-only, try a simpler approach
                response = client.models.generate_content(
                    model=GEMINI_MODEL_NAME,
                    contents=prompt
                )
        
        progress(0.8, desc="🖼️ 마무리 중...")
        
        # Since Gemini doesn't generate images, we'll need to handle this differently
        # For demonstration, create a placeholder or use a different service
        if hasattr(response, 'text') and response.text:
            # Return the text response for now
            # In production, you'd call an actual image generation API here
            description = response.text
            
            # Create a placeholder image with the description
            from PIL import Image, ImageDraw, ImageFont
            img = Image.new('RGB', (512, 512), color='white')
            draw = ImageDraw.Draw(img)
            
            # Add text to explain the situation
            text = "Gemini API는 이미지 생성을 지원하지 않습니다.\n\n"
            text += "생성된 설명:\n" + description[:200] + "..."
            
            # Simple text wrapping
            y_position = 50
            for line in text.split('\n'):
                draw.text((20, y_position), line, fill='black')
                y_position += 30
            
            # Save the placeholder image
            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile:
                img.save(tmpfile.name)
                progress(1.0, desc="✅ 완료!")
                return tmpfile.name
        else:
            raise ValueError("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 analysis.
    """
    if not images:
        raise gr.Error("'여러 이미지' 탭에서 최소 한 개의 이미지를 업로드해주세요.")
    
    try:
        progress(0.2, desc="🎨 이미지 준비 중...")
        contents = []
        for image_path in images:
            if isinstance(image_path, (list, tuple)):
                image_path = image_path[0]
            contents.append(Image.open(image_path))
        contents.append(f"Analyze these images based on: {prompt}")

        progress(0.5, desc="✨ 분석 중...")
        
        # Simple API call without generation_config
        response = client.models.generate_content( 
            model=GEMINI_MODEL_NAME,
            contents=contents,
        )
        
        progress(0.8, desc="🖼️ 마무리 중...")
        
        # Create a result image with the analysis
        if hasattr(response, 'text') and response.text:
            from PIL import Image, ImageDraw
            img = Image.new('RGB', (512, 512), color='white')
            draw = ImageDraw.Draw(img)
            
            text = "다중 이미지 분석 결과:\n\n"
            text += response.text[:300] + "..."
            
            y_position = 50
            for line in text.split('\n'):
                draw.text((20, y_position), line, fill='black')
                y_position += 30
            
            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile:
                img.save(tmpfile.name)
                progress(1.0, desc="✅ 완료!")
                return tmpfile.name
        else:
            raise ValueError("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;
}

/* Pro Message Styling */
.pro-message {
    background: linear-gradient(135deg, #fef3c7 0%, #fde68a 100%);
    border-radius: 1rem;
    padding: 2rem;
    text-align: center;
    border: 2px solid #f59e0b;
}

.dark .pro-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 Gemini API로 구동되는 AI 이미지 분석기
        </p>
    </div>
    ''')
    
    # Important Notice
    gr.HTML('''
    <div style="background: #fef2f2; border-radius: 0.5rem; padding: 1rem; margin-bottom: 1.5rem; 
                border-left: 4px solid #ef4444;">
        <p style="margin: 0; color: #991b1b; font-weight: 600;">
            ⚠️ 참고: Gemini API는 이미지 생성이 아닌 분석을 제공합니다. 
            실제 이미지 생성을 원하시면 DALL-E, Midjourney, Stable Diffusion 등을 사용하세요.
        </p>
    </div>
    ''')
    
    # Pro User Notice
    gr.HTML('''
    <div style="background: linear-gradient(135deg, #fef3c7 0%, #fde68a 100%); 
                border-radius: 0.5rem; padding: 1rem; margin-bottom: 1.5rem; 
                border-left: 4px solid #f59e0b;">
        <p style="margin: 0; color: #92400e; font-weight: 600;">
            🌟 이 스페이스는 Hugging Face PRO 사용자 전용입니다. 
            <a href="https://huggingface.co/pro" target="_blank" 
               style="color: #dc2626; text-decoration: underline;">
               PRO 구독하기
            </a>
        </p>
    </div>
    ''')

    pro_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=[
                            ["이 이미지에서 무엇이 보이나요?"],
                            ["이미지의 색상과 구성을 분석해주세요"],
                            ["이 장면을 더 극적으로 만들려면 어떻게 해야 할까요?"],
                            ["이미지의 분위기와 감정을 설명해주세요"],
                            ["기술적인 관점에서 이 이미지를 평가해주세요"],
                        ],
                        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>이미지를 업로드하여 AI 분석을 받을 수 있습니다</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 💜 by Hugging Face PRO | Powered by Google Gemini API
            </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_pro_status(oauth_token):
            raise gr.Error("액세스 거부: 이 서비스는 PRO 사용자 전용입니다.")
        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_pro_status(oauth_token):
            return gr.update(visible=True), gr.update(visible=False)
        else:
            message = '''
            <div class="pro-message">
                <h2>✨ PRO 사용자 전용 기능</h2>
                <p style="font-size: 1.1rem; margin: 1rem 0;">
                    이 강력한 AI 이미지 분석 도구는 Hugging Face <strong>PRO</strong> 멤버 전용입니다.
                </p>
                <p style="margin: 1rem 0;">
                    PRO 구독으로 다음을 누리세요:
                </p>
                <ul style="text-align: left; display: inline-block; margin: 1rem 0;">
                    <li>🚀 Google Gemini API 무제한 액세스</li>
                    <li>⚡ 빠른 이미지 분석</li>
                    <li>🎨 상세한 이미지 설명</li>
                    <li>🔧 다중 이미지 비교 분석</li>
                </ul>
                <a href="https://huggingface.co/pro" target="_blank" 
                   style="display: inline-block; margin-top: 1rem; padding: 1rem 2rem;
                          background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                          color: white; text-decoration: none; border-radius: 0.5rem;
                          font-weight: bold; font-size: 1.1rem;">
                    🌟 지금 PRO 멤버 되기
                </a>
            </div>
            '''
            return gr.update(visible=False), gr.update(visible=True, value=message)

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

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