import spaces
import gradio as gr
import random
import os
import time
import torch
from diffusers import FluxPipeline
from transformers import pipeline

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {DEVICE}")

DEFAULT_HEIGHT = 1024
DEFAULT_WIDTH = 1024
DEFAULT_GUIDANCE_SCALE = 3.5
DEFAULT_NUM_INFERENCE_STEPS = 15
DEFAULT_MAX_SEQUENCE_LENGTH = 512
HF_TOKEN = os.environ.get("HF_ACCESS_TOKEN")

# Cache for the pipelines
CACHED_PIPE = None
CACHED_LLM_PIPE = None

def load_llm_pipeline():
    """Load the LLM pipeline for prompt enhancement"""
    global CACHED_LLM_PIPE
    if CACHED_LLM_PIPE is not None:
        return CACHED_LLM_PIPE
    
    print("Loading LLM pipeline for prompt enhancement...")
    try:
        # Note: Using a smaller model that's actually available
        # You can replace this with "openai/gpt-oss-120b" if you have access
        llm_pipe = pipeline(
            "text-generation", 
            model="microsoft/Phi-3-mini-4k-instruct",  # Alternative smaller model
            torch_dtype=torch.bfloat16,
            device_map="auto"
        )
        CACHED_LLM_PIPE = llm_pipe
        print("LLM pipeline loaded successfully")
        return llm_pipe
    except Exception as e:
        print(f"Error loading LLM pipeline: {e}")
        # Fallback to a simpler model if the main one fails
        try:
            llm_pipe = pipeline(
                "text-generation",
                model="gpt2",  # Fallback to GPT-2
                device_map="auto"
            )
            CACHED_LLM_PIPE = llm_pipe
            print("Loaded fallback LLM pipeline (GPT-2)")
            return llm_pipe
        except Exception as e2:
            print(f"Error loading fallback LLM pipeline: {e2}")
            return None

def enhance_prompt(prompt, progress=gr.Progress()):
    """Enhance the prompt using LLM"""
    if not prompt:
        return prompt, "Please enter a prompt first."
    
    progress(0.3, desc="Enhancing prompt with AI...")
    
    try:
        llm_pipe = load_llm_pipeline()
        if llm_pipe is None:
            return prompt, "LLM pipeline not available, using original prompt."
        
        # Create enhancement prompt
        messages = [
            {
                "role": "system", 
                "content": "You are a helpful assistant that enhances image generation prompts. Make prompts more detailed, artistic, and visually descriptive while keeping the core concept. Add details about lighting, style, colors, mood, and composition. Keep the enhanced prompt under 200 words."
            },
            {
                "role": "user", 
                "content": f"Enhance this image generation prompt, making it more detailed and artistic: '{prompt}'"
            }
        ]
        
        # Generate enhanced prompt
        result = llm_pipe(
            messages,
            max_new_tokens=200,
            temperature=0.7,
            do_sample=True,
            top_p=0.9
        )
        
        # Extract the enhanced prompt from the response
        if isinstance(result, list) and len(result) > 0:
            enhanced = result[0].get('generated_text', '')
            # Extract only the assistant's response
            if isinstance(enhanced, list):
                for msg in enhanced:
                    if msg.get('role') == 'assistant':
                        enhanced = msg.get('content', prompt)
                        break
            elif isinstance(enhanced, str):
                # Clean up the response if needed
                enhanced = enhanced.strip()
                if enhanced.startswith("Enhanced prompt:"):
                    enhanced = enhanced.replace("Enhanced prompt:", "").strip()
            
            if enhanced and enhanced != prompt:
                return enhanced, "Prompt enhanced successfully!"
            else:
                return prompt, "Using original prompt."
        else:
            return prompt, "Enhancement failed, using original prompt."
            
    except Exception as e:
        print(f"Error during prompt enhancement: {e}")
        return prompt, f"Enhancement error: {e}. Using original prompt."

def load_bnb_4bit_pipeline():
    """Load the 4-bit quantized pipeline"""
    global CACHED_PIPE
    if CACHED_PIPE is not None:
        return CACHED_PIPE
    
    print("Loading 4-bit BNB pipeline...")
    MODEL_ID = "derekl35/FLUX.1-dev-nf4"
    
    start_time = time.time()
    try:
        pipe = FluxPipeline.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.bfloat16
        )
        pipe.enable_model_cpu_offload()
        end_time = time.time()
        mem_reserved = torch.cuda.memory_reserved(0)/1024**3 if DEVICE == "cuda" else 0
        print(f"4-bit BNB pipeline loaded in {end_time - start_time:.2f}s. Memory reserved: {mem_reserved:.2f} GB")
        CACHED_PIPE = pipe
        return pipe
    except Exception as e:
        print(f"Error loading 4-bit BNB pipeline: {e}")
        raise

@spaces.GPU(duration=240)
def generate_image(prompt, use_enhancement=False, progress=gr.Progress(track_tqdm=True)):
    """Generate image using 4-bit quantized model with optional prompt enhancement"""
    if not prompt:
        return None, prompt, "Please enter a prompt."
    
    enhanced_prompt = prompt
    enhancement_status = ""
    
    # Enhance prompt if requested
    if use_enhancement:
        enhanced_prompt, enhancement_status = enhance_prompt(prompt, progress)
    
    progress(0.5, desc="Loading 4-bit quantized model...")
    
    try:
        # Load the 4-bit pipeline
        pipe = load_bnb_4bit_pipeline()
        
        # Set up generation parameters
        pipe_kwargs = {
            "prompt": enhanced_prompt,
            "height": DEFAULT_HEIGHT,
            "width": DEFAULT_WIDTH,
            "guidance_scale": DEFAULT_GUIDANCE_SCALE,
            "num_inference_steps": DEFAULT_NUM_INFERENCE_STEPS,
            "max_sequence_length": DEFAULT_MAX_SEQUENCE_LENGTH,
        }
        
        # Generate seed
        seed = random.getrandbits(64)
        print(f"Using seed: {seed}")
        
        progress(0.7, desc="Generating image...")
        
        # Generate image
        gen_start_time = time.time()
        image = pipe(**pipe_kwargs, generator=torch.manual_seed(seed)).images[0]
        gen_end_time = time.time()
        
        print(f"Image generated in {gen_end_time - gen_start_time:.2f} seconds")
        mem_reserved = torch.cuda.memory_reserved(0)/1024**3 if DEVICE == "cuda" else 0
        print(f"Memory reserved: {mem_reserved:.2f} GB")
        
        status_msg = f"Generation complete! (Seed: {seed})"
        if enhancement_status:
            status_msg = f"{enhancement_status} | {status_msg}"
        
        return image, enhanced_prompt, status_msg
        
    except Exception as e:
        print(f"Error during generation: {e}")
        return None, enhanced_prompt, f"Error: {e}"

@spaces.GPU(duration=60)
def enhance_only(prompt, progress=gr.Progress()):
    """Only enhance the prompt without generating an image"""
    enhanced_prompt, status = enhance_prompt(prompt, progress)
    return enhanced_prompt, status

# Create Gradio interface
with gr.Blocks(title="FLUXllama Enhanced", theme=gr.themes.Soft()) as demo:
    gr.HTML(
        """
        <div style='text-align: center; margin-bottom: 20px;'>
            <h1>FLUXllama Enhanced</h1>
            <p>FLUX.1-dev 4-bit Quantized Version with AI Prompt Enhancement</p>
        </div>
        """
    )
    
    gr.HTML(
        """
        <div class='container' style='display:flex; justify-content:center; gap:12px; margin-bottom: 20px;'>
            <a href="https://huggingface.co/spaces/openfree/Best-AI" target="_blank">
                <img src="https://img.shields.io/static/v1?label=OpenFree&message=BEST%20AI%20Services&color=%230000ff&labelColor=%23000080&logo=huggingface&logoColor=%23ffa500&style=for-the-badge" alt="OpenFree badge">
            </a>
    
            <a href="https://discord.gg/openfreeai" target="_blank">
                <img src="https://img.shields.io/static/v1?label=Discord&message=Openfree%20AI&color=%230000ff&labelColor=%23800080&logo=discord&logoColor=white&style=for-the-badge" alt="Discord badge">
            </a>
        </div>
        """
    )
    
    with gr.Column():
        prompt_input = gr.Textbox(
            label="Enter your prompt",
            placeholder="e.g., A photorealistic portrait of an astronaut on Mars",
            lines=3
        )
        
        with gr.Row():
            enhance_checkbox = gr.Checkbox(
                label="🎨 Use AI Prompt Enhancement",
                value=False,
                info="Automatically enhance your prompt for better results"
            )
            enhance_only_button = gr.Button("✨ Enhance Only", variant="secondary", scale=1)
        
        enhanced_prompt_display = gr.Textbox(
            label="Enhanced Prompt (will appear after enhancement)",
            lines=3,
            interactive=False,
            visible=True
        )
        
        with gr.Row():
            generate_button = gr.Button("🚀 Generate Image", variant="primary", scale=2)
            generate_enhanced_button = gr.Button("🎨 Enhance & Generate", variant="primary", scale=2)
    
    output_image = gr.Image(
        label="Generated Image (4-bit Quantized)",
        type="pil",
        height=600
    )
    
    status_text = gr.Textbox(
        label="Status",
        interactive=False,
        lines=1
    )
    
    # Connect components
    generate_button.click(
        fn=lambda p: generate_image(p, use_enhancement=False),
        inputs=[prompt_input],
        outputs=[output_image, enhanced_prompt_display, status_text]
    )
    
    generate_enhanced_button.click(
        fn=lambda p: generate_image(p, use_enhancement=True),
        inputs=[prompt_input],
        outputs=[output_image, enhanced_prompt_display, status_text]
    )
    
    enhance_only_button.click(
        fn=enhance_only,
        inputs=[prompt_input],
        outputs=[enhanced_prompt_display, status_text]
    )
    
    # Enter key to submit (with enhancement checkbox consideration)
    prompt_input.submit(
        fn=generate_image,
        inputs=[prompt_input, enhance_checkbox],
        outputs=[output_image, enhanced_prompt_display, status_text]
    )
    
    # Example prompts
    gr.Examples(
        examples=[
            "A photorealistic portrait of an astronaut on Mars",
            "Water-color painting of a cat wearing sunglasses",
            "Neo-tokyo cyberpunk cityscape at night, rain-soaked streets, 8K",
            "A majestic dragon flying over a medieval castle at sunset",
            "Abstract art representing the concept of time and space",
            "Detailed oil painting of a steampunk clockwork city",
            "Underwater scene with bioluminescent creatures in deep ocean",
            "Japanese garden in autumn with falling maple leaves"
        ],
        inputs=prompt_input
    )
    
    gr.HTML(
        """
        <div style='text-align: center; margin-top: 20px; padding: 20px; background-color: #f0f0f0; border-radius: 10px;'>
            <h3>✨ Prompt Enhancement Feature</h3>
            <p>This app now includes AI-powered prompt enhancement! The enhancement feature will:</p>
            <ul style='text-align: left; display: inline-block;'>
                <li>Add artistic details and visual descriptions</li>
                <li>Specify lighting, mood, and atmosphere</li>
                <li>Include style and composition elements</li>
                <li>Make your prompts more effective for image generation</li>
            </ul>
            <p><strong>How to use:</strong></p>
            <p>1. Enter a simple prompt</p>
            <p>2. Click "✨ Enhance Only" to preview the enhanced version</p>
            <p>3. Click "🎨 Enhance & Generate" to enhance and generate in one step</p>
            <p>4. Or check the enhancement checkbox and click Generate</p>
        </div>
        """
    )

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