import spaces
import gradio as gr
import torch
from PIL import Image
from transformers import AutoProcessor
from longcat_image.models import LongCatImageTransformer2DModel
from longcat_image.pipelines import LongCatImageEditPipeline, LongCatImagePipeline
import numpy as np

# 1. DEFINE THE CUSTOM THEME CLASS (Fixed: Removing invalid keys)
class AppleStyleTheme(gr.themes.Base):
    def __init__(
        self,
        # Core Colors: Apple uses a specific blue, crisp grays, and clean white
        primary_hue=gr.themes.colors.blue,
        secondary_hue=gr.themes.colors.gray,
        neutral_hue=gr.themes.colors.neutral,
        # Core Sizing: Larger radius for that "iOS Card" look
        radius_size=gr.themes.sizes.radius_lg, 
        # Core Fonts: System fonts are key for the Apple feel
        font=["system-ui", "-apple-system", "BlinkMacSystemFont", "Segoe UI", "sans-serif"],
        font_mono=gr.themes.GoogleFont("IBM Plex Mono"),
        spacing_size=gr.themes.sizes.spacing_lg,
        text_size=gr.themes.sizes.text_md,
        **kwargs
    ):
        super().__init__(
            primary_hue=primary_hue,
            secondary_hue=secondary_hue,
            neutral_hue=neutral_hue,
            radius_size=radius_size,
            font=font,
            font_mono=font_mono,
            spacing_size=spacing_size,
            text_size=text_size,
            **kwargs
        )
        # Use .set() to modify specific variables. 
        # Removed 'tabs_border_color' which caused the error.
        self.set(
            # Global Backgrounds
            body_background_fill="*neutral_50", 
            block_background_fill="white",
            
            # Block (Card) Styles - Clean, thin borders, no deep shadows
            block_border_color="*neutral_200",
            block_border_width="1px",
            block_shadow="none",
            
            # Primary Button Styles (The "Apple Blue" button)
            button_primary_background_fill="*primary_600",
            button_primary_background_fill_hover="*primary_500",
            button_primary_text_color="white",
            button_primary_border_color="*primary_600",
            
            # Input Styles
            input_background_fill="*neutral_50",
            input_border_color="*neutral_200",
            input_shadow="none",
            input_shadow_focus="0 0 0 2px *primary_200",
        )

# Instantiate the theme
apple_theme = AppleStyleTheme()

# --- Model Loading ---
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Text-to-Image Model
t2i_model_id = 'meituan-longcat/LongCat-Image'
print(f"🔄 Loading Text-to-Image model from {t2i_model_id}...")
t2i_text_processor = AutoProcessor.from_pretrained(
    t2i_model_id,
    subfolder='tokenizer'
)

t2i_transformer = LongCatImageTransformer2DModel.from_pretrained(
    t2i_model_id,
    subfolder='transformer',
    torch_dtype=torch.bfloat16,
    use_safetensors=True
).to(device)

t2i_pipe = LongCatImagePipeline.from_pretrained(
    t2i_model_id,
    transformer=t2i_transformer,
    text_processor=t2i_text_processor,
)
t2i_pipe.to(device, torch.bfloat16)

print(f"✅ Text-to-Image model loaded successfully")

# Image Edit Model
edit_model_id = 'meituan-longcat/LongCat-Image-Edit'
print(f"🔄 Loading Image Edit model from {edit_model_id}...")
edit_text_processor = AutoProcessor.from_pretrained(
    edit_model_id,
    subfolder='tokenizer'
)

edit_transformer = LongCatImageTransformer2DModel.from_pretrained(
    edit_model_id,
    subfolder='transformer',
    torch_dtype=torch.bfloat16,
    use_safetensors=True
).to(device)

edit_pipe = LongCatImageEditPipeline.from_pretrained(
    edit_model_id,
    transformer=edit_transformer,
    text_processor=edit_text_processor,
)
edit_pipe.to(device, torch.bfloat16)

print(f"✅ Image Edit model loaded successfully on {device}")

# --- Core Functions ---
@spaces.GPU(duration=120)
def generate_image(
    prompt: str,
    width: int,
    height: int,
    seed: int,
    progress=gr.Progress()
):
    """Generate image from text prompt"""
    if not prompt or prompt.strip() == "":
        raise gr.Error("Please enter a prompt")
    try:
        progress(0.1, desc="Preparing generation...")
        progress(0.2, desc="Generating image...")
        generator = torch.Generator("cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
        with torch.inference_mode():
            output = t2i_pipe(
                prompt,
                negative_prompt="",
                height=height,
                width=width,
                guidance_scale=4.5,
                num_inference_steps=50,
                num_images_per_prompt=1,
                generator=generator,
                enable_cfg_renorm=True,
                enable_prompt_rewrite=True
            )
        progress(1.0, desc="Done!")
        return output.images[0]
    except Exception as e:
        raise gr.Error(f"Error during image generation: {str(e)}")

@spaces.GPU(duration=120)
def edit_image(
    input_image: Image.Image,
    prompt: str,
    seed: int,
    progress=gr.Progress()
):
    """Edit image based on text prompt"""
    if input_image is None:
        raise gr.Error("Please upload an image first")
    if not prompt or prompt.strip() == "":
        raise gr.Error("Please enter an edit instruction")
    try:
        progress(0.1, desc="Preparing image...")
        if input_image.mode != 'RGB':
            input_image = input_image.convert('RGB')
        progress(0.2, desc="Generating edited image...")
        generator = torch.Generator("cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
        with torch.inference_mode():
            output = edit_pipe(
                input_image,
                prompt,
                negative_prompt="",
                guidance_scale=4.5,
                num_inference_steps=50,
                num_images_per_prompt=1,
                generator=generator
            )
        progress(1.0, desc="Done!")
        return output.images[0]
    except Exception as e:
        raise gr.Error(f"Error during image editing: {str(e)}")

# --- Examples ---
edit_example_image_url = "https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo.png"
edit_example_data = [
    [edit_example_image_url, "Add a mustache", 42],
]

t2i_example_prompts = [
    ["一个年轻的亚裔女性，身穿黄色针织衫，搭配白色项链。她的双手放在膝盖上，表情恬静。背景是一堵粗糙的砖墙，午后的阳光温暖地洒在她身上，营造出一种宁静而温馨的氛围。", 1344, 768, 43],
    ["A serene mountain landscape at sunset with golden clouds", 1344, 768, 42],
    ["A cute robot sitting at a desk, digital art style", 1024, 1024, 44],
]

# Build Gradio interface
# Pass theme=apple_theme here. fill_width=True ensures it expands horizontally.
with gr.Blocks(fill_width=True) as demo:
    gr.HTML("""
        <div style="text-align: center; padding: 40px 20px 30px 20px;">
            <h1 style="font-size: 48px; font-weight: 700; margin: 0; color: #1d1d1f;">
                LongCat Studio
            </h1>
            <p style="font-size: 20px; color: #86868b; margin-top: 12px; font-weight: 400;">
                AI-powered image generation and editing
            </p>
        </div>
    """)

    with gr.Tabs(selected=0):
        # Image Edit Tab
        with gr.TabItem("Edit Image", id=0):
            # gr.Row with NO variant or scale arguments on the row itself
            with gr.Row(): 
                # Left Column: Inputs
                with gr.Column(scale=1, min_width=0, variant="panel"): 
                    gr.Markdown("### 🖼️ Input Image & Controls")
                    input_image = gr.Image(
                        label="Upload Image",
                        type="pil",
                        sources=["upload", "clipboard"],
                        height=450,
                    )

                    prompt = gr.Textbox(
                        label="What would you like to change?",
                        placeholder="e.g., Add a mustache, Change to sunset, Make it vintage...",
                        lines=2,
                        max_lines=3
                    )

                    seed = gr.Slider(
                        minimum=0,
                        maximum=999999,
                        value=42,
                        step=1,
                        label="Seed",
                        visible=False
                    )

                    edit_btn = gr.Button("Edit Image", variant="primary", size="lg")

                # Right Column: Output
                with gr.Column(scale=1, min_width=0, variant="panel"):
                    gr.Markdown("### ✨ Result")
                    output_image = gr.Image(
                        label="Result",
                        type="pil",
                        height=450,
                    )

            gr.HTML("<div style='margin: 30px 0 20px 0;'></div>")

            gr.Examples(
                examples=edit_example_data,
                inputs=[input_image, prompt, seed],
                outputs=output_image,
                fn=edit_image,
                cache_examples=False,
                label="Try an example",
                examples_per_page=3
            )

        # Text-to-Image Tab
        with gr.TabItem("Generate Image", id=1):
            with gr.Row(): 
                # Left Column: Inputs
                with gr.Column(scale=1, min_width=0, variant="panel"):
                    gr.Markdown("### 🎨 Generation Controls")
                    t2i_prompt = gr.Textbox(
                        label="Describe your image",
                        placeholder="e.g., A serene mountain landscape at sunset...",
                        lines=4,
                        max_lines=6
                    )

                    t2i_width = gr.Slider(
                        minimum=512,
                        maximum=2048,
                        value=1344,
                        step=64,
                        label="Width",
                    )

                    t2i_height = gr.Slider(
                        minimum=512,
                        maximum=2048,
                        value=768,
                        step=64,
                        label="Height",
                    )

                    t2i_seed = gr.Slider(
                        minimum=0,
                        maximum=999999,
                        value=42,
                        step=1,
                        label="Seed",
                        visible=False
                    )

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

                # Right Column: Output
                with gr.Column(scale=1, min_width=0, variant="panel"):
                    gr.Markdown("### ✨ Result")
                    t2i_output = gr.Image(
                        label="Result",
                        type="pil",
                        height=550,
                    )

            gr.HTML("<div style='margin: 30px 0 20px 0;'></div>")

            gr.Examples(
                examples=t2i_example_prompts,
                inputs=[t2i_prompt, t2i_width, t2i_height, t2i_seed],
                outputs=t2i_output,
                fn=generate_image,
                cache_examples=False,
                label="Try an example",
                examples_per_page=3
            )

    # Event handlers
    generate_btn.click(
        fn=generate_image,
        inputs=[t2i_prompt, t2i_width, t2i_height, t2i_seed],
        outputs=t2i_output,
    )

    edit_btn.click(
        fn=edit_image,
        inputs=[input_image, prompt, seed],
        outputs=output_image,
    )

    # Footer
    gr.HTML("""
        <div style="text-align: center; margin-top: 60px; padding: 30px 20px; border-top: 1px solid #d2d2d7;">
            <p style="color: #86868b; font-size: 13px; margin: 0;">
                Powered by LongCat • Built with 
                <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="color: #007aff; text-decoration: none;">anycoder</a>
            </p>
        </div>
    """)

# Launch the app
if __name__ == "__main__":
    demo.launch(
        mcp_server=True,
        theme=apple_theme,
    )