import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from peft import PeftModel

BASE_MODEL  = "facebook/nllb-200-distilled-1.3B"
LORA_REPO   = "kawkumputer/pulaar-ai-nllb-1.3b-v9"
DEVICE      = "cpu"

LANG_FR  = "fra_Latn"
LANG_AR  = "arb_Arab"
LANG_PUL = "fuv_Latn"

DIRECTION_MAP = {
    "fr → pul": (LANG_FR,  LANG_PUL),
    "ar → pul": (LANG_AR,  LANG_PUL),
    "pul → fr": (LANG_PUL, LANG_FR),
    "pul → ar": (LANG_PUL, LANG_AR),
}

print("Chargement du tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

print("Chargement du modèle de base...")
base_model = AutoModelForSeq2SeqLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True,
)

print(f"Application des adaptateurs LoRA depuis {LORA_REPO}...")
model = PeftModel.from_pretrained(base_model, LORA_REPO)
model.eval()
print("Modèle prêt.")


def translate(text: str, direction: str) -> str:
    if not text or not text.strip():
        return ""
    src_lang, tgt_lang = DIRECTION_MAP[direction]
    tokenizer.src_lang = src_lang
    inputs = tokenizer(
        text.strip(),
        return_tensors="pt",
        truncation=True,
        max_length=256,
    )
    forced_bos = tokenizer.convert_tokens_to_ids(tgt_lang)
    with torch.no_grad():
        out = model.generate(
            **inputs,
            forced_bos_token_id=forced_bos,
            num_beams=4,
            max_length=256,
            repetition_penalty=1.3,
            no_repeat_ngram_size=3,
            early_stopping=True,
        )
    return tokenizer.decode(out[0], skip_special_tokens=True)


with gr.Blocks(title="PulaarAI — Traducteur Pulaar") as demo:

    gr.Markdown(
        "# PulaarAI — Traducteur Pulaar\n"
        "**Traduction fr/ar ↔ Pulaar** · Dialecte Fouta-Toro · "
        "Développé par Hamath Kane, Abou Sy & Bocar Amadou Ba (ARPRIM)\n\n"
        "Basé sur **NLLB-200-distilled-1.3B** + LoRA fine-tuné · "
        "⚠️ Inférence CPU — comptez 20-40 secondes par traduction"
    )

    gr.Markdown(
        "> **Modèle actif : v9** — NLLB-1.3B · dataset quotidien + ARPRIM (juridique/santé/histoire) · "
        "chrF++ Coran fr→pul : **37.67** · chrF++ Quotidien fr→pul : **42.99** (+4.45 vs v7)"
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_text = gr.Textbox(
                label="Texte source",
                placeholder="Saisissez le texte à traduire...",
                lines=5,
            )
            direction = gr.Dropdown(
                choices=list(DIRECTION_MAP.keys()),
                value="fr → pul",
                label="Direction",
            )
            with gr.Row():
                btn_clear = gr.Button("Effacer", variant="secondary")
                btn       = gr.Button("Traduire", variant="primary", scale=2)

        with gr.Column(scale=1):
            output_text = gr.Textbox(
                label="Traduction",
                lines=5,
                interactive=False,
            )

    btn.click(fn=translate, inputs=[input_text, direction], outputs=output_text)
    input_text.submit(fn=translate, inputs=[input_text, direction], outputs=output_text)
    btn_clear.click(fn=lambda: ("", ""), outputs=[input_text, output_text])

demo.launch(theme=gr.themes.Soft())