import os
import random
import uuid
import json
import time
import asyncio
import re
from threading import Thread

import gradio as gr
import spaces
import torch
import numpy as np
from PIL import Image
import edge_tts

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TextIteratorStreamer,
    Qwen2VLForConditionalGeneration,
    AutoProcessor,
)
from transformers.image_utils import load_image
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler

DESCRIPTION = """
# Gen Vision 🎃
Separate Tabs for Chat, Image Generation (LoRA), Qwen2 VL OCR and Text-to-Speech
"""

css = '''
h1 {
  text-align: center;
  display: block;
}

#duplicate-button {
  margin: auto;
  color: #fff;
  background: #1565c0;
  border-radius: 100vh;
}
'''

MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))

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

# -----------------------
# Progress Bar Helper
# -----------------------
def progress_bar_html(label: str) -> str:
    """
    Returns an HTML snippet for a thin progress bar with a label.
    The progress bar is styled as a dark red animated bar.
    """
    return f'''
<div style="display: flex; align-items: center;">
    <span style="margin-right: 10px; font-size: 14px;">{label}</span>
    <div style="width: 110px; height: 5px; background-color: #DDA0DD; border-radius: 2px; overflow: hidden;">
        <div style="width: 100%; height: 100%; background-color: #FF00FF; animation: loading 1.5s linear infinite;"></div>
    </div>
</div>
<style>
@keyframes loading {{
    0% {{ transform: translateX(-100%); }}
    100% {{ transform: translateX(100%); }}
}}
</style>
    '''

# -----------------------
# Text Generation Setup (Chat)
# -----------------------
model_id = "prithivMLmods/FastThink-0.5B-Tiny"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)
model.eval()

# -----------------------
# TTS Setup
# -----------------------
TTS_VOICES = [
    "en-US-JennyNeural",
    "en-US-GuyNeural",
]

async def text_to_speech(text: str, voice: str, output_file="output.mp3"):
    """Convert text to speech using Edge TTS and save as MP3"""
    communicate = edge_tts.Communicate(text, voice)
    await communicate.save(output_file)
    return output_file

# -----------------------
# Utility: Clean Chat History
# -----------------------
def clean_chat_history(chat_history):
    """
    Filter out any chat entries whose "content" is not a string.
    """
    cleaned = []
    for msg in chat_history:
        if isinstance(msg, dict) and isinstance(msg.get("content"), str):
            cleaned.append(msg)
    return cleaned

# -----------------------
# Qwen2 VL OCR Setup
# -----------------------
OCR_MODEL_ID = "prithivMLmods/Qwen2-VL-OCR2-2B-Instruct" 
processor = AutoProcessor.from_pretrained(OCR_MODEL_ID, trust_remote_code=True)
model_m = Qwen2VLForConditionalGeneration.from_pretrained(
    OCR_MODEL_ID,
    trust_remote_code=True,
    torch_dtype=torch.float16
).to("cuda").eval()

# -----------------------
# Stable Diffusion Image Generation Setup (LoRA)
# -----------------------
MAX_SEED = np.iinfo(np.int32).max
USE_TORCH_COMPILE = False
ENABLE_CPU_OFFLOAD = False

if torch.cuda.is_available():
    pipe = StableDiffusionXLPipeline.from_pretrained(
        "SG161222/RealVisXL_V4.0_Lightning",
        torch_dtype=torch.float16,
        use_safetensors=True,
    )
    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
    
    # LoRA options with one example for each.
    LORA_OPTIONS = {
        "Realism": ("prithivMLmods/Canopus-Realism-LoRA", "Canopus-Realism-LoRA.safetensors", "rlms"),
        "Pixar": ("prithivMLmods/Canopus-Pixar-Art", "Canopus-Pixar-Art.safetensors", "pixar"),
        "Photoshoot": ("prithivMLmods/Canopus-Photo-Shoot-Mini-LoRA", "Canopus-Photo-Shoot-Mini-LoRA.safetensors", "photo"),
        "Clothing": ("prithivMLmods/Canopus-Clothing-Adp-LoRA", "Canopus-Dress-Clothing-LoRA.safetensors", "clth"),
        "Interior": ("prithivMLmods/Canopus-Interior-Architecture-0.1", "Canopus-Interior-Architecture-0.1δ.safetensors", "arch"),
        "Fashion": ("prithivMLmods/Canopus-Fashion-Product-Dilation", "Canopus-Fashion-Product-Dilation.safetensors", "fashion"),
        "Minimalistic": ("prithivMLmods/Pegasi-Minimalist-Image-Style", "Pegasi-Minimalist-Image-Style.safetensors", "minimalist"),
        "Modern": ("prithivMLmods/Canopus-Modern-Clothing-Design", "Canopus-Modern-Clothing-Design.safetensors", "mdrnclth"),
        "Animaliea": ("prithivMLmods/Canopus-Animaliea-Artism", "Canopus-Animaliea-Artism.safetensors", "Animaliea"),
        "Wallpaper": ("prithivMLmods/Canopus-Liquid-Wallpaper-Art", "Canopus-Liquid-Wallpaper-Minimalize-LoRA.safetensors", "liquid"),
        "Cars": ("prithivMLmods/Canes-Cars-Model-LoRA", "Canes-Cars-Model-LoRA.safetensors", "car"),
        "PencilArt": ("prithivMLmods/Canopus-Pencil-Art-LoRA", "Canopus-Pencil-Art-LoRA.safetensors", "Pencil Art"),
        "ArtMinimalistic": ("prithivMLmods/Canopus-Art-Medium-LoRA", "Canopus-Art-Medium-LoRA.safetensors", "mdm"),
    }

    # Load all LoRA weights
    for model_name, weight_name, adapter_name in LORA_OPTIONS.values():
        pipe.load_lora_weights(model_name, weight_name=weight_name, adapter_name=adapter_name)
    pipe.to("cuda")
else:
    pipe = StableDiffusionXLPipeline.from_pretrained(
        "SG161222/RealVisXL_V4.0_Lightning",
        torch_dtype=torch.float32,
        use_safetensors=True,
    ).to(device)

def save_image(img: Image.Image) -> str:
    """Save a PIL image with a unique filename and return the path."""
    unique_name = str(uuid.uuid4()) + ".png"
    img.save(unique_name)
    return unique_name

def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed

@spaces.GPU(duration=180, enable_queue=True)
def generate_image(prompt: str, negative_prompt: str, seed: int, width: int, height: int, guidance_scale: float, randomize_seed: bool, lora_model: str):
    seed = int(randomize_seed_fn(seed, randomize_seed))
    effective_negative_prompt = negative_prompt  # Use provided negative prompt if any
    model_name, weight_name, adapter_name = LORA_OPTIONS[lora_model]
    pipe.set_adapters(adapter_name)
    outputs = pipe(
         prompt=prompt,
         negative_prompt=effective_negative_prompt,
         width=width,
         height=height,
         guidance_scale=guidance_scale,
         num_inference_steps=28,
         num_images_per_prompt=1,
         cross_attention_kwargs={"scale": 0.65},
         output_type="pil",
    )
    images = outputs.images
    image_paths = [save_image(img) for img in images]
    return image_paths, seed

# -----------------------
# Chat Generation Function (Text-only)
# -----------------------
def generate_chat(input_text: str, chat_history: list, max_new_tokens: int, temperature: float, top_p: float, top_k: int, repetition_penalty: float):
    conversation = clean_chat_history(chat_history)
    conversation.append({"role": "user", "content": input_text})
    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
    input_ids = input_ids.to(model.device)
    streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
    generation_kwargs = {
        "input_ids": input_ids,
        "streamer": streamer,
        "max_new_tokens": max_new_tokens,
        "do_sample": True,
        "top_p": top_p,
        "top_k": top_k,
        "temperature": temperature,
        "num_beams": 1,
        "repetition_penalty": repetition_penalty,
    }
    t = Thread(target=model.generate, kwargs=generation_kwargs)
    t.start()
    outputs = []
    for new_text in streamer:
        outputs.append(new_text)
    final_response = "".join(outputs)
    chat_history.append({"role": "assistant", "content": final_response})
    return chat_history

# -----------------------
# Qwen2 VL OCR Function (Multimodal)
# -----------------------
def generate_ocr(text: str, files, max_new_tokens: int):
    if files:
        if isinstance(files, list) and len(files) > 1:
            images = [load_image(image) for image in files]
        elif isinstance(files, list) and len(files) == 1:
            images = [load_image(files[0])]
        else:
            images = [load_image(files)]
        messages = [{
            "role": "user",
            "content": [*([{"type": "image", "image": image} for image in images]),
                        {"type": "text", "text": text}]
        }]
        prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = processor(text=[prompt], images=images, return_tensors="pt", padding=True).to("cuda")
        streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
        generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens}
        thread = Thread(target=model_m.generate, kwargs=generation_kwargs)
        thread.start()
        buffer = ""
        for new_text in streamer:
            buffer += new_text
        return buffer
    else:
        return "No images provided."

# -----------------------
# Text-to-Speech Function
# -----------------------
def generate_tts(text: str, voice: str):
    output_file = asyncio.run(text_to_speech(text, voice))
    return output_file

# -----------------------
# Gradio Interface with Tabs
# -----------------------
with gr.Blocks(css=css, title="Gen Vision") as demo:
    gr.Markdown(DESCRIPTION)
    
    with gr.Tab("Chat Interface"):
        with gr.Row():
            chat_history = gr.Chatbot(label="Chat History")
        with gr.Row():
            chat_input = gr.Textbox(placeholder="Enter your message", label="Your Message")
        with gr.Row():
            max_new_tokens_slider = gr.Slider(label="Max New Tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS)
            temperature_slider = gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6)
        with gr.Row():
            top_p_slider = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9)
            top_k_slider = gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50)
            repetition_penalty_slider = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2)
        send_btn = gr.Button("Send")
        send_btn.click(
            fn=generate_chat,
            inputs=[chat_input, chat_history, max_new_tokens_slider, temperature_slider, top_p_slider, top_k_slider, repetition_penalty_slider],
            outputs=chat_history,
        )
    
    with gr.Tab("Image Generation"):
        image_prompt = gr.Textbox(label="Prompt", placeholder="Enter image prompt")
        negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Enter negative prompt")
        seed_input = gr.Number(label="Seed", value=0)
        width_slider = gr.Slider(label="Width", minimum=256, maximum=2048, step=64, value=1024)
        height_slider = gr.Slider(label="Height", minimum=256, maximum=2048, step=64, value=1024)
        guidance_scale_slider = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=3.0)
        randomize_checkbox = gr.Checkbox(label="Randomize Seed", value=True)
        lora_dropdown = gr.Dropdown(label="LoRA Style", choices=list(LORA_OPTIONS.keys()), value="Realism")
        generate_img_btn = gr.Button("Generate Image")
        img_output = gr.Image(label="Generated Image")
        seed_output = gr.Number(label="Used Seed")
        generate_img_btn.click(
            fn=generate_image,
            inputs=[image_prompt, negative_prompt, seed_input, width_slider, height_slider, guidance_scale_slider, randomize_checkbox, lora_dropdown],
            outputs=[img_output, seed_output],
        )
    
    with gr.Tab("Qwen 2 VL OCR"):
        ocr_text = gr.Textbox(label="Text Prompt", placeholder="Enter prompt for OCR")
        file_input = gr.File(label="Upload Images", file_count="multiple")
        ocr_max_new_tokens = gr.Slider(label="Max New Tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS)
        ocr_btn = gr.Button("Run OCR")
        ocr_output = gr.Textbox(label="OCR Output")
        ocr_btn.click(
            fn=generate_ocr,
            inputs=[ocr_text, file_input, ocr_max_new_tokens],
            outputs=ocr_output,
        )
    
    with gr.Tab("Text-to-Speech"):
        tts_text = gr.Textbox(label="Text", placeholder="Enter text for TTS")
        voice_dropdown = gr.Dropdown(label="Voice", choices=TTS_VOICES, value=TTS_VOICES[0])
        tts_btn = gr.Button("Generate Audio")
        tts_audio = gr.Audio(label="Audio Output", type="filepath")
        tts_btn.click(
            fn=generate_tts,
            inputs=[tts_text, voice_dropdown],
            outputs=tts_audio,
        )

demo.queue(max_size=20).launch(share=True)