import gradio as gr
import torch
from transformers import BlipProcessor, BlipForConditionalGeneration
import re
import os


MODEL_NAME = "Salesforce/blip-image-captioning-base"

print(f"📦 Using model: {MODEL_NAME}")
print(f"💾 Models stored in: {os.path.expanduser('~/.cache/huggingface/transformers/')}")

device = "cuda" if torch.cuda.is_available() else "cpu"

# Initialize BLIP model
try:
    processor = BlipProcessor.from_pretrained(MODEL_NAME)
    model = BlipForConditionalGeneration.from_pretrained(MODEL_NAME).to(device)
    MODEL_LOADED = True
    print(f"✅ Loaded BLIP model on {device} (~400MB)")
except Exception as e:
    print(f"❌ Error loading model: {e}")
    MODEL_LOADED = False
    model = None
    processor = None

def generate_caption_blip(image, style="detailed", max_length=50):
    """Generate caption using BLIP model"""
    try:
        inputs = processor(image, return_tensors="pt").to(device)
        
        # Different generation parameters for different styles
        if style == "simple":
            generated_ids = model.generate(**inputs, max_length=25, num_beams=3)
        elif style == "detailed":
            generated_ids = model.generate(**inputs, max_length=max_length, num_beams=4, do_sample=True, temperature=0.7)
        else:  # artistic
            generated_ids = model.generate(**inputs, max_length=max_length, num_beams=5, do_sample=True, temperature=0.9)
        
        caption = processor.decode(generated_ids[0], skip_special_tokens=True)
        return enhance_caption(caption, style)
        
    except Exception as e:
        return f"BLIP Error: {str(e)}"

def enhance_caption(caption, style):
    """Clean up and enhance captions"""
    # Remove common prefixes
    caption = re.sub(r'^(a|an|the|this is|there is)\s+', '', caption, flags=re.IGNORECASE)
    caption = re.sub(r'(image|picture|photo)\s+(of|showing)\s+', '', caption, flags=re.IGNORECASE)
    
    # Style enhancements
    if style == "simple":
        return caption.strip()
    elif style == "detailed":
        enhancers = ["detailed", "high-quality", "clear", "professional"]
        import random
        return f"{random.choice(enhancers)} {caption.strip()}"
    else:  # artistic
        enhancers = ["stunning", "beautiful", "artistic", "masterpiece", "breathtaking"]
        import random
        return f"{random.choice(enhancers)} {caption.strip()}"

def generate_lightweight_prompt(image, style="detailed", max_length=50):
    """Main function - always uses BLIP model if loaded, else fallback"""
    if not MODEL_LOADED:
        return generate_basic_fallback(image, style)
    try:
        return generate_caption_blip(image, style, max_length)
    except Exception as e:
        return generate_basic_fallback(image, style)

def generate_basic_fallback(image, style):
    """Ultra-basic fallback (no models needed)"""
    import numpy as np
    
    try:
        img_array = np.array(image)
        height, width = img_array.shape[:2]
        
        # Basic analysis
        aspect_ratio = width / height
        if aspect_ratio > 1.5:
            comp = "wide landscape"
        elif aspect_ratio < 0.7:
            comp = "portrait orientation"
        else:
            comp = "balanced composition"
        
        # Color analysis
        if len(img_array.shape) == 3:
            avg_brightness = np.mean(img_array)
            if avg_brightness > 180:
                lighting = "bright, well-lit"
            elif avg_brightness > 100:
                lighting = "evenly lit"
            else:
                lighting = "moody, darker"
        else:
            lighting = "black and white"
        
        base_descriptions = {
            "simple": f"{comp} image",
            "detailed": f"A {lighting} {comp} photograph with good clarity",
            "artistic": f"Artistic {comp} with {lighting} atmospheric mood"
        }
        
        return base_descriptions.get(style, "processed image")
        
    except Exception as e:
        return f"Basic analysis: {style} style image"

def create_lightweight_interface():
    """Lightweight Gradio interface (no model info or delete instructions)"""
    
    with gr.Blocks(title="Lightweight Image-to-Prompt", theme=gr.themes.Monochrome()) as demo:
        gr.Markdown("# ⚡ Lightweight Image-to-Prompt")
        gr.Markdown("*Optimized for speed and storage efficiency*")
        
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(
                    type="pil",
                    label="📷 Upload Image",
                    height=300
                )
                
                with gr.Row():
                    style_radio = gr.Radio(
                        choices=["simple", "detailed", "artistic"],
                        value="detailed",
                        label="Style",
                        scale=2
                    )
                    
                    length_slider = gr.Slider(
                        minimum=20,
                        maximum=80,
                        value=50,
                        step=10,
                        label="Length",
                        scale=1
                    )
                
                generate_btn = gr.Button("Generate", variant="primary")
                
            with gr.Column():
                prompt_output = gr.Textbox(
                    label="Generated Prompt",
                    lines=5,
                    placeholder="Upload an image..."
                )
                
                with gr.Row():
                    copy_btn = gr.Button("Copy", size="sm")
                    clear_btn = gr.Button("Clear", size="sm")
        
        # Event handlers
        def process_image(image, style, length):
            if image is None:
                return "Upload an image first!"
            return generate_lightweight_prompt(image, style, length)
        
        # Auto-generate on upload
        image_input.change(
            fn=process_image,
            inputs=[image_input, style_radio, length_slider],
            outputs=prompt_output
        )
        
        generate_btn.click(
            fn=process_image,
            inputs=[image_input, style_radio, length_slider],
            outputs=prompt_output
        )
        
        copy_btn.click(None, inputs=prompt_output, js="(text) => navigator.clipboard.writeText(text)")
        clear_btn.click(lambda: "", outputs=prompt_output)
    
    return demo

if __name__ == "__main__":
    print("🚀 Starting Lightweight Image-to-Prompt...")
    print(f"💾 Model cache: ~/.cache/huggingface/transformers/")
    
    demo = create_lightweight_interface()
    demo.launch()