import os
import torch
import gradio as gr
import sentencepiece as spm
from huggingface_hub import hf_hub_download
from modeling_flash import NordicFlash

REPO = "NodeNestor/bifrost-flash-430m"          # private — pulled with the HF_TOKEN space secret
TOKEN = os.environ.get("HF_TOKEN")

print("Downloading Bifrost Flash assets…", flush=True)
weights = hf_hub_download(REPO, "model.safetensors", token=TOKEN)
spm_path = hf_hub_download(REPO, "nordic_unigram_65k.model", token=TOKEN)
sp = spm.SentencePieceProcessor(); sp.load(spm_path)
model = NordicFlash.from_checkpoint(weights, device="cpu", dtype=torch.float32)
print("Model ready.", flush=True)

TARGETS = {
    "Swedish": "sv", "Danish": "da", "Norwegian Bokmål": "nb",
    "Norwegian Nynorsk": "nn", "Finnish": "fi", "Icelandic": "is", "English": "en",
}
LID = {"en": 65000, "sv": 65001, "da": 65002, "nb": 65003, "nn": 65004, "fi": 65005, "is": 65006}


def translate(text, target):
    text = (text or "").strip()
    if not text:
        return ""
    ids = sp.encode(text, out_type=int)[:1000]
    out = model.translate(ids, LID[TARGETS[target]], max_new=128)
    return sp.decode(out)


with gr.Blocks(title="Bifrost — Nordic translation") as demo:
    gr.Markdown(
        "# 🌉 Bifrost Flash 430M\n"
        "Open **Nordic ↔ English** translation — Swedish · Danish · Norwegian (Bokmål/Nynorsk) · "
        "Finnish · Icelandic ↔ English. Source language is auto-detected; pick a target.\n\n"
        "*The 430M distilled model, running on CPU. Part of [NodeNestor](https://nodenestor.com).*"
    )
    with gr.Row():
        inp = gr.Textbox(label="Text", lines=4, placeholder="Hello, how are you today?")
        out = gr.Textbox(label="Translation", lines=4)
    target = gr.Dropdown(list(TARGETS), value="Swedish", label="Translate into")
    btn = gr.Button("Translate", variant="primary")
    btn.click(translate, [inp, target], out)
    inp.submit(translate, [inp, target], out)
    gr.Examples(
        [["Hello, how are you today?", "Swedish"],
         ["The weather is nice and the sun is shining.", "Danish"],
         ["Vädret är fint idag.", "English"],
         ["God morgen, min venn.", "Finnish"]],
        [inp, target],
    )

demo.queue(max_size=16).launch(server_name="0.0.0.0", server_port=7860)