import os
import json
import base64
import logging
from typing import Optional
from datetime import datetime
from pathlib import Path
from io import BytesIO

import gradio as gr
from PIL import Image, ImageDraw

# Google Gemini API (New SDK - Nano Banana)
from google import genai
from google.genai import types

from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Create directory for generated images
GENERATED_DIR = Path("generated_images")
GENERATED_DIR.mkdir(exist_ok=True)

# Initialize Dataset Manager
HF_TOKEN = os.getenv("HF_TOKEN")
DATASET_REPO_ID = os.getenv("DATASET_REPO_ID")
dataset_manager = None

if HF_TOKEN and DATASET_REPO_ID:
    try:
        from dataset_manager import DatasetManager
        dataset_manager = DatasetManager(DATASET_REPO_ID, HF_TOKEN)
        logger.info(f"Dataset manager initialized for repository: {DATASET_REPO_ID}")
    except Exception as e:
        logger.warning(f"Could not initialize dataset manager: {e}")
else:
    if not HF_TOKEN:
        logger.info("HF_TOKEN not set. Dataset saving feature disabled.")
    if not DATASET_REPO_ID:
        logger.info("DATASET_REPO_ID not set. Dataset saving feature disabled.")

# Available Gemini models
AVAILABLE_MODELS = {
    "gemini-2.5-flash-image": {
        "name": "Gemini 2.5 Flash Image",
        "description": "Fast, low-cost ($0.039/image), 10 aspect ratios",
        "cost": "Low"
    },
    "gemini-3-pro-image-preview": {
        "name": "Gemini 3 Pro Image Preview",
        "description": "High-quality, 2K/4K resolution, Google Search grounding",
        "cost": "High"
    }
}

def generate_image_with_gemini(prompt: str, gemini_api_key: str, model: str = "gemini-2.5-flash-image",
                               aspect_ratio: str = "1:1", size: str = "1K") -> Image.Image:
    """
    Generate image using Gemini with user-provided API key and model (New SDK)

    Args:
        prompt: 画像生成プロンプト
        gemini_api_key: Gemini APIキー
        model: モデル名
        aspect_ratio: アスペクト比 (1:1, 4:3, 3:4, 16:9, 9:16, 3:2)
        size: 画像サイズ (1K, 2K, 4K) - Gemini 3 Proのみ有効
    """
    if not gemini_api_key or not gemini_api_key.strip():
        logger.warning("No API key provided, using placeholder image generation")
        return generate_placeholder_image(prompt, 1024, 1024)

    try:
        # 新SDK: Clientベースのアーキテクチャ
        client = genai.Client(api_key=gemini_api_key.strip())

        # Validate model name
        if model not in AVAILABLE_MODELS:
            logger.warning(f"Invalid model '{model}', using default")
            model = "gemini-2.5-flash-image"

        # プロンプトをそのまま使用（Style機能削除）
        enhanced_prompt = prompt

        # Add camera and technical details for better results
        if "portrait" in prompt.lower():
            enhanced_prompt += ". Shot with 85mm lens, shallow depth of field, golden hour lighting"
        elif "landscape" in prompt.lower():
            enhanced_prompt += ". Wide-angle shot, dramatic lighting, high dynamic range"
        elif "product" in prompt.lower():
            enhanced_prompt += ". Professional product photography, clean white background, studio lighting"

        logger.info(f"Generating image with {model}: {enhanced_prompt[:100]}...")

        # ✅ モデル別のImageConfig設定
        if model == "gemini-3-pro-image-preview":
            # Gemini 3 Pro: image_sizeパラメータをサポート
            logger.info(f"Using Gemini 3 Pro with image_size={size}, aspect_ratio={aspect_ratio}")

            config = types.GenerateContentConfig(
                temperature=1.0,
                response_modalities=[types.Modality.TEXT, types.Modality.IMAGE],
                image_config=types.ImageConfig(
                    aspect_ratio=aspect_ratio,
                    image_size=size,  # ✅ Gemini 3 Proでのみ指定
                )
            )
        else:
            # Gemini 2.5 Flash: aspect_ratioのみサポート、image_sizeは指定しない
            logger.info(f"Using Gemini 2.5 Flash with aspect_ratio={aspect_ratio} (1024px固定)")

            config = types.GenerateContentConfig(
                temperature=1.0,
                response_modalities=[types.Modality.TEXT, types.Modality.IMAGE],
                image_config=types.ImageConfig(
                    aspect_ratio=aspect_ratio,
                    # image_sizeは指定しない（デフォルト1024px）
                )
            )

        # 新SDK: generate_contentの呼び出し
        response = client.models.generate_content(
            model=model,
            contents=enhanced_prompt,
            config=config
        )

        # Process response
        if response.candidates:
            for candidate in response.candidates:
                for part in candidate.content.parts:
                    if hasattr(part, 'inline_data') and part.inline_data:
                        # Image data is returned as inline_data
                        image_data = part.inline_data.data
                        mime_type = part.inline_data.mime_type

                        if mime_type and mime_type.startswith('image/'):
                            image = Image.open(BytesIO(image_data))
                            return image
                    elif hasattr(part, 'text') and part.text:
                        logger.info(f"Gemini text response: {part.text[:200]}")

        # Fallback to placeholder if no image generated
        return generate_placeholder_image(enhanced_prompt, 1024, 1024)

    except Exception as e:
        logger.error(f"Error generating image with Gemini: {e}")
        return generate_placeholder_image(prompt, 1024, 1024)

def generate_placeholder_image(prompt: str, width: int = 1024, height: int = 1024) -> Image.Image:
    """Generate a beautiful placeholder image with gradient and text"""
    # Create gradient background
    img = Image.new('RGB', (width, height))
    pixels = img.load()

    # Create a more vibrant gradient
    for y in range(height):
        for x in range(width):
            # Diagonal gradient with vibrant colors
            r = int((x / width) * 180 + 75)
            g = int((y / height) * 120 + 60)
            b = int(((x + y) / (width + height)) * 200 + 55)
            pixels[x, y] = (r, g, b)

    # Add text overlay
    draw = ImageDraw.Draw(img)

    # Create semi-transparent overlay for text background
    overlay = Image.new('RGBA', (width, height), (0, 0, 0, 0))
    overlay_draw = ImageDraw.Draw(overlay)

    # Draw a semi-transparent rectangle for text background
    rect_height = height // 3
    rect_y = (height - rect_height) // 2
    overlay_draw.rectangle(
        [(0, rect_y), (width, rect_y + rect_height)],
        fill=(0, 0, 0, 120)
    )

    # Composite overlay onto main image
    img = Image.alpha_composite(img.convert('RGBA'), overlay).convert('RGB')
    draw = ImageDraw.Draw(img)

    # Draw text
    text_lines = [
        "🍌 NanoBanana Generator",
        "",
        "Generated prompt:",
        f'"{prompt[:60]}..."' if len(prompt) > 60 else f'"{prompt}"',
        "",
        f"Size: {width}x{height}",
        "",
        "⚠️ Add Gemini API Key for real AI generation"
    ]

    try:
        # Calculate text position
        line_height = height // 20
        start_y = (height - len(text_lines) * line_height) // 2

        for i, line in enumerate(text_lines):
            text_bbox = draw.textbbox((0, 0), line)
            text_width = text_bbox[2] - text_bbox[0]
            position = ((width - text_width) // 2, start_y + i * line_height)
            draw.text(position, line, fill=(255, 255, 255))
    except:
        pass

    return img

# Gradio Interface functions
def gradio_generate(gemini_api_key: str, model: str, aspect_ratio: str, size: str, prompt: str,
                   save_to_dataset: bool = True, dataset_folder: str = "", custom_filename: str = ""):
    """Generate image through Gradio interface using Nano Banana"""
    try:
        if not prompt:
            return None, "❌ プロンプトを入力してください"

        if not gemini_api_key or not gemini_api_key.strip():
            return None, "❌ Gemini APIキーを入力してください"

        # Validate model
        if model not in AVAILABLE_MODELS:
            return None, f"❌ 無効なモデルが選択されました"

        # Generate image using Gemini
        # Gemini 2.5 Flashの場合、sizeは無視される（1024px固定）
        image = generate_image_with_gemini(prompt, gemini_api_key, model, aspect_ratio, size)

        # Save image locally with custom or auto-generated filename
        if custom_filename and custom_filename.strip():
            # Sanitize custom filename
            clean_name = os.path.splitext(custom_filename.strip())[0]
            clean_name = "".join(c if c.isalnum() or c in '-_' else '_' for c in clean_name)
            if not clean_name:
                clean_name = f"gradio_gen_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            filename = f"{clean_name}.png"
        else:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"gradio_gen_{timestamp}.png"

        filepath = GENERATED_DIR / filename
        image.save(filepath)

        model_info = AVAILABLE_MODELS.get(model, {})
        status = f"✅ 生成成功！ファイル名: {filename}"
        status += f"\n🎨 モデル: {model_info.get('name', model)}"
        status += f"\n📐 アスペクト比: {aspect_ratio}"
        if model == "gemini-3-pro-image-preview":
            status += f"\n📏 サイズ: {size}"

        # Save to dataset if enabled
        if dataset_manager and save_to_dataset:
            try:
                metadata = {
                    "aspect_ratio": aspect_ratio,
                    "size": size if model == "gemini-3-pro-image-preview" else "1K",
                    "model": model,
                    "generation_type": "text-to-image"
                }

                # Use provided folder or None (will default to date)
                folder_name = dataset_folder if dataset_folder.strip() else None

                # Use custom filename for dataset as well
                dataset_filename = custom_filename.strip() if custom_filename.strip() else None

                dataset_url = dataset_manager.save_image(
                    image=image,
                    prompt=prompt,
                    folder_name=folder_name,
                    filename=dataset_filename,
                    metadata=metadata
                )

                if dataset_url:
                    status += f"\n📁 Dataset保存: {folder_name or datetime.now().strftime('%Y_%m_%d')}"
                    status += f"\n🔗 URL: {dataset_url}"
            except Exception as dataset_error:
                status += f"\n⚠️ Dataset保存失敗: {str(dataset_error)}"

        return image, status

    except Exception as e:
        return None, f"❌ エラー: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="NanoBanana Gemini Image Generator V9", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🍌 NanoBanana - Gemini画像生成 (Version 9)

        Google Gemini AIでテキストから画像を生成します。

        [Google AI Studio](https://aistudio.google.com/app/apikey)で無料APIキーを取得してください。
        """
    )

    # Gemini API Key入力
    gemini_api_key_input = gr.Textbox(
        label="Gemini API Key",
        placeholder="AIza... で始まるAPIキーを入力",
        type="password",
        value="",
        interactive=True
    )

    # Model選択（Radioボタン）
    model_radio = gr.Radio(
        label="Model",
        choices=[
            ("Gemini 2.5 Flash（高速・1024px固定）", "gemini-2.5-flash-image"),
            ("Gemini 3 Pro（高品質・4K対応）", "gemini-3-pro-image-preview")
        ],
        value="gemini-2.5-flash-image",
        interactive=True
    )

    # Generation Tab
    gr.Markdown("### 画像生成")

    with gr.Row():
        with gr.Column():
            gen_prompt = gr.Textbox(
                label="Prompt",
                placeholder="例: 夕焼けの富士山、フォトリアリスティック、4K画質",
                lines=4
            )

            with gr.Row():
                gen_aspect_ratio = gr.Dropdown(
                    label="Aspect Ratio",
                    choices=["1:1", "4:3", "3:4", "16:9", "9:16", "3:2"],
                    value="1:1",
                    interactive=True
                )
                gen_size = gr.Dropdown(
                    label="Size（Gemini 3 Pro用）",
                    choices=["1K", "2K", "4K"],
                    value="1K",
                    visible=False,  # 初期は非表示（Gemini 2.5 Flash選択時）
                    interactive=True
                )

            # Dataset save options
            with gr.Accordion("📁 Dataset Options", open=False):
                gen_save_dataset = gr.Checkbox(
                    label="Save to Dataset Repository",
                    value=True if dataset_manager else False,
                    interactive=bool(dataset_manager)
                )
                gen_dataset_folder = gr.Textbox(
                    label="Folder Name (optional)",
                    placeholder="例: portraits（空欄の場合は日付フォルダ）",
                    value="",
                    interactive=bool(dataset_manager)
                )
                gen_custom_filename = gr.Textbox(
                    label="Custom Filename (optional)",
                    placeholder="例: my-artwork（拡張子不要）",
                    value="",
                    interactive=True
                )

                if not dataset_manager:
                    gr.Markdown("⚠️ Dataset保存は無効です。HF_TOKENとDATASET_REPO_IDを環境変数に設定してください。")

            gen_button = gr.Button("🚀 Generate Image", variant="primary", size="lg")

        with gr.Column():
            gen_output = gr.Image(label="Generated Image", type="pil")
            gen_status = gr.Textbox(label="Status", interactive=False)

    # Professional examples
    gr.Examples(
        examples=[
            ["富士山と桜、フォトリアリスティック、夕焼け、4K画質", "1:1"],
            ["可愛い猫のイラスト、アニメスタイル、パステルカラー", "1:1"],
            ["夕日に向かって走る犬、シネマティック", "16:9"],
        ],
        inputs=[gen_prompt, gen_aspect_ratio],
        label="Example Prompts"
    )

    # Size表示の動的制御（Gemini 3 Pro選択時のみ表示）
    def update_size_visibility(model):
        if model == "gemini-3-pro-image-preview":
            return gr.update(visible=True)
        else:
            return gr.update(visible=False)

    model_radio.change(
        fn=update_size_visibility,
        inputs=[model_radio],
        outputs=[gen_size]
    )

    gen_button.click(
        fn=gradio_generate,
        inputs=[gemini_api_key_input, model_radio, gen_aspect_ratio, gen_size, gen_prompt,
               gen_save_dataset, gen_dataset_folder, gen_custom_filename],
        outputs=[gen_output, gen_status]
    )

    # ===== Gradio純正APIエンドポイント =====

    # API専用の画像生成関数（UIを持たない）
    def api_generate(
        prompt: str,
        gemini_api_key: str,
        model: str = "gemini-2.5-flash-image",
        aspect_ratio: str = "1:1",
        size: str = "1K",
        save_to_dataset: bool = True,
        dataset_folder: str = "",
        custom_filename: str = "",
        return_image_data: bool = False
    ) -> dict:
        """
        API endpoint for image generation

        Args:
            prompt: 画像生成プロンプト
            gemini_api_key: Gemini APIキー
            model: モデル名 (gemini-2.5-flash-image または gemini-3-pro-image-preview)
            aspect_ratio: アスペクト比 (1:1, 4:3, 3:4, 16:9, 9:16, 3:2)
            size: 画像サイズ (1K, 2K, 4K) - Gemini 3 Proのみ有効
            save_to_dataset: Datasetに保存するか
            dataset_folder: Datasetフォルダ名
            custom_filename: カスタムファイル名
            return_image_data: Base64画像データを含めるか
        """
        try:
            if not prompt:
                return {"error": "Prompt is required", "success": False}

            if not gemini_api_key or not gemini_api_key.strip():
                return {"error": "gemini_api_key is required", "success": False}

            # Validate model
            if model not in AVAILABLE_MODELS:
                return {"error": f"Invalid model. Available: {list(AVAILABLE_MODELS.keys())}", "success": False}

            # Validate aspect_ratio
            valid_aspect_ratios = ["1:1", "4:3", "3:4", "16:9", "9:16", "3:2"]
            if aspect_ratio not in valid_aspect_ratios:
                return {"error": f"Invalid aspect_ratio. Available: {valid_aspect_ratios}", "success": False}

            # Validate size (only for Gemini 3 Pro)
            valid_sizes = ["1K", "2K", "4K"]
            if size not in valid_sizes:
                return {"error": f"Invalid size. Available: {valid_sizes}", "success": False}

            # Generate image
            image = generate_image_with_gemini(prompt, gemini_api_key, model, aspect_ratio, size)

            # Save image locally
            if custom_filename and custom_filename.strip():
                clean_name = os.path.splitext(custom_filename.strip())[0]
                clean_name = "".join(c if c.isalnum() or c in '-_' else '_' for c in clean_name)
                if not clean_name:
                    clean_name = f"api_gen_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
                filename = f"{clean_name}.png"
            else:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"api_gen_{timestamp}.png"

            filepath = GENERATED_DIR / filename
            image.save(filepath)

            # Save to dataset if enabled
            dataset_url = None
            if dataset_manager and save_to_dataset:
                try:
                    metadata = {
                        "aspect_ratio": aspect_ratio,
                        "size": size if model == "gemini-3-pro-image-preview" else "1K",
                        "model": model,
                        "generation_type": "text-to-image"
                    }
                    dataset_url = dataset_manager.save_image(
                        image=image,
                        prompt=prompt,
                        folder_name=dataset_folder if dataset_folder.strip() else None,
                        filename=custom_filename.strip() if custom_filename.strip() else None,
                        metadata=metadata
                    )
                except Exception as dataset_error:
                    logger.error(f"Failed to save to dataset: {dataset_error}")

            response_data = {
                "success": True,
                "filename": filename,
                "local_path": f"/file=generated_images/{filename}",
                "prompt": prompt,
                "aspect_ratio": aspect_ratio,
                "size": size if model == "gemini-3-pro-image-preview" else "1K (fixed)",
                "model": model
            }

            if dataset_url:
                response_data["dataset_url"] = dataset_url

            # Base64エンコードされた画像データを含める（オプション）
            if return_image_data:
                import base64
                buffer = BytesIO()
                image.save(buffer, format="PNG")
                img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
                response_data["image_base64"] = img_base64

            return response_data

        except Exception as e:
            logger.error(f"API generation error: {e}")
            return {"error": str(e), "success": False}

    # API専用エンドポイントとして公開（UIには表示されない）
    gr.api(api_generate, api_name="generate")

    # Health check endpoint
    def api_health() -> dict:
        """Health check endpoint"""
        from datetime import datetime
        return {
            "status": "healthy",
            "timestamp": datetime.utcnow().isoformat() + "Z",
            "version": "9.1.0",
            "available_models": AVAILABLE_MODELS
        }

    gr.api(api_health, api_name="health")

    # Models endpoint
    def api_models() -> dict:
        """Get available models"""
        return {"models": AVAILABLE_MODELS}

    gr.api(api_models, api_name="models")

    # Footer
    gr.Markdown(
        """
        ---
        Powered by **Google Gemini AI** 🍌
        """
    )

# Run Gradio app directly (no FastAPI mounting)
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )