import os
import gradio as gr
import torch
import gdown
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from PIL import Image


# --- Download from Google Drive ---
def download_from_gdrive(file_id, dest_path):
    if not os.path.exists(dest_path):
        print(f"📥 Downloading {dest_path}...")
        gdown.download(f"https://drive.google.com/uc?id={file_id}", dest_path, quiet=False)
        print(f"✅ Downloaded {dest_path} from Google Drive")
    else:
        print(f"✅ {dest_path} already exists")


# Download base model and LoRA
print("🚀 Starting BitKun LoRA Generator...")
ckpt_id = "1OdP2SDB6MyR6JDK6_ekyqFHz91rkT2ZK"  # Your base model ID
download_from_gdrive(ckpt_id, "AnyLoRA_noVae_fp16-pruned.ckpt")

# ⚙️ Configuration
project_name = "bitkun"
epoch_number = 10
lora_path = f"{project_name}-{epoch_number:02d}.safetensors"
base_model_path = "AnyLoRA_noVae_fp16-pruned.ckpt"
output_folder = "generated_images"
os.makedirs(output_folder, exist_ok=True)

# Auto-detect device and set appropriate dtype
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32

print(f"🔧 Using device: {device} with dtype: {dtype}")

# 🚀 Load Base Model with optimizations
print("📦 Loading base model...")
try:
    pipe = StableDiffusionPipeline.from_single_file(
        base_model_path,
        torch_dtype=dtype,
        safety_checker=None,
        requires_safety_checker=False,
        use_safetensors=True
    ).to(device)

    # ⏩ Use faster scheduler
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(
        pipe.scheduler.config,
        use_karras_sigmas=True,
        algorithm_type="dpmsolver++"
    )

    # Enable memory efficient attention if available
    if hasattr(pipe, "enable_attention_slicing"):
        pipe.enable_attention_slicing()

    # Enable memory efficient attention
    if hasattr(pipe, "enable_xformers_memory_efficient_attention"):
        try:
            pipe.enable_xformers_memory_efficient_attention()
            print("✅ XFormers memory efficient attention enabled")
        except:
            print("⚠️ XFormers not available, using standard attention")

    print("✅ Base model loaded successfully!")

except Exception as e:
    print(f"❌ Error loading base model: {e}")
    raise e

# 🔄 Load LoRA
lora_loaded = False
try:
    if os.path.exists(lora_path):
        pipe.load_lora_weights(lora_path, adapter_name="default")
        pipe.set_adapters(["default"], adapter_weights=[0.8])
        lora_loaded = True
        print("✅ LoRA loaded and pipeline ready!")
    else:
        print(f"⚠️ LoRA file not found: {lora_path}")
        print("Pipeline will run with base model only.")
except Exception as e:
    print(f"⚠️ Could not load LoRA weights: {e}")
    print("Pipeline will run with base model only.")


# 🎨 Optimized Generation Function
def generate_bitkun(prompt, negative_prompt, num_images, steps, guidance_scale, width, height):
    if not prompt.strip():
        return [], "⚠️ Please enter a prompt!"

    # Add bitkun to prompt if not present
    if "bitkun" not in prompt.lower():
        prompt = f"bitkun, {prompt}"

    seed = 42
    images = []

    for i in range(num_images):
        try:
            print(f"🎨 Generating image {i + 1}/{num_images}...")
            generator = torch.Generator(device=device).manual_seed(seed + i)

            # Use autocast only for CUDA
            if device == "cuda":
                with torch.autocast("cuda"):
                    result = pipe(
                        prompt=prompt,
                        negative_prompt=negative_prompt,
                        num_inference_steps=steps,
                        guidance_scale=guidance_scale,
                        width=width,
                        height=height,
                        generator=generator
                    )
            else:
                result = pipe(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    num_inference_steps=steps,
                    guidance_scale=guidance_scale,
                    width=width,
                    height=height,
                    generator=generator
                )

            image = result.images[0]

            # Save image
            filename = f"{project_name}_custom_{i + 1}.png"
            filepath = os.path.join(output_folder, filename)
            image.save(filepath)
            images.append(image)

        except Exception as e:
            error_msg = f"❌ Error generating image {i + 1}: {str(e)}"
            print(error_msg)
            continue

    if not images:
        final_status = "❌ Failed to generate any images. Please try again with different settings."
    else:
        lora_status = "with LoRA" if lora_loaded else "without LoRA"
        final_status = f"🎉 Successfully generated {len(images)}/{num_images} image(s) {lora_status}!"

    return images, final_status


# 🌐 Gradio UI with Speed Presets
with gr.Blocks(title="BitKun LoRA Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🐶 BitKun LoRA Generator 🎨 (FAST VERSION)
    ### 自由なプロンプト対応 / Custom Prompt Support
    """)

    # Show LoRA status
    lora_status_text = "✅ LoRA loaded successfully!" if lora_loaded else "⚠️ Running with base model only (LoRA not found)"
    gr.Markdown(f"**Status:** {lora_status_text}")

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 🎯 Generation Settings")

            custom_prompt = gr.Textbox(
                label="📝 プロンプト / Prompt",
                placeholder="例: happy, smiling, cartoon style, colorful background",
                lines=3,
                info="'bitkun' will be automatically added if not included"
            )

            negative_prompt = gr.Textbox(
                label="🚫 ネガティブプロンプト / Negative Prompt",
                value="realistic, human skin, photo, blurry, distorted, extra limbs, bad anatomy",
                lines=2
            )

            # Speed Presets
            gr.Markdown("### ⚡ Speed Presets")
            with gr.Row():
                speed_preset = gr.Radio(
                    choices=[
                        ("🚀 Ultra Fast (10 steps, 256x256)", "ultra_fast"),
                        ("⚡ Fast (15 steps, 384x384)", "fast"),
                        ("🎯 Balanced (20 steps, 512x512)", "balanced"),
                        ("🎨 Quality (25 steps, 512x512)", "quality"),
                        ("🔧 Custom", "custom")
                    ],
                    value="fast",
                    label="Choose Speed vs Quality"
                )

            with gr.Row():
                num_images = gr.Slider(
                    label="🖼️ Number of Images",
                    minimum=1,
                    maximum=3,
                    value=1,
                    step=1,
                    info="More images = longer processing time"
                )

            # Advanced settings (hidden by default)
            with gr.Accordion("🔧 Advanced Settings", open=False):
                steps = gr.Slider(
                    label="🔄 Inference Steps",
                    minimum=5,
                    maximum=50,
                    value=15,
                    step=1,
                    info="More steps = higher quality but slower"
                )

                guidance_scale = gr.Slider(
                    label="🎚️ Guidance Scale",
                    minimum=1.0,
                    maximum=15.0,
                    value=7.5,
                    step=0.5,
                    info="Higher = more prompt adherence"
                )

                with gr.Row():
                    width = gr.Slider(
                        label="📏 Width",
                        minimum=256,
                        maximum=768,
                        value=384,
                        step=64
                    )

                    height = gr.Slider(
                        label="📐 Height",
                        minimum=256,
                        maximum=768,
                        value=384,
                        step=64
                    )

            generate_btn = gr.Button(
                "🎨 画像を生成 / Generate Images",
                variant="primary",
                size="lg"
            )

        with gr.Column(scale=1):
            gr.Markdown("### 🖼️ Generated Images")

            gallery = gr.Gallery(
                label="生成画像 / Generated Images",
                columns=2,
                rows=2,
                height=400,
                show_label=False
            )

            status_text = gr.Textbox(
                label="📊 Generation Status",
                interactive=False,
                lines=2
            )


    # Speed preset change handler
    def update_settings(preset):
        if preset == "ultra_fast":
            return 10, 7.0, 256, 256
        elif preset == "fast":
            return 15, 7.5, 384, 384
        elif preset == "balanced":
            return 20, 7.5, 512, 512
        elif preset == "quality":
            return 25, 8.0, 512, 512
        else:  # custom
            return 15, 7.5, 384, 384


    speed_preset.change(
        fn=update_settings,
        inputs=[speed_preset],
        outputs=[steps, guidance_scale, width, height]
    )

    # Example prompts section
    gr.Markdown("""
    ### 📝 Quick Examples / クイック例:

    **⚡ For fastest results, try these short prompts:**
    - `happy, smiling`
    - `sad, crying`  
    - `angry, red face`
    - `surprised, shocked`
    - `sleepy, tired`
    - `superhero, cape`
    """)

    # Button click event
    generate_btn.click(
        fn=generate_bitkun,
        inputs=[custom_prompt, negative_prompt, num_images, steps, guidance_scale, width, height],
        outputs=[gallery, status_text],
        show_progress=True
    )

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