import spaces
import re
import time

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


REPOS = {
    "translate": ("Sandroeth/cali-id-en-translate", "472c3e80063867ef775fb41d058b9ffec13f44bd"),
    "base": ("Sandroeth/cali-0.1B", "c9de368e7e62184e2a70e76f4b3cedcfa170b73a"),
}
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MODELS = {}
for key, (repo, revision) in REPOS.items():
    tokenizer = AutoTokenizer.from_pretrained(repo, revision=revision, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        repo, revision=revision, trust_remote_code=True, dtype=torch.float32,
    ).to(DEVICE).eval()
    MODELS[key] = (tokenizer, model)
    print(f"Loaded {repo}@{revision} on {DEVICE}, dtype={model.dtype}", flush=True)


def run_model(key: str, prompt: str, max_new_tokens: int, **options) -> tuple[str, str]:
    tokenizer, model = MODELS[key]
    inputs = tokenizer(prompt, return_tensors="pt", return_token_type_ids=False).to(DEVICE)
    input_length = inputs["input_ids"].shape[1]
    if input_length + max_new_tokens > model.config.seq_len:
        raise gr.Error("Text is too long. Shorten the input or reduce the output limit (1024 tokens total).")
    start = time.perf_counter()
    with torch.inference_mode():
        output = model.generate(
            **inputs, max_new_tokens=max_new_tokens, use_cache=False,
            pad_token_id=tokenizer.eos_token_id, **options,
        )
    generated = output[0][input_length:]
    text = tokenizer.decode(generated, skip_special_tokens=True).strip()
    elapsed = time.perf_counter() - start
    print(f"{key}: input={input_length}, output={len(generated)}, seconds={elapsed:.2f}", flush=True)
    return text, f"{input_length} input tokens · {len(generated)} output tokens · {elapsed:.2f} seconds"


@spaces.GPU(duration=60)
def translate(text: str, direction: str = "Indonesian → English", max_new_tokens: int = 64,
              first_sentence: bool = True) -> tuple[str, str]:
    """Translate a short sentence between Indonesian and English using CALI's trained prompt format."""
    if not text or not text.strip():
        raise gr.Error("Enter the text you want to translate.")
    if direction not in ("Indonesian → English", "English → Indonesian"):
        raise gr.Error("Invalid translation direction.")
    max_new_tokens = int(max_new_tokens)
    if not 16 <= max_new_tokens <= 256:
        raise gr.Error("Choose between 16 and 256 output tokens.")
    if len(text.split()) > 200:
        raise gr.Error("Use up to 200 words; this model is designed for short texts.")
    tag = "[id→en]" if direction == "Indonesian → English" else "[en→id]"
    result, stats = run_model("translate", f"{tag}\n{text.strip()}\n→", max_new_tokens,
                              do_sample=False, repetition_penalty=1.5, no_repeat_ngram_size=3)
    if first_sentence:
        result = re.split(r"(?<=[.!?])\s+", result)[0]
    return result, stats


@spaces.GPU(duration=60)
def complete(prompt: str, max_new_tokens: int = 96, temperature: float = 0.7,
             top_p: float = 0.9) -> tuple[str, str]:
    """Continue an Indonesian or English text prompt with the experimental CALI base language model."""
    if not prompt or not prompt.strip():
        raise gr.Error("Enter a text prompt to continue.")
    max_new_tokens = int(max_new_tokens)
    if not 16 <= max_new_tokens <= 256 or not 0 <= temperature <= 1.5 or not 0.1 <= top_p <= 1:
        raise gr.Error("Generation settings are outside the supported range.")
    options = {"do_sample": temperature > 0, "repetition_penalty": 1.1}
    if temperature > 0:
        options.update(temperature=float(temperature), top_p=float(top_p))
    return run_model("base", prompt, max_new_tokens, **options)


with gr.Blocks(title="cali") as demo:
    gr.Markdown("# cali\nTwo lightweight models by **Sandroeth** for exploring Indonesian and English.")
    with gr.Tab("Translation"):
        gr.Markdown("### From one language to another\nBest for everyday sentences. Very short inputs (1–2 words), technical terms, and long passages may produce less accurate translations.")
        direction = gr.Radio(["Indonesian → English", "English → Indonesian"], value="Indonesian → English", label="Translation direction")
        with gr.Row():
            source = gr.Textbox(label="Source text", lines=5, placeholder="Dia pergi ke pasar setiap pagi.", max_length=6000)
            translated = gr.Textbox(label="Translation", lines=5, interactive=False)
        with gr.Accordion("Settings", open=False):
            tokens = gr.Slider(16, 256, value=64, step=1, label="Maximum output tokens")
            first = gr.Checkbox(value=True, label="Return the first sentence only (matches the official example)")
        translate_button = gr.Button("Translate", variant="primary")
        translation_stats = gr.Textbox(label="Stats", interactive=False)
        translate_button.click(translate, [source, direction, tokens, first], [translated, translation_stats], api_name="translate")
        gr.Examples([
            ["Dia pergi ke pasar setiap pagi.", "Indonesian → English"],
            ["Saya makan nasi.", "Indonesian → English"],
            ["The weather is very cold today.", "English → Indonesian"],
            ["She is happy.", "English → Indonesian"],
        ], inputs=[source, direction], outputs=[translated, translation_stats], fn=translate,
            cache_examples=True, cache_mode="lazy")
    with gr.Tab("Explore CALI"):
        gr.Markdown("### Continue a piece of text\nCALI-0.1B is an experimental base model. Start a sentence or paragraph and let it continue. Outputs may repeat or become inconsistent.")
        prompt = gr.Textbox(label="Prompt", lines=4, placeholder="Indonesia adalah negara yang")
        continuation = gr.Textbox(label="Generated continuation", lines=6, interactive=False)
        with gr.Accordion("Settings", open=False):
            base_tokens = gr.Slider(16, 256, value=96, step=1, label="Maximum output tokens")
            temperature = gr.Slider(0, 1.5, value=0.7, step=0.05, label="Temperature (0 = deterministic)")
            top_p = gr.Slider(0.1, 1, value=0.9, step=0.05, label="Top-p")
        generate_button = gr.Button("Generate continuation", variant="primary")
        base_stats = gr.Textbox(label="Stats", interactive=False)
        generate_button.click(complete, [prompt, base_tokens, temperature, top_p], [continuation, base_stats], api_name="complete")
        gr.Examples([["Indonesia adalah negara yang"], ["The sun rises in the"]],
                    inputs=[prompt], outputs=[continuation, base_stats], fn=complete,
                    cache_examples=True, cache_mode="lazy")
    gr.Markdown("Model: [CALI Translate](https://huggingface.co/Sandroeth/cali-id-en-translate) · [CALI-0.1B](https://huggingface.co/Sandroeth/cali-0.1B) · Apache 2.0 license")

if __name__ == "__main__":
    demo.queue(max_size=20, default_concurrency_limit=1).launch(mcp_server=True)