import random
import os
import uuid
from datetime import datetime
import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import DiffusionPipeline
from PIL import Image

# Create permanent storage directory
SAVE_DIR = "saved_images"  # Gradio will handle the persistence
if not os.path.exists(SAVE_DIR):
    os.makedirs(SAVE_DIR, exist_ok=True)

device = "cuda" if torch.cuda.is_available() else "cpu"
repo_id = "black-forest-labs/FLUX.1-dev"
adapter_id = "openfree/van-gogh"  # Changed to van Gogh model

pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
pipeline.load_lora_weights(adapter_id)
pipeline = pipeline.to(device)

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

def save_generated_image(image, prompt):
    # Generate unique filename with timestamp
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    unique_id = str(uuid.uuid4())[:8]
    filename = f"{timestamp}_{unique_id}.png"
    filepath = os.path.join(SAVE_DIR, filename)
    
    # Save the image
    image.save(filepath)
    
    # Save metadata
    metadata_file = os.path.join(SAVE_DIR, "metadata.txt")
    with open(metadata_file, "a", encoding="utf-8") as f:
        f.write(f"{filename}|{prompt}|{timestamp}\n")
    
    return filepath

@spaces.GPU(duration=60)
def inference(
    prompt,
    seed=42,
    randomize_seed=True,
    width=1024,
    height=768,
    guidance_scale=3.5,
    num_inference_steps=30,
    lora_scale=1.0,
    progress=None,
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(int(seed))
    
    image = pipeline(
        prompt=prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
        joint_attention_kwargs={"scale": lora_scale},
    ).images[0]
    
    # Save the generated image
    filepath = save_generated_image(image, prompt)
    
    # Return just the image and seed
    return image, seed

# Updated examples with 1880s clothing style for Van Gogh
examples = [
    "van Gogh's painting of a lively outdoor gathering in the 1880s, with men in formal top hats, frock coats, and women in bustled dresses with elaborate hats, enjoying a summer afternoon. The scene captures the Belle Époque atmosphere with swirling brushstrokes and vibrant colors typical of van Gogh, highlighting the fashionable attire of the period. [trigger]",
    "van Gogh's intimate portrait of a young woman from the 1880s, with her hair styled in a fashionable updo, wearing a high-necked dress with lace details and leg-of-mutton sleeves. She wears delicate jewelry and has the expressive, bold brushstrokes and vibrant color palette characteristic of van Gogh's portraiture. [trigger]",
    "van Gogh's painting of two young girls in 1880s attire seated at a piano. One plays while the other stands nearby, both dressed in white frocks with ribbon details, sashes, and high collars typical of the period. The interior setting features swirling patterns, intense colors, and the bold impasto technique typical of van Gogh's style. [trigger]",
    "van Gogh's painting of an elegant 1880s boating party, with gentlemen in striped boating blazers, straw boater hats, and formal trousers, alongside ladies in bustled day dresses with parasols. The scene captures the leisure activities of French society during the Belle Époque era, rendered with the artist's characteristic swirling brushstrokes and vibrant contrasting colors. [trigger]",
    "van Gogh's painting of children playing in an 1880s garden scene, dressed in formal period children's wear including sailor suits for boys and pinafores with full skirts for girls. Their Victorian-era clothing contrasts with their playful activities, set against van Gogh's vibrant treatment of nature with bold, expressive brushstrokes and dramatic color combinations. [trigger]",
    "van Gogh's depiction of bathers in 1880s swimming attire, showing the modest bathing costumes of the period. Women wear full-coverage dark bathing dresses with stockings, while men are in knee-length swimming suits. The figures are arranged against a backdrop of water and landscape rendered in van Gogh's distinctive swirling brushwork and bold chromatic contrasts. [trigger]"
]

# Improved custom CSS with better visuals - updated colors for Van Gogh theme
custom_css = """
:root {
    --color-primary: #1A3A8F;
    --color-secondary: #E09900;
    --background-fill-primary: linear-gradient(to right, #f6f9fc, #f0f4e3);
}

footer {
    visibility: hidden;
}

.gradio-container {
    background: var(--background-fill-primary);
}

.title {
    color: var(--color-primary) !important;
    font-size: 3rem !important;
    font-weight: 700 !important;
    text-align: center;
    margin: 1rem 0;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.05);
    font-family: 'Playfair Display', serif;
}

.subtitle {
    color: #4A5568 !important;
    font-size: 1.2rem !important;
    text-align: center;
    margin-bottom: 1.5rem;
    font-style: italic;
}

.collection-link {
    text-align: center;
    margin-bottom: 2rem;
    font-size: 1.1rem;
}

.collection-link a {
    color: var(--color-primary);
    text-decoration: underline;
    transition: color 0.3s ease;
}

.collection-link a:hover {
    color: var(--color-secondary);
}

.model-description {
    background-color: rgba(255, 255, 255, 0.8);
    border-radius: 12px;
    padding: 24px;
    margin: 20px 0;
    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.05);
    border-left: 5px solid var(--color-primary);
}

button.primary {
    background-color: var(--color-primary) !important;
    transition: all 0.3s ease;
}

button:hover {
    transform: translateY(-2px);
    box-shadow: 0 5px 15px rgba(0,0,0,0.1);
}

.input-container {
    border-radius: 10px;
    box-shadow: 0 2px 8px rgba(0,0,0,0.05);
}

.advanced-settings {
    margin-top: 1rem;
    padding: 1rem;
    border-radius: 10px;
    background-color: rgba(255, 255, 255, 0.6);
}

.example-region {
    background-color: rgba(255, 255, 255, 0.5);
    border-radius: 10px;
    padding: 1rem;
    margin-top: 1rem;
}
"""

with gr.Blocks(css=custom_css, analytics_enabled=False) as demo:
    gr.HTML('<div class="title">Vincent van Gogh STUDIO</div>')
    
    # Add collection link below title
    gr.HTML('<div class="collection-link"><a href="https://huggingface.co/collections/openfree/painting-art-ai-681453484ec15ef5978bbeb1" target="_blank">View the full Painting Art AI Collection</a></div>')
    
    # Model description with the requested content
    with gr.Group(elem_classes="model-description"):
        gr.HTML('<p>Generate beautiful artwork in the style of Vincent van Gogh. Add [trigger] at the end of your prompt for best results.</p>')

    # Simplified structure without tabs and gallery
    with gr.Column(elem_id="col-container"):
        with gr.Row(elem_classes="input-container"):
            prompt = gr.Text(
                label="Prompt",
                max_lines=1,
                placeholder="Enter your prompt (add [trigger] at the end)",
                value=examples[0]  # Set default text instead of generating an image
            )
            run_button = gr.Button("Generate", variant="primary", scale=0)

        result = gr.Image(label="Generated Image")
        seed_output = gr.Number(label="Seed", visible=True)

        with gr.Accordion("Advanced Settings", open=False, elem_classes="advanced-settings"):
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=42,
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=768,
                )

            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.0,
                    maximum=10.0,
                    step=0.1,
                    value=3.5,
                )
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=30,
                )
                lora_scale = gr.Slider(
                    label="LoRA scale",
                    minimum=0.0,
                    maximum=1.0,
                    step=0.1,
                    value=1.0,
                )

        with gr.Group(elem_classes="example-region"):
            gr.Markdown("### Examples")
            gr.Examples(
                examples=examples,
                inputs=prompt,
                outputs=None,  # Don't auto-run examples
                fn=None,  # No function to run for examples - just fill the prompt
                cache_examples=False,  # Disable caching
            )

    # Event handlers
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=inference,
        inputs=[
            prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            lora_scale,
        ],
        outputs=[result, seed_output],
    )
    
    # No preloading or automatic image generation

demo.queue()
demo.launch()