from __future__ import annotations

import os
import threading

os.environ.setdefault("OPENBLAS_NUM_THREADS", "4")
os.environ.setdefault("OMP_NUM_THREADS", "4")
os.environ.setdefault("MKL_NUM_THREADS", "4")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("GRADIO_SSR_MODE", "false")

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

MODEL_ID = os.environ.get("TINY_AYA_MODEL", "CohereLabs/tiny-aya-global")
DEFAULT_MAX_TOKENS = int(os.environ.get("TINY_AYA_MAX_TOKENS", "400"))

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
if torch.cuda.is_available():
    model.to("cuda")
model.eval()

_lock = threading.Lock()


def _messages(system: str, user: str):
    messages = []
    if system and system.strip():
        messages.append({"role": "system", "content": system.strip()})
    messages.append({"role": "user", "content": (user or "").strip()})
    return messages


@spaces.GPU(duration=120)
def generate(system: str, user: str, max_tokens: int = DEFAULT_MAX_TOKENS, temperature: float = 0.8):
    if not user or not user.strip():
        raise gr.Error("user prompt required")
    max_tokens = max(1, min(int(max_tokens or DEFAULT_MAX_TOKENS), 1024))
    temperature = max(0.0, min(float(temperature if temperature is not None else 0.8), 2.0))
    inputs = tokenizer.apply_chat_template(
        _messages(system, user),
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)
    with _lock, torch.inference_mode():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            do_sample=temperature > 0,
            temperature=max(temperature, 1e-5),
            top_p=0.95,
            pad_token_id=tokenizer.eos_token_id,
        )
    return tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True).strip()


@spaces.GPU(duration=120)
def generate_stream(system: str, user: str, max_tokens: int = DEFAULT_MAX_TOKENS, temperature: float = 0.8):
    if not user or not user.strip():
        raise gr.Error("user prompt required")
    max_tokens = max(1, min(int(max_tokens or DEFAULT_MAX_TOKENS), 1024))
    temperature = max(0.0, min(float(temperature if temperature is not None else 0.8), 2.0))
    inputs = tokenizer.apply_chat_template(
        _messages(system, user),
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

    def run():
        with _lock, torch.inference_mode():
            model.generate(
                **inputs,
                max_new_tokens=max_tokens,
                do_sample=temperature > 0,
                temperature=max(temperature, 1e-5),
                top_p=0.95,
                pad_token_id=tokenizer.eos_token_id,
                streamer=streamer,
            )

    thread = threading.Thread(target=run, daemon=True)
    thread.start()
    acc = ""
    for token in streamer:
        acc += token
        yield acc
    thread.join(timeout=1)


with gr.Blocks(title="Tiny Army Tiny Aya ZeroGPU") as demo:
    gr.Markdown("# Tiny Army Tiny Aya ZeroGPU")
    system = gr.Textbox(label="System", lines=5)
    user = gr.Textbox(label="User", lines=5)
    max_tokens = gr.Slider(1, 1024, value=DEFAULT_MAX_TOKENS, step=1, label="Max new tokens")
    temperature = gr.Slider(0, 2, value=0.8, step=0.05, label="Temperature")
    btn = gr.Button("Generate")
    out = gr.Textbox(label="Output", lines=10)
    btn.click(generate, inputs=[system, user, max_tokens, temperature], outputs=out, api_name="generate")
    btn.click(generate_stream, inputs=[system, user, max_tokens, temperature], outputs=out, api_name="generate_stream")


if __name__ == "__main__":
    demo.launch()