import spaces
import os
import tempfile
import gradio as gr
from dotenv import load_dotenv
import torch
from scipy.io.wavfile import write
from diffusers import DiffusionPipeline
from transformers import pipeline
from pathlib import Path
from PIL import Image  # <-- ADDED THIS IMPORT
import io               # <-- ADDED THIS IMPORT

load_dotenv()
hf_token = os.getenv("HF_TKN")

device_id = 0 if torch.cuda.is_available() else -1

# Correctly initialize the modern, reliable captioning pipeline
captioning_pipeline = pipeline(
    "image-to-text",
    model="Salesforce/blip-image-captioning-large", 
    device=device_id
)

# Initialize the audio pipeline
pipe = DiffusionPipeline.from_pretrained(
    "cvssp/audioldm2",
    use_auth_token=hf_token
)


# === THIS IS THE CORRECTED FUNCTION ===
@spaces.GPU(duration=120)
def analyze_image_with_free_model(image_file_bytes):
    try:
        # No more temp files!
        # Open the image data directly from memory using Pillow
        image = Image.open(io.BytesIO(image_file_bytes))

        # Pass the Pillow Image object directly to the pipeline. This is the robust method.
        results = captioning_pipeline(image)
        
        if not results or not isinstance(results, list):
            return "Error: Could not generate caption.", True
        
        caption = results[0].get("generated_text", "").strip()
        if not caption:
            return "No caption was generated.", True
        return caption, False

    except Exception as e:
        print(f"ERROR in analyze_image_with_free_model: {e}") # Print error to logs
        return f"Error analyzing image: {e}", True


@spaces.GPU(duration=120)
def get_audioldm_from_caption(caption):
    try:
        pipe.to("cuda")
        audio_output = pipe(
            prompt=caption,
            num_inference_steps=50,
            guidance_scale=7.5
        )
        pipe.to("cpu")
        audio = audio_output.audios[0]

        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_wav:
            write(temp_wav.name, 16000, audio)
            return temp_wav.name

    except Exception as e:
        print(f"Error generating audio from caption: {e}")
        return None

# --- Gradio UI (No changes needed here) ---
css = """
#col-container{
    margin: 0 auto;
    max-width: 800px;
    }
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML("""
    <h1 style="text-align: center;">🎶 Generate Sound Effects from Image</h1>
    <p style="text-align: center;">
        ⚡ Powered by <a href="https://bilsimaging.com" target="_blank">Bilsimaging</a>
    </p>
        """)

    gr.Markdown("""
    Welcome to this unique sound effect generator! This tool allows you to upload an image and generate a 
    descriptive caption and a corresponding sound effect, all using free, open-source models on Hugging Face.
    
    **💡 How it works:**
    1. **Upload an image**: Choose an image that you'd like to analyze.
    2. **Generate Description**: Click on 'Generate Description' to get a textual description of your uploaded image.
    3. **Generate Sound Effect**: Based on the image description, click on 'Generate Sound Effect' to create a 
       sound effect that matches the image context.
    
    Enjoy the journey from visual to auditory sensation with just a few clicks!
    """)

    image_upload = gr.File(label="Upload Image", type="binary")
    generate_description_button = gr.Button("Generate Description")
    caption_display = gr.Textbox(label="Image Description", interactive=False)
    generate_sound_button = gr.Button("Generate Sound Effect")
    audio_output = gr.Audio(label="Generated Sound Effect")

    gr.Markdown("""
    ## 👥 How You Can Contribute
    We welcome contributions and suggestions for improvements. Your feedback is invaluable 
    to the continuous enhancement of this application. 
    
    For support, questions, or to contribute, please contact us at 
    [contact@bilsimaging.com](mailto:contact@bilsimaging.com).
    
    Support our work and get involved by donating through 
    [Ko-fi](https://ko-fi.com/bilsimaging). - Bilel Aroua
    """)

    gr.Markdown("""
    ## 📢 Stay Connected
    This app is a testament to the creative possibilities that emerge when technology meets art. 
    Enjoy exploring the auditory landscape of your images!
    """)
    
    # --- Gradio event handlers (I've updated the function called here) ---
    def update_caption(image_file_bytes):
        # We pass the bytes from the uploader directly to our corrected function
        description, _ = analyze_image_with_free_model(image_file_bytes)
        return description

    def generate_sound(description):
        if not description or description.startswith("Error"):
            return None
        audio_path = get_audioldm_from_caption(description)
        return audio_path

    generate_description_button.click(
        fn=update_caption,
        inputs=image_upload,
        outputs=caption_display
    )

    generate_sound_button.click(
        fn=generate_sound,
        inputs=caption_display,
        outputs=audio_output
    )
    
    gr.HTML('<a href="https://visitorbadge.io/status?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2FBils%2FGenerate-Sound-Effects-from-Image"><img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2FBils%2FGenerate-Sound-Effects-from-Image&countColor=%23263759" /></a>')
    html = gr.HTML()

demo.launch(debug=True, share=True)