"""
Voice-controlled image generator (Gradio) for deployment on Hugging Face Spaces.
"""

import os
import random
from typing import Optional
from contextlib import contextmanager

import torch
from PIL import Image
import gradio as gr
from transformers import pipeline
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

HF_TOKEN = os.environ.get("HF_TOKEN")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

ASR_MODEL = "openai/whisper-small"
SD_MODEL = "runwayml/stable-diffusion-v1-5"

asr = None
pipe = None

def load_asr():
    global asr
    if asr is None:
        device_index = 0 if DEVICE == "cuda" else -1
        asr = pipeline("automatic-speech-recognition", model=ASR_MODEL, device=device_index, chunk_length_s=30)
    return asr

def load_sd():
    global pipe
    if pipe is None:
        pipe = StableDiffusionPipeline.from_pretrained(
            SD_MODEL,
            use_auth_token=HF_TOKEN if HF_TOKEN else None,
            safety_checker=None,
        )
        try:
            pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
        except Exception:
            pass
        pipe = pipe.to(DEVICE)
        try:
            pipe.enable_xformers_memory_efficient_attention()
        except Exception:
            pass
    return pipe

@contextmanager
def nullcontext():
    yield

# Gradio component (return numpy array + sample rate)
mic = gr.Audio(type="numpy", label="Record or upload audio")

def transcribe_audio(audio_input):
    if audio_input is None:
        return ""
    # audio_input = (numpy array, sample_rate)
    audio_array, sample_rate = audio_input
    asr_model = load_asr()
    try:
        result = asr_model({"array": audio_array, "sampling_rate": sample_rate})
        return result.get("text", "").strip() if isinstance(result, dict) else str(result).strip()
    except Exception as e:
        return f"Error: {str(e)}"

def generate_image(prompt: str, negative_prompt: str, steps: int, guidance: float, seed: int):
    if not prompt.strip():
        return None, "Please provide a prompt."
    pipe = load_sd()
    if seed is None or int(seed) < 0:
        seed = random.randint(1, 2**32 - 1)
    gen = torch.Generator(device=DEVICE).manual_seed(int(seed)) if DEVICE == "cuda" else torch.Generator().manual_seed(int(seed))
    autocast_ctx = torch.autocast(device_type="cuda") if DEVICE == "cuda" else nullcontext()
    with autocast_ctx:
        out = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt or None,
            num_inference_steps=int(steps),
            guidance_scale=float(guidance),
            generator=gen,
        )
        image = out.images[0]
    return image, f"Seed: {seed} | Steps: {steps} | Guidance: {guidance}"

# --- Gradio UI ---
with gr.Blocks(title="Voice-controlled Image Generator") as demo:
    gr.Markdown("## 🎙️ Voice-controlled Image Generator\nRecord or upload a voice prompt → transcribe → edit → generate images.")
    with gr.Row():
        with gr.Column(scale=2):
            mic = gr.Audio(type="filepath", label="Record or upload audio")  # ✅ fixed
            transcribe_btn = gr.Button("Transcribe")
            prompt_box = gr.Textbox(label="Prompt (editable)", placeholder="Speak or type your prompt here", lines=3)
            negative = gr.Textbox(label="Negative prompt (optional)", placeholder="Things to avoid", lines=2)
            with gr.Row():
                steps = gr.Slider(minimum=10, maximum=60, value=28, step=1, label="Sampling steps")
                guidance = gr.Slider(minimum=1.0, maximum=25.0, value=7.5, step=0.5, label="Guidance scale (CFG)")
            with gr.Row():
                seed_in = gr.Number(label="Seed (leave -1 for random)", value=-1)
                gen_btn = gr.Button("Generate Image")
            status = gr.Textbox(label="Status / Info", interactive=False)
        with gr.Column(scale=3):
            image_out = gr.Image(label="Generated image", type="pil")
            gallery = gr.gallery = gr.Gallery(
    label="Generated Captions",
    show_label=True,
    elem_id="gallery",
    columns=[1],  # replaces grid=[1]
    height="auto"
)


    def on_transcribe(audio_data):
    transcript = transcribe_audio(audio_data)
    return transcript, "Transcription complete." if transcript else "Couldn't transcribe."

    def on_generate(prompt_text, negative_prompt, steps_val, guidance_val, seed_val):
        seed_int = int(seed_val) if seed_val is not None else -1
        img, info = generate_image(prompt_text, negative_prompt, steps_val, guidance_val, seed_int)
        return img, [img] if img else [], info

    transcribe_btn.click(fn=on_transcribe, inputs=[mic], outputs=[prompt_box, status])
    gen_btn.click(fn=on_generate, inputs=[prompt_box, negative, steps, guidance, seed_in], outputs=[image_out, gallery, status])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))