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 cv2
import translators as ts

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

# Constants
MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
MAX_SEED = np.iinfo(np.int32).max

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

# Helper function to return a progress bar HTML snippet.
def progress_bar_html(label: str) -> str:
    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: #F0FFF0; border-radius: 2px; overflow: hidden;">
        <div style="width: 100%; height: 100%; background-color: #00FF00 ; animation: loading 1.5s linear infinite;"></div>
    </div>
</div>
<style>
@keyframes loading {{
    0% {{ transform: translateX(-100%); }}
    100% {{ transform: translateX(100%); }}
}}
</style>
    '''

# TEXT MODEL - Utiliser Napoleon 4B au lieu de FastThink
model_id = "baconnier/Napoleon_4B_V0.0" 
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)
model.eval()

# MULTIMODAL (OCR) MODELS - Garder Qwen2-VL pour OCR
MODEL_ID_VL = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct" 
processor = AutoProcessor.from_pretrained(MODEL_ID_VL, trust_remote_code=True)
model_m = Qwen2VLForConditionalGeneration.from_pretrained(
    MODEL_ID_VL,
    trust_remote_code=True,
    torch_dtype=torch.float16
).to("cuda").eval()

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

bad_words = json.loads(os.getenv('BAD_WORDS', "[]"))
bad_words_negative = json.loads(os.getenv('BAD_WORDS_NEGATIVE', "[]"))
default_negative = os.getenv("default_negative", "")

def check_text(prompt, negative=""):
    for i in bad_words:
        if i in prompt:
            return True
    for i in bad_words_negative:
        if i in negative:
            return True
    return False

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

CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"

dtype = torch.float16 if device.type == "cuda" else torch.float32

# NAPOLEON 4B MULTIMODAL MODEL - Remplacer Gemma3 par Napoleon
napoleon_model_id = "baconnier/Napoleon_4B_V0.0"
napoleon_model = AutoModelForCausalLM.from_pretrained(
    napoleon_model_id, device_map="auto", torch_dtype=torch.bfloat16
).eval()
napoleon_processor = AutoProcessor.from_pretrained(napoleon_model_id)

# Fonction de traduction
def translate_text(text, target_lang="fr", source_lang="auto"):
    try:
        return ts.deepl(text, from_language=source_lang, to_language=target_lang)
    except:
        try:
            return ts.google(text, from_language=source_lang, to_language=target_lang)
        except:
            return text  # Retourner le texte original en cas d'échec

# VIDEO PROCESSING HELPER
def downsample_video(video_path):
    vidcap = cv2.VideoCapture(video_path)
    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = vidcap.get(cv2.CAP_PROP_FPS)
    frames = []
    # Sample 10 evenly spaced frames.
    frame_indices = np.linspace(0, total_frames - 1, 10, dtype=int)
    for i in frame_indices:
        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
        success, image = vidcap.read()
        if success:
            # Convert from BGR to RGB and then to PIL Image.
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            pil_image = Image.fromarray(image)
            timestamp = round(i / fps, 2)
            frames.append((pil_image, timestamp))
    vidcap.release()
    return frames

# MAIN GENERATION FUNCTION
@spaces.GPU
def generate(
    input_dict: dict,
    chat_history: list[dict],
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
    temperature: float = 0.6,
    top_p: float = 0.9,
    top_k: int = 50,
    repetition_penalty: float = 1.2,
):
    text = input_dict["text"]
    files = input_dict.get("files", [])

    lower_text = text.lower().strip()

    # NAPOLEON 4B TEXT & MULTIMODAL (image) Branch
    if lower_text.startswith("@napoleon"):
        # Remove the napoleon flag from the prompt.
        prompt_clean = re.sub(r"@napoleon", "", text, flags=re.IGNORECASE).strip().strip('"')
        
        # Traduire en français si le texte n'est pas déjà en français
        prompt_clean_fr = translate_text(prompt_clean, target_lang="fr")
        
        if files:
            # If image files are provided, load them.
            images = [load_image(f) for f in files]
            messages = [{
                "role": "user",
                "content": [
                    *[{"type": "image", "image": image} for image in images],
                    {"type": "text", "text": prompt_clean_fr},
                ]
            }]
        else:
            messages = [
                {"role": "system", "content": [{"type": "text", "text": "Vous êtes un assistant utile qui parle français."}]},
                {"role": "user", "content": [{"type": "text", "text": prompt_clean_fr}]}
            ]
        
        inputs = napoleon_processor.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=True,
            return_dict=True, return_tensors="pt"
        ).to(napoleon_model.device, dtype=torch.bfloat16)
        
        streamer = TextIteratorStreamer(
            napoleon_processor.tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True
        )
        
        generation_kwargs = {
            **inputs,
            "streamer": streamer,
            "max_new_tokens": max_new_tokens,
            "do_sample": True,
            "temperature": temperature,
            "top_p": top_p,
            "top_k": top_k,
            "repetition_penalty": repetition_penalty,
        }
        
        thread = Thread(target=napoleon_model.generate, kwargs=generation_kwargs)
        thread.start()
        
        buffer = ""
        yield progress_bar_html("Traitement avec Napoleon 4B")
        for new_text in streamer:
            buffer += new_text
            time.sleep(0.01)
            yield buffer
        return

    # NAPOLEON 4B VIDEO Branch
    if lower_text.startswith("@video"):
        # Remove the video flag from the prompt.
        prompt_clean = re.sub(r"@video", "", text, flags=re.IGNORECASE).strip().strip('"')
        
        # Traduire en français si le texte n'est pas déjà en français
        prompt_clean_fr = translate_text(prompt_clean, target_lang="fr")
        
        if files:
            # Assume the first file is a video.
            video_path = files[0]
            frames = downsample_video(video_path)
            
            messages = [
                {"role": "system", "content": [{"type": "text", "text": "Vous êtes un assistant utile qui parle français."}]},
                {"role": "user", "content": [{"type": "text", "text": prompt_clean_fr}]}
            ]
            
            # Append each frame as an image with a timestamp label.
            for frame in frames:
                image, timestamp = frame
                image_path = f"video_frame_{uuid.uuid4().hex}.png"
                image.save(image_path)
                messages[1]["content"].append({"type": "text", "text": f"Image à {timestamp}s:"})
                messages[1]["content"].append({"type": "image", "url": image_path})
        else:
            messages = [
                {"role": "system", "content": [{"type": "text", "text": "Vous êtes un assistant utile qui parle français."}]},
                {"role": "user", "content": [{"type": "text", "text": prompt_clean_fr}]}
            ]
            
        inputs = napoleon_processor.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=True,
            return_dict=True, return_tensors="pt"
        ).to(napoleon_model.device, dtype=torch.bfloat16)
        
        streamer = TextIteratorStreamer(
            napoleon_processor.tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True
        )
        
        generation_kwargs = {
            **inputs,
            "streamer": streamer,
            "max_new_tokens": max_new_tokens,
            "do_sample": True,
            "temperature": temperature,
            "top_p": top_p,
            "top_k": top_k,
            "repetition_penalty": repetition_penalty,
        }
        
        thread = Thread(target=napoleon_model.generate, kwargs=generation_kwargs)
        thread.start()
        
        buffer = ""
        yield progress_bar_html("Traitement vidéo avec Napoleon 4B")
        for new_text in streamer:
            buffer += new_text
            time.sleep(0.01)
            yield buffer
        return

    # Otherwise, handle text/chat generation.
    conversation = clean_chat_history(chat_history)
    conversation.append({"role": "user", "content": text})
    
    if files:
        images = [load_image(image) for image in files] if len(files) > 1 else [load_image(files[0])]
        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 = ""
        yield progress_bar_html("Traitement avec Qwen2VL OCR")
        for new_text in streamer:
            buffer += new_text
            buffer = buffer.replace("<|im_end|>", "")
            time.sleep(0.01)
            yield buffer
    else:
        # Traduire le texte en français pour Napoleon
        text_fr = translate_text(text, target_lang="fr")
        conversation_fr = clean_chat_history(chat_history)
        conversation_fr.append({"role": "user", "content": text_fr})
        
        input_ids = tokenizer.apply_chat_template(conversation_fr, add_generation_prompt=True, return_tensors="pt")
        if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
            input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
            gr.Warning(f"Texte d'entrée tronqué car plus long que {MAX_INPUT_TOKEN_LENGTH} tokens.")
        
        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)
            yield "".join(outputs)

        final_response = "".join(outputs)
        yield final_response

demo = gr.ChatInterface(
    fn=generate,
    additional_inputs=[
        gr.Slider(label="Nombre maximum de tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS),
        gr.Slider(label="Température", minimum=0.1, maximum=4.0, step=0.1, value=0.6),
        gr.Slider(label="Top-p (échantillonnage nucleus)", 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="Pénalité de répétition", minimum=1.0, maximum=2.0, step=0.05, value=1.2),
    ],
    examples=[
        [
            {
                "text": "@napoleon Créez une histoire courte basée sur les images.",
                "files": [
                    "examples/1111.jpg",
                    "examples/2222.jpg",
                    "examples/3333.jpg",
                ],
            }
        ],
        [{"text": "@napoleon Expliquez cette image", "files": ["examples/3.jpg"]}],
        [{"text": "@video Expliquez le contenu de cette publicité", "files": ["examples/videoplayback.mp4"]}],
        [{"text": "@napoleon Quel personnage de film est-ce?", "files": ["examples/9999.jpg"]}],
        ["@napoleon Expliquez la température critique d'une substance"],
        [{"text": "@napoleon Transcription de cette lettre", "files": ["examples/222.png"]}],
        [{"text": "@video Expliquez le contenu de la vidéo en détail", "files": ["examples/breakfast.mp4"]}],
        [{"text": "@video Décrivez la vidéo", "files": ["examples/Missing.mp4"]}],
        [{"text": "@video Expliquez ce qui se passe dans cette vidéo", "files": ["examples/oreo.mp4"]}],
        [{"text": "@video Résumez les événements de cette vidéo", "files": ["examples/sky.mp4"]}],
        [{"text": "@video Qu'y a-t-il dans cette vidéo?", "files": ["examples/redlight.mp4"]}],
        ["Programme Python pour la rotation de tableau"],
        ["@napoleon Expliquez la température critique d'une substance"]
    ],
    cache_examples=False,
    type="messages",
    description="# **Napoleon 4B `@napoleon pour le multimodal, @video pour la compréhension vidéo`**",
    fill_height=True,
    textbox=gr.MultimodalTextbox(label="Saisir votre question", file_types=["image", "video"], file_count="multiple", placeholder="Utilisez @napoleon pour le multimodal, @video pour l'analyse vidéo !"),
    stop_btn="Arrêter la génération",
    multimodal=True,
)

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