import gradio as gr
import os
import random
import torch
from PIL import Image, ImageDraw, ImageFont
from diffusers import DiffusionPipeline, AutoencoderKL
import spaces

# Global pipeline
_pipeline = None
_device = "cuda" if torch.cuda.is_available() else "cpu"

def load_lora():
    global _pipeline
    if _pipeline is not None:
        return _pipeline
    
    # Load SDXL + LoRA
    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
    _pipeline = DiffusionPipeline.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0",
        vae=vae,
        torch_dtype=torch.float16,
        use_safetensors=True,
        variant="fp16"
    )
    
    # Load LoRA with 1.0 influence (reduced from 1.5 to avoid Bitcoin bias)
    if os.path.exists("./adapter_model.safetensors"):
        _pipeline.load_lora_weights(".", weight_name="adapter_model.safetensors")
        _pipeline.fuse_lora(lora_scale=1.0)
    elif os.path.exists("./crypto_cover_styles_lora.safetensors"):
        _pipeline.load_lora_weights(".", weight_name="crypto_cover_styles_lora.safetensors")
        _pipeline.fuse_lora(lora_scale=1.0)
    
    _pipeline = _pipeline.to(_device)
    if _device == "cuda":
        try:
            _pipeline.enable_memory_efficient_attention()
            _pipeline.enable_vae_slicing()
        except:
            pass
    
    return _pipeline

def add_title_watermark(image, title):
    if image.size != (1800, 900):
        image = image.resize((1800, 900), Image.Resampling.LANCZOS)
    
    if not title:
        return image
    
    # Convert and create overlay
    if image.mode != 'RGBA':
        image = image.convert('RGBA')
    overlay = Image.new('RGBA', image.size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(overlay)
    
    # Add title
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf", 72)
    except:
        font = ImageFont.load_default()
    
    title = title.upper()
    bbox = draw.textbbox((0, 0), title, font=font)
    text_width = bbox[2] - bbox[0]
    text_height = bbox[3] - bbox[1]
    
    x = (1800 - text_width) // 2
    y = (900 - text_height) // 2 - 40
    draw.text((x, y), title, font=font, fill=(255, 255, 255, 255))
    
    # Add watermark
    if os.path.exists("./genfinity-watermark.png"):
        try:
            watermark = Image.open("./genfinity-watermark.png").convert('RGBA')
            wm_width = int(1800 * 0.15)
            wm_ratio = watermark.size[1] / watermark.size[0]
            wm_height = int(wm_width * wm_ratio)
            watermark = watermark.resize((wm_width, wm_height), Image.Resampling.LANCZOS)
            
            wm_x = 1800 - wm_width - 30
            wm_y = 900 - wm_height - 30
            overlay.paste(watermark, (wm_x, wm_y), watermark)
        except:
            pass
    
    result = Image.alpha_composite(image, overlay)
    return result.convert('RGB')

@spaces.GPU
def generate(prompt: str, title: str = ""):
    try:
        pipeline = load_lora()
        
        # Improved prompting to reduce Bitcoin bias and encourage diversity
        enhanced_prompt = f"{prompt}, professional editorial magazine cover design, high quality cryptocurrency illustration, detailed, trending"
        negative_prompt = "low quality, blurry, text, watermark, signature, bad anatomy, poorly drawn, ugly, deformed, repetitive coin stacks, bitcoin dominance"
        
        with torch.autocast(_device):
            result = pipeline(
                prompt=enhanced_prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=20,  # Reduced from 30 for faster generation
                guidance_scale=7.5,
                width=1024,
                height=1024,
                generator=torch.Generator(_device).manual_seed(random.randint(0, 2**32-1))
            )
        
        image = result.images[0]
        final_image = add_title_watermark(image, title)
        
        return final_image
        
    except Exception as e:
        error_img = Image.new('RGB', (1800, 900), color=(50, 50, 50))
        draw = ImageDraw.Draw(error_img)
        draw.text((900, 450), f"Error: {str(e)[:50]}...", fill=(255, 255, 255), anchor="mm")
        return error_img

# Simple interface
demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="News Prompt", value="XRP reaches new all-time highs"),
        gr.Textbox(label="Cover Title", value="XRP Breaks Records")
    ],
    outputs=gr.Image(label="Generated Cover (1800x900)", type="pil"),
    title="🚀 Your Trained LoRA Crypto Cover Generator",
    description="*Using YOUR trained LoRA + VALOR title system + Genfinity watermark (1800x900)*"
)

if __name__ == "__main__":
    demo.launch()