import gradio as gr
import torch
import torch.nn as nn
import sentencepiece as spm
import math
import os
from groq import Groq

# ── config ──────────────────────────────────────────────────────────────────
VOCAB_SIZE = 4000
D_MODEL    = 256
NHEAD      = 4
NUM_LAYERS = 3
DIM_FF     = 512
DROPOUT    = 0.1
MAX_LEN    = 128
PAD_ID, BOS_ID, EOS_ID = 0, 2, 3
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# ── model definition ─────────────────────────────────────────────────────────
class PE(nn.Module):
    def __init__(self, d, mx=512):
        super().__init__()
        pe  = torch.zeros(mx, d)
        pos = torch.arange(0, mx).float().unsqueeze(1)
        div = torch.exp(torch.arange(0, d, 2).float() * (-math.log(10000.0) / d))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe.unsqueeze(0))

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

class Transltr(nn.Module):
    def __init__(self):
        super().__init__()
        self.se   = nn.Embedding(VOCAB_SIZE, D_MODEL, padding_idx=PAD_ID)
        self.te   = nn.Embedding(VOCAB_SIZE, D_MODEL, padding_idx=PAD_ID)
        self.pe   = PE(D_MODEL)
        self.tf   = nn.Transformer(D_MODEL, NHEAD, NUM_LAYERS, NUM_LAYERS,
                                   DIM_FF, DROPOUT, batch_first=True)
        self.proj = nn.Linear(D_MODEL, VOCAB_SIZE)

    def forward(self, src, tgt, tgt_mask=None, src_key_padding_mask=None,
                tgt_key_padding_mask=None, memory_key_padding_mask=None):
        tmsk = self.tf.generate_square_subsequent_mask(tgt.size(1)).to(DEVICE)
        sm   = (src == PAD_ID)
        tm   = (tgt == PAD_ID)
        se   = self.pe(self.se(src) * math.sqrt(D_MODEL))
        te   = self.pe(self.te(tgt) * math.sqrt(D_MODEL))
        return self.proj(self.tf(se, te, tgt_mask=tmsk,
                                 src_key_padding_mask=sm,
                                 tgt_key_padding_mask=tm,
                                 memory_key_padding_mask=sm))

# ── load model & tokenizers ──────────────────────────────────────────────────
src_sp = spm.SentencePieceProcessor()
src_sp.load("eng_tok.model")
tgt_sp = spm.SentencePieceProcessor()
tgt_sp.load("tiv_tok.model")

model = Transltr().to(DEVICE)
ckpt  = torch.load("transltr_model.pt", map_location=DEVICE)
model.load_state_dict(ckpt["model_state"])
model.eval()

groq_client = Groq(api_key=os.environ["GROQ_API_KEY"])

# ── inference ────────────────────────────────────────────────────────────────
def translate(text: str) -> str:
    ids = [BOS_ID] + src_sp.encode(text) + [EOS_ID]
    src = torch.tensor([ids], device=DEVICE)
    out = [BOS_ID]
    with torch.no_grad():
        for _ in range(MAX_LEN):
            tgt     = torch.tensor([out], device=DEVICE)
            logits  = model(src, tgt)
            next_id = logits[0, -1].argmax().item()
            if next_id == EOS_ID:
                break
            out.append(next_id)
    return tgt_sp.decode(out[1:])

def speech_to_tiv(audio_path):
    if audio_path is None:
        return "", ""
    with open(audio_path, "rb") as f:
        transcription = groq_client.audio.transcriptions.create(
            file=("audio.wav", f),
            model="whisper-large-v3",
            language="en",
        )
    english = transcription.text.strip()
    tiv     = translate(english)
    return english, tiv

def text_to_tiv(text):
    if not text.strip():
        return ""
    return translate(text.strip())

# ── UI ───────────────────────────────────────────────────────────────────────
css = """
@import url('https://fonts.googleapis.com/css2?family=Space+Mono:wght@400;700&family=Syne:wght@400;700;800&display=swap');

body { background: #0a0a0a !important; }

.gradio-container {
    max-width: 820px !important;
    margin: 0 auto !important;
    font-family: 'Syne', sans-serif !important;
    background: #0a0a0a !important;
}

#title {
    text-align: center;
    padding: 2.5rem 0 1rem;
    color: #f0ebe0;
}

#title h1 {
    font-size: 3rem;
    font-weight: 800;
    letter-spacing: -2px;
    margin: 0;
    background: linear-gradient(135deg, #f0ebe0, #c8a96e);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}

#title p {
    color: #888;
    font-family: 'Space Mono', monospace;
    font-size: 0.75rem;
    letter-spacing: 3px;
    text-transform: uppercase;
    margin-top: 0.4rem;
}

.tab-nav button {
    font-family: 'Space Mono', monospace !important;
    font-size: 0.7rem !important;
    letter-spacing: 2px !important;
    text-transform: uppercase !important;
    color: #888 !important;
    background: transparent !important;
    border: none !important;
    border-bottom: 2px solid transparent !important;
    padding: 0.75rem 1.5rem !important;
}

.tab-nav button.selected {
    color: #c8a96e !important;
    border-bottom: 2px solid #c8a96e !important;
}

.output-box {
    background: #141414;
    border: 1px solid #2a2a2a;
    border-radius: 8px;
    padding: 1.25rem;
    font-family: 'Space Mono', monospace;
    font-size: 1.1rem;
    color: #f0ebe0;
    min-height: 80px;
    line-height: 1.8;
}

.label-text {
    font-family: 'Space Mono', monospace;
    font-size: 0.65rem;
    letter-spacing: 3px;
    text-transform: uppercase;
    color: #555;
    margin-bottom: 0.5rem;
}

footer { display: none !important; }
"""

with gr.Blocks(css=css, title="Transltr — English to Tiv") as demo:
    gr.HTML("""
    <div id="title">
        <h1>transltr</h1>
        <p>English &rarr; Tiv &nbsp;·&nbsp; Gospel-trained neural translation</p>
    </div>
    """)

    with gr.Tabs():
        with gr.Tab("🎙 speak"):
            gr.Markdown("Record your voice in English and get the Tiv translation instantly.")
            audio_in = gr.Audio(sources=["microphone"], type="filepath", label="Speak in English")
            btn_voice = gr.Button("Translate ↗", variant="primary")
            with gr.Row():
                with gr.Column():
                    gr.HTML('<div class="label-text">English (transcribed)</div>')
                    eng_out = gr.Textbox(show_label=False, interactive=False,
                                         placeholder="Your speech will appear here…",
                                         elem_classes=["output-box"])
                with gr.Column():
                    gr.HTML('<div class="label-text">Tiv translation</div>')
                    tiv_out1 = gr.Textbox(show_label=False, interactive=False,
                                          placeholder="Tiv translation will appear here…",
                                          elem_classes=["output-box"])
            btn_voice.click(speech_to_tiv, inputs=audio_in, outputs=[eng_out, tiv_out1])

        with gr.Tab("⌨ type"):
            gr.Markdown("Type English text and translate to Tiv.")
            text_in = gr.Textbox(placeholder="Type English here…", label="English", lines=3)
            btn_text = gr.Button("Translate ↗", variant="primary")
            gr.HTML('<div class="label-text">Tiv translation</div>')
            tiv_out2 = gr.Textbox(show_label=False, interactive=False,
                                   placeholder="Tiv translation will appear here…",
                                   elem_classes=["output-box"])
            btn_text.click(text_to_tiv, inputs=text_in, outputs=tiv_out2)
            text_in.submit(text_to_tiv, inputs=text_in, outputs=tiv_out2)

    gr.HTML("""
    <div style="text-align:center;margin-top:2rem;font-family:'Space Mono',monospace;
                font-size:0.6rem;letter-spacing:2px;color:#333;text-transform:uppercase;">
        Trained on the Tiv Gospels &nbsp;·&nbsp; Matt · Mark · Luke · John
    </div>
    """)

demo.launch()