import os
from threading import Thread
import gradio as gr
import spaces
import torch
import edge_tts
import asyncio
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from transformers.image_utils import load_image
from huggingface_hub import InferenceClient
import time

# Load text-only model and tokenizer
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()

# Load multimodal (OCR) model and processor
MODEL_ID = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct" 
processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model_m = Qwen2VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    torch_dtype=torch.float16
).to("cuda").eval()

TTS_VOICES = [
    "en-US-JennyNeural",  # @tts1
    "en-US-GuyNeural",    # @tts2
]

def image_gen(prompt):
    """Generate image using API"""
    try:
        client = InferenceClient("prithivMLmods/STABLE-HAMSTER")
        return client.text_to_image(prompt)
    except:
        client_flux = InferenceClient("black-forest-labs/FLUX.1-schnell")
        return client_flux.text_to_image(prompt)

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

def clean_chat_history(chat_history):
    return [msg for msg in chat_history if isinstance(msg, dict) and isinstance(msg.get("content"), str)]

@spaces.GPU
def generate(input_dict: dict, chat_history: list[dict], max_new_tokens=1024, temperature=0.6, top_p=0.9, top_k=50, repetition_penalty=1.2):
    """Generates chatbot responses with multimodal input, TTS, and image generation."""
    text = input_dict["text"]
    files = input_dict.get("files", [])
    images = [load_image(file) for file in files] if files else []
    
    if text.startswith("@tts"):
        voice_index = next((i for i in range(1, 3) if text.startswith(f"@tts{i}")), None)
        if voice_index:
            voice = TTS_VOICES[voice_index - 1]
            text = text.replace(f"@tts{voice_index}", "").strip()
            conversation = [{"role": "user", "content": text}]
        else:
            voice = None
    elif text.startswith("@image"):
        query = text.replace("@image", "").strip()
        yield "Generating Image, Please wait..."
        image = image_gen(query)
        yield gr.Image(image)
    else:
        conversation = clean_chat_history(chat_history) + [{"role": "user", "content": text}]
        if images:
            messages = [{
                "role": "user",
                "content": [
                    *[{"type": "image", "image": img} for img 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)
            thread = Thread(target=model_m.generate, kwargs={**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens})
            thread.start()
            buffer = ""
            for new_text in streamer:
                buffer += new_text.replace("<|im_end|>", "")
                yield buffer
        else:
            input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(model.device)
            streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
            thread = Thread(target=model.generate, 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,
            })
            thread.start()
            response = "".join([new_text for new_text in streamer])
            yield response
            if voice:
                output_file = asyncio.run(text_to_speech(response, voice))
                yield gr.Audio(output_file, autoplay=True)

demo = gr.ChatInterface(
    fn=generate,
    additional_inputs=[
        gr.Slider(label="Max new tokens", minimum=1, maximum=2048, step=1, value=1024),
        gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6),
        gr.Slider(label="Top-p", minimum=0.05, maximum=1.0, step=0.05, value=0.9),
        gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50),
        gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2),
    ],
    examples=[
        ["@tts1 Who is Nikola Tesla?"],
        [{"text": "Extract JSON from the image", "files": ["examples/document.jpg"]}],
        ["@image futuristic city at sunset"],
        ["A train travels 60 kilometers per hour. How far will it travel in 5 hours?"],
    ],
    cache_examples=False,
    description="# QwQ Edge 💬",
    fill_height=True,
    textbox=gr.MultimodalTextbox(label="Query Input", file_types=["image"], file_count="multiple"),
    stop_btn="Stop Generation",
    multimodal=True,
)

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