import os
import sys
import cv2
import time
import copy
import random
import torch
import spaces
import requests
import subprocess
import gradio as gr
from PIL import Image
import importlib.util
from threading import Thread
from typing import Iterable, Optional, Tuple, List

def check_and_install_package(package_name, import_name=None, pip_name=None):
    """Check if a package is installed, and if not, install it."""
    if import_name is None:
        import_name = package_name
    if pip_name is None:
        pip_name = package_name

    spec = importlib.util.find_spec(import_name)
    if spec is None:
        print(f"Installing {package_name}...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", pip_name, "-q"])
        print(f"✓ {package_name} installed successfully")
    return True

print("Checking and installing transformers==4.57.3 ...")
check_and_install_package("transformers", "transformers", "transformers==4.57.3")
print("Done!")

from transformers import (
    Qwen2_5_VLForConditionalGeneration,
    AutoModelForImageTextToText,
    AutoModelForCausalLM,
    AutoProcessor,
    TextIteratorStreamer,
)

try:
    from vibevoice.modular.modeling_vibevoice_streaming_inference import (
        VibeVoiceStreamingForConditionalGenerationInference,
    )
    from vibevoice.processor.vibevoice_streaming_processor import (
        VibeVoiceStreamingProcessor,
    )
except ImportError:
    print("CRITICAL WARNING: 'vibevoice' modules not found. Ensure the vibevoice repository structure is present.")
    VibeVoiceStreamingForConditionalGenerationInference = None
    VibeVoiceStreamingProcessor = None

from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes

colors.orange_red = colors.Color(
    name="orange_red",
    c50="#FFF0E5",
    c100="#FFE0CC",
    c200="#FFC299",
    c300="#FFA366",
    c400="#FF8533",
    c500="#FF4500",
    c600="#E63E00",
    c700="#CC3700",
    c800="#B33000",
    c900="#992900",
    c950="#802200",
)

class OrangeRedTheme(Soft):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.gray,
        secondary_hue: colors.Color | str = colors.orange_red,
        neutral_hue: colors.Color | str = colors.slate,
        text_size: sizes.Size | str = sizes.text_lg,
        font: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
        ),
        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
        ),
    ):
        super().__init__(
            primary_hue=primary_hue,
            secondary_hue=secondary_hue,
            neutral_hue=neutral_hue,
            text_size=text_size,
            font=font,
            font_mono=font_mono,
        )
        super().set(
            background_fill_primary="*primary_50",
            background_fill_primary_dark="*primary_900",
            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
            button_primary_text_color="white",
            button_primary_text_color_hover="white",
            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
            button_secondary_text_color="black",
            button_secondary_text_color_hover="white",
            button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)",
            button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)",
            button_secondary_background_fill_dark="linear-gradient(90deg, *primary_500, *primary_600)",
            button_secondary_background_fill_hover_dark="linear-gradient(90deg, *primary_500, *primary_500)",
            slider_color="*secondary_500",
            slider_color_dark="*secondary_600",
            block_title_text_weight="600",
            block_border_width="3px",
            block_shadow="*shadow_drop_lg",
            button_primary_shadow="*shadow_drop_lg",
            button_large_padding="11px",
            color_accent_soft="*primary_100",
            block_label_background_fill="*primary_200",
        )

orange_red_theme = OrangeRedTheme()

css = """
#main-title h1 {
    font-size: 2.3em !important;
}
#output-title h2 {
    font-size: 2.1em !important;
}
.generating {
    border: 2px solid #4682B4;
}
"""

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Using Main Device: {device}")

QWEN_VL_MODEL_ID = "Qwen/Qwen2.5-VL-7B-Instruct"
print(f"Loading OCR Model: {QWEN_VL_MODEL_ID}...")

qwen_processor = AutoProcessor.from_pretrained(QWEN_VL_MODEL_ID, trust_remote_code=True)
qwen_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    QWEN_VL_MODEL_ID,
    #attn_implementation="flash_attention_2",
    trust_remote_code=True,
    torch_dtype=torch.float16
).to(device).eval()

print("Model loaded successfully.")

TTS_MODEL_PATH = "microsoft/VibeVoice-Realtime-0.5B"
print(f"Loading TTS Model: {TTS_MODEL_PATH}...")
print("VibeVoice Model loaded successfully.")

tts_processor = VibeVoiceStreamingProcessor.from_pretrained(TTS_MODEL_PATH)

tts_model = VibeVoiceStreamingForConditionalGenerationInference.from_pretrained(
    TTS_MODEL_PATH,
    torch_dtype=torch.float16,
    device_map="cuda",
    attn_implementation="sdpa",
)
tts_model.eval()
tts_model.set_ddpm_inference_steps(num_steps=5)

class VoiceMapper:
    """Maps speaker names to voice file paths"""
    def __init__(self):
        self.setup_voice_presets()
        new_dict = {}
        for name, path in self.voice_presets.items():
            if "_" in name: name = name.split("_")[0]
            if "-" in name: name = name.split("-")[-1]
            new_dict[name] = path
        self.voice_presets.update(new_dict)

    def setup_voice_presets(self):
        voices_dir = os.path.join(os.path.dirname(__file__), "demo/voices/streaming_model")
        if not os.path.exists(voices_dir):
            print(f"Warning: Voices directory not found at {voices_dir}")
            self.voice_presets = {}
            self.available_voices = {}
            return

        self.voice_presets = {}
        pt_files = [f for f in os.listdir(voices_dir) if f.lower().endswith(".pt") and os.path.isfile(os.path.join(voices_dir, f))]

        for pt_file in pt_files:
            name = os.path.splitext(pt_file)[0]
            full_path = os.path.join(voices_dir, pt_file)
            self.voice_presets[name] = full_path

        self.voice_presets = dict(sorted(self.voice_presets.items()))
        self.available_voices = {name: path for name, path in self.voice_presets.items() if os.path.exists(path)}
        print(f"Found {len(self.available_voices)} voice files.")

    def get_voice_path(self, speaker_name: str) -> str:
        if speaker_name in self.voice_presets:
            return self.voice_presets[speaker_name]
        speaker_lower = speaker_name.lower()
        for preset_name, path in self.voice_presets.items():
            if preset_name.lower() in speaker_lower or speaker_lower in preset_name.lower():
                return path
        if self.voice_presets:
            return list(self.voice_presets.values())[0]
        return ""

VOICE_MAPPER = VoiceMapper()
print("TTS Model loaded successfully.")

@spaces.GPU
def process_pipeline(
    image: Image.Image,
    query: str,
    speaker_name: str,
    cfg_scale: float,
    ocr_max_tokens: int,
    ocr_temp: float,
    progress=gr.Progress()
):
    """
    Combined pipeline: Image - Text -> TTS - Audio
    """
    if image is None:
        return "Please upload an image.", None, "Error: No image provided."
    
    progress(0.2, desc="Analyzing Image ()...")
    
    if not query: 
        query = "Analyze the content perfectly."
        
    messages = [{
        "role": "user",
        "content": [
            {"type": "image"},
            {"type": "text", "text": query},
        ]
    }]
    
    prompt_full = qwen_processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    
    inputs = qwen_processor(
        text=[prompt_full],
        images=[image],
        return_tensors="pt",
        padding=True
    ).to(device)
    
    generated_ids = qwen_model.generate(
        **inputs,
        max_new_tokens=ocr_max_tokens,
        do_sample=True,
        temperature=ocr_temp,
        top_p=0.9,
    )
    
    generated_ids_trimmed = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    extracted_text = qwen_processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0]
    
    extracted_text = extracted_text.replace("<|im_end|>", "").strip()
    
    progress(0.5, desc=f"Analyzing completed. Converting to speech ({len(extracted_text)} chars)...")
    
    if not extracted_text:
        return extracted_text, None, "produced no text."
        
    try:
        full_script = extracted_text.replace("'", "'").replace('"', '"').replace('"', '"')
        
        voice_path = VOICE_MAPPER.get_voice_path(speaker_name)
        if not voice_path:
             return extracted_text, None, "Error: Voice file not found."

        all_prefilled_outputs = torch.load(voice_path, map_location="cuda", weights_only=False)
        
        tts_inputs = tts_processor.process_input_with_cached_prompt(
            text=full_script,
            cached_prompt=all_prefilled_outputs,
            padding=True,
            return_tensors="pt",
            return_attention_mask=True,
        )
        
        tts_model.to("cuda")
        for k, v in tts_inputs.items():
            if torch.is_tensor(v):
                tts_inputs[k] = v.to("cuda")
                
        with torch.cuda.amp.autocast():
            outputs = tts_model.generate(
                **tts_inputs,
                max_new_tokens=None,
                cfg_scale=cfg_scale,
                tokenizer=tts_processor.tokenizer,
                generation_config={"do_sample": False},
                verbose=False,
                all_prefilled_outputs=copy.deepcopy(all_prefilled_outputs)
            )
            
        tts_model.to("cpu")
        torch.cuda.empty_cache()
        
        if outputs.speech_outputs and outputs.speech_outputs[0] is not None:
            sample_rate = 24000
            
            output_dir = "./outputs"
            os.makedirs(output_dir, exist_ok=True)
            output_path = os.path.join(output_dir, f"generated_{int(time.time())}.wav")
            
            tts_processor.save_audio(
                outputs.speech_outputs[0].cpu(),
                output_path=output_path,
            )
            
            status = f"✅ Success! Text Length: {len(extracted_text)} chars."
            return extracted_text, output_path, status
        else:
            return extracted_text, None, "TTS Generation failed (no output)."

    except Exception as e:
        tts_model.to("cpu")
        torch.cuda.empty_cache()
        import traceback
        return extracted_text, None, f"Error during TTS: {str(e)}"

url = "https://huggingface.co/datasets/strangervisionhf/image-examples/resolve/main/2.jpg?download=true"
example_image = Image.open(requests.get(url, stream=True).raw).convert("RGB")

with gr.Blocks() as demo:
    gr.Markdown("# **Vision-to-VibeVoice-en**", elem_id="main-title")
    gr.Markdown("Perform vision-to-audio inference with [Qwen2.5VL](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct) + [VibeVoice-Realtime-0.5B](https://huggingface.co/microsoft/VibeVoice-Realtime-0.5B).")
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 1. Vision Input")
            image_upload = gr.Image(type="pil", label="Upload Image", value=example_image, height=300)
            image_query = gr.Textbox(label="Enter the prompt", value="Give a short description indicating whether the image is safe or unsafe.", placeholder="E.g., Read this page...")
            
            gr.Markdown("### 2. Voice Settings")
            voice_choices = list(VOICE_MAPPER.available_voices.keys())
            if not voice_choices: voice_choices = ["Default"]
                
            speaker_dropdown = gr.Dropdown(
                 choices=voice_choices,
                value=voice_choices[0],
                label="Speaker Voice"
            )
                
            cfg_slider = gr.Slider(minimum=1.0, maximum=3.0, value=1.5, step=0.1, label="CFG Scale (Speech Fidelity)")
                
            with gr.Accordion("Advanced Options", open=False):
                max_new_tokens = gr.Slider(label="Max Tokens", minimum=128, maximum=4096, step=128, value=1024)
                temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, step=0.1, value=0.1)

            submit_btn = gr.Button("Generate Speech", variant="primary", size="lg")

        with gr.Column(scale=1):
            gr.Markdown("### 3. Results", elem_id="output-title")
            
            text_output = gr.Textbox(
                label="Extracted Text (Editable)", 
                interactive=True, 
                lines=11, 
            )
            
            audio_output = gr.Audio(
                label="Generated Speech",
                type="filepath",
                interactive=False
            )
            
            status_output = gr.Textbox(label="Status Log", lines=2)
            
            gr.Examples(
                examples=[["Caption the image...", "https://huggingface.co/datasets/merve/vlm_test_images/resolve/main/venice.jpg"]],
                inputs=[image_query, image_upload],
                label="Example"
            )

    submit_btn.click(
        fn=process_pipeline,
        inputs=[
            image_upload, 
            image_query, 
            speaker_dropdown, 
            cfg_slider,
            max_new_tokens,
            temperature
        ],
        outputs=[text_output, audio_output, status_output]
    )

if __name__ == "__main__":
    demo.queue(max_size=40).launch(css=css, theme=orange_red_theme, mcp_server=True, ssr_mode=False, show_error=True)