#!/usr/bin/env python3

import gradio as gr
import numpy as np
import spaces
import torch
import random
import base64
import io
import math
from PIL import Image
from diffusers import FluxKontextPipeline
from diffusers import FluxTransformer2DModel
from diffusers.utils import load_image
from diffusers import EulerDiscreteScheduler

from huggingface_hub import hf_hub_download

# Load pipeline with device detection
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

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

pipe = FluxKontextPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Kontext-dev", 
    torch_dtype=dtype
).to(device)

# Keep the default scheduler that works with FluxKontextPipeline
print(f"📅 Using default scheduler: {pipe.scheduler.__class__.__name__}")

pipe.load_lora_weights(
    "DoozyWo/Kontext_avatar_LoRA", 
    weight_name="Avataar_LoRA_000003000.safetensors", 
    adapter_name="avatar_lora"
)
pipe.set_adapters(["avatar_lora"], adapter_weights=[1.0])

MAX_SEED = np.iinfo(np.int32).max

def scale_image_to_megapixels(image, target_megapixels=1.0):
    """Scale image to target megapixels like ComfyUI ImageScaleToTotalPixels"""
    original_width, original_height = image.size
    original_pixels = original_width * original_height
    target_pixels = target_megapixels * 1000000
    
    if original_pixels <= target_pixels:
        return image
    
    scale_factor = math.sqrt(target_pixels / original_pixels)
    new_width = int(original_width * scale_factor)
    new_height = int(original_height * scale_factor)
    
    # Ensure dimensions are divisible by 8 for better model performance
    new_width = (new_width // 8) * 8
    new_height = (new_height // 8) * 8
    
    return image.resize((new_width, new_height), Image.Resampling.BICUBIC)

@spaces.GPU
def transform_to_avatar(input_image, custom_prompt="", guidance_scale=1.0, num_inference_steps=25, seed=42, randomize_seed=False, lora_strength=1.0):
    """
    Transform image to Avatar character using FLUX Kontext + Avatar LoRA
    
    Args:
        input_image: Input PIL Image
        custom_prompt: Custom prompt for additional details
        guidance_scale: How closely to follow the prompt
        num_inference_steps: Quality vs speed
        seed: Random seed
        randomize_seed: Whether to randomize seed
        lora_strength: Strength of Avatar LoRA
    
    Returns:
        tuple: (result_images_list, seed_used, prompt_used)
    """
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
        
    input_image = input_image.convert("RGB")
    
    # Scale image to 1 megapixel like ComfyUI
    scaled_image = scale_image_to_megapixels(input_image, target_megapixels=1.0)
    print(f"📏 Image scaled from {input_image.size} to {scaled_image.size}")
    
    # Update LoRA strength if needed
    if lora_strength != 1.0:
        pipe.set_adapters(["avatar_lora"], adapter_weights=[lora_strength])
    
    # Construct Avatar transformation prompt (matching ComfyUI approach)
    base_prompt = "Turn this into a photorealistic Na'vi character from Avatar, with blue bioluminescent skin, large eyes, and set in the glowing jungle of Pandora."
    
    if custom_prompt.strip():
        full_prompt = f"{base_prompt} {custom_prompt.strip()}"
    else:
        full_prompt = base_prompt
    
    # Use the exact prompt structure from ComfyUI
    prompt_with_template = full_prompt  # No additional template wrapping
    
    print(f"Avatar transformation prompt: {prompt_with_template}")
    
    # Match ComfyUI parameters: cfg=1, steps=25
    # Let the pipeline handle dimensions automatically
    image = pipe(
        image=scaled_image, 
        prompt=prompt_with_template,
        guidance_scale=guidance_scale,  # CFG=1 from ComfyUI
        num_inference_steps=num_inference_steps,
        generator=torch.Generator(device=pipe.device).manual_seed(seed),
    ).images[0]
    
    return [input_image, image], seed, prompt_with_template

# Avatar style suggestions
AVATAR_STYLES = {
    "warrior": "fierce Na'vi warrior with battle scars, tribal war paint, intimidating expression",
    "mystical": "ethereal Na'vi with glowing tattoos, magical aura, spiritual energy radiating", 
    "forest_guardian": "nature-connected Na'vi, tree bark textures, forest camouflage patterns",
    "ceremonial": "ornate Na'vi with jewelry, ritual face paint, sacred symbols, regal bearing",
    "hunter": "stealth Na'vi with camouflage markings, predator instincts, keen alert expression",
    "shaman": "wise Na'vi spiritual leader, glowing eyes, mystical energy swirling around",
    "royalty": "regal Na'vi with elaborate headdress, noble posture, ornate decorations",
    "bioluminescent": "Na'vi with enhanced glowing patterns, bright bio-lights, phosphorescent skin"
}

def update_prompt_from_style(style_option):
    """Update the prompt textbox based on style selection"""
    if style_option == "custom":
        return ""
    else:
        return AVATAR_STYLES.get(style_option, "")

css="""
#col-container {
    margin: 0 auto;
    max-width: 1020px;
}
"""

with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# Avatar Transformation Studio 🧞‍♂️
        """)
        gr.Markdown(f"""Transform any portrait into a stunning Na'vi character using FLUX Kontext + Avatar LoRA ✨
         """)

        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Upload portrait for Avatar transformation", type="pil")
                
                with gr.Row():
                    
                    style_dropdown = gr.Dropdown(
                        choices=["custom"] + list(AVATAR_STYLES.keys()),
                        value="warrior",
                        label="Choose Avatar Style",
                        scale=2
                    )
                
                custom_prompt = gr.Textbox(
                    label="Custom Details",
                    info="Additional details for your Avatar transformation",
                    show_label=True,
                    max_lines=3,
                    placeholder="Select a style above or add custom details like 'glowing tattoos', 'forest setting', etc.",
                    value="fierce Na'vi warrior with battle scars, tribal war paint, intimidating expression",
                    container=True
                )
                
                run_button = gr.Button("🧞‍♂️ Transform to Avatar", scale=0, variant="primary")
                
                with gr.Accordion("Advanced Settings", open=False):
            
                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=42,
                    )
                    
                    randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
                    
                    guidance_scale = gr.Slider(
                        label="Guidance Scale (CFG)",
                        minimum=1,
                        maximum=10,
                        step=0.1,
                        value=1.0,
                        info="ComfyUI uses CFG=1"
                    )
                    
                    num_inference_steps = gr.Slider(
                        label="Inference Steps",
                        minimum=10,
                        maximum=50,
                        step=1,
                        value=25,
                    )
                    
                    lora_strength = gr.Slider(
                        label="Avatar LoRA Strength",
                        minimum=0.5,
                        maximum=1.5,
                        step=0.1,
                        value=1.0,
                    )
                    
            with gr.Column():
                result = gr.ImageSlider(label="Avatar Transformation Result", show_label=False, interactive=False)
                final_prompt = gr.Textbox(label="Transformation prompt used", info="The actual prompt used for the Avatar transformation")
        
        # Update prompt when dropdown changes
        style_dropdown.change(
            fn=update_prompt_from_style,
            inputs=[style_dropdown],
            outputs=[custom_prompt]
        )
        
        gr.Examples(
            examples=[
                ["./examples/portrait1.jpg", "fierce warrior with glowing blue war paint", "warrior", 1.0, 25, 42, False, 1.0],
                ["./examples/portrait2.jpg", "mystical shaman with ethereal glow and floating spirits", "shaman", 1.0, 25, 42, False, 1.0],
                ["./examples/portrait3.jpg", "regal princess with ornate jewelry and ceremonial markings", "royalty", 1.0, 25, 42, False, 1.0],
            ],
            inputs=[input_image, custom_prompt, style_dropdown, guidance_scale, num_inference_steps, seed, randomize_seed, lora_strength],
            outputs=[result, seed, final_prompt],
            fn=transform_to_avatar,
            cache_examples="lazy"
        )
    
        gr.on(
            triggers=[run_button.click, custom_prompt.submit],
            fn=transform_to_avatar,
            inputs=[input_image, custom_prompt, guidance_scale, num_inference_steps, seed, randomize_seed, lora_strength],
            outputs=[result, seed, final_prompt]
        )

# MCP Server functionality
from mcp.server import NotificationOptions, Server
from mcp.server.models import InitializationOptions
import mcp.server.stdio
import mcp.types as types
import asyncio

app = Server("avatar-transformation")

@app.list_tools()
async def handle_list_tools() -> types.ListToolsResult:
    return types.ListToolsResult(
        tools=[
            types.Tool(
                name="transform_to_avatar",
                description="Transform portrait to Avatar character using FLUX Kontext + Avatar LoRA",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "image_base64": {"type": "string", "description": "Base64 encoded image"},
                        "custom_prompt": {"type": "string", "default": ""},
                        "guidance_scale": {"type": "number", "default": 1.0},
                        "num_inference_steps": {"type": "integer", "default": 25},
                        "seed": {"type": "integer", "default": 42},
                        "randomize_seed": {"type": "boolean", "default": False},
                        "lora_strength": {"type": "number", "default": 1.0}
                    },
                    "required": ["image_base64"]
                }
            )
        ]
    )

@app.call_tool()
async def handle_call_tool(request: types.CallToolRequest) -> types.CallToolResult:
    if request.name == "transform_to_avatar":
        try:
            args = request.arguments or {}
            
            # Decode base64 image
            image_data = args.get("image_base64")
            if not image_data:
                return types.CallToolResult(
                    content=[types.TextContent(type="text", text="Error: No image provided")],
                    isError=True
                )
            
            image_bytes = base64.b64decode(image_data)
            input_image = Image.open(io.BytesIO(image_bytes))
            
            # Transform image
            result_images, used_seed, used_prompt = transform_to_avatar(
                input_image=input_image,
                custom_prompt=args.get("custom_prompt", ""),
                guidance_scale=args.get("guidance_scale", 1.0),
                num_inference_steps=args.get("num_inference_steps", 25),
                seed=args.get("seed", 42),
                randomize_seed=args.get("randomize_seed", False),
                lora_strength=args.get("lora_strength", 1.0)
            )
            
            # Convert result to base64
            result_image = result_images[1]  # Transformed image
            buffer = io.BytesIO()
            result_image.save(buffer, format='PNG')
            result_base64 = base64.b64encode(buffer.getvalue()).decode()
            
            return types.CallToolResult(
                content=[types.TextContent(
                    type="text", 
                    text=f"Avatar transformation complete!\nSeed used: {used_seed}\nPrompt: {used_prompt}\nResult image (base64): {result_base64}"
                )]
            )
            
        except Exception as e:
            return types.CallToolResult(
                content=[types.TextContent(type="text", text=f"Error: {str(e)}")],
                isError=True
            )
    
    return types.CallToolResult(
        content=[types.TextContent(type="text", text=f"Unknown tool: {request.name}")],
        isError=True
    )

async def run_mcp_server():
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await app.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="avatar-transformation",
                server_version="1.0.0",
                capabilities=app.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        )

if __name__ == "__main__":
    # Check if running as MCP server or Gradio demo
    import sys
    if "--mcp" in sys.argv:
        asyncio.run(run_mcp_server())
    else:
        # Launch with MCP server support
        demo.launch(mcp_server=True)