import os

import gradio as gr
import spaces
import torch

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel


# ============================================================
# CONFIGURATION
# ============================================================

BASE_MODEL_ID = "unsloth/gemma-4-E4B-it-unsloth-bnb-4bit"
ADAPTER_ID = "arbib/gemma4-e4b-chamaliya-30k-l4-adapter_v2"

HF_TOKEN = os.environ.get("HF_TOKEN")


# ============================================================
# CHARGEMENT DU TOKENIZER
# ============================================================

print("Chargement du tokenizer...")

tokenizer = AutoTokenizer.from_pretrained(
    BASE_MODEL_ID,
    token=HF_TOKEN,
    trust_remote_code=True,
)

if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id


# ============================================================
# CHARGEMENT DU MODÈLE 4-BIT
# ============================================================

print("Chargement du modèle Gemma-4 E4B...")

model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_ID,
    token=HF_TOKEN,
    trust_remote_code=True,
    device_map="cuda",
    torch_dtype=torch.float16,
)

print("Chargement de l’adaptateur Chamaliya...")

model = PeftModel.from_pretrained(
    model,
    ADAPTER_ID,
    token=HF_TOKEN,
    torch_device="cpu",
    low_cpu_mem_usage=False,
)

model.eval()

print("Modèle Chamaliya chargé avec succès.")


# ============================================================
# NETTOYAGE
# ============================================================

def nettoyer_reponse(texte):
    coupures = [
        "<turn|>",
        "<|turn>",
        "<end_of_turn>",
        "shift_to_english",
        "\nModel\n",
        "\nmodel\n",
        "User:",
        "Human:",
        "Assistant:",
    ]

    for coupure in coupures:
        position = texte.find(coupure)

        if position != -1:
            texte = texte[:position]

    return texte.strip()


# ============================================================
# INFÉRENCE ZEROGPU
# ============================================================

@spaces.GPU(duration=120)
def generer_reponse(message, historique):
    if not message or not message.strip():
        return "كتب ليا شي سؤال باش نجاوبك."

    consigne = (
        "جاوبني مباشرة غير بالدارجة الشمالية ديال تطوان وطنجة "
        "وماتستعملش كلمات د الداخل. "
        "استعمل كلمات فحال جبرت، موطع، مزيون، دقم، لكانيف، فحال، "
        "وتبع قاعدة ديالوم، عندوم، وقتوم.\n"
        "السؤال ديالي هو: "
    )

    messages = []

    if historique:
        for element in historique[-6:]:
            if not isinstance(element, dict):
                continue

            role = element.get("role")
            content = element.get("content")

            if role in ["user", "assistant"] and isinstance(content, str):
                messages.append({
                    "role": role,
                    "content": content,
                })

    messages.append({
        "role": "user",
        "content": consigne + message,
    })

    texte = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    inputs = tokenizer(
        texte,
        return_tensors="pt",
    ).to("cuda")

    longueur_entree = inputs["input_ids"].shape[-1]

    with torch.inference_mode():
        outputs = model.generate(
            **inputs,
            max_new_tokens=220,
            temperature=0.15,
            top_p=0.80,
            repetition_penalty=1.25,
            no_repeat_ngram_size=3,
            do_sample=True,
            use_cache=True,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.pad_token_id,
        )

    nouveaux_tokens = outputs[0][longueur_entree:]

    reponse = tokenizer.decode(
        nouveaux_tokens,
        skip_special_tokens=True,
    )

    return nettoyer_reponse(reponse)


# ============================================================
# INTERFACE
# ============================================================

demo = gr.ChatInterface(
    fn=generer_reponse,
    title="Darija Chamaliya",
    description="مساعد ذكي كيهضر بالدارجة الشمالية ديال تطوان وطنجة.",
    textbox=gr.Textbox(
        placeholder="كتب السؤال ديالك هنا...",
        rtl=True,
    ),
    examples=[
        "شني خبارك؟",
        "شرح ليا الذكاء الاصطناعي",
        "عطيني نصيحة باش نتعلم البرمجة",
    ],
)

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