import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import spaces
import json
from threading import Thread, Lock
from typing import Optional
from quota_guard import QuotaBlocked, verify_free_quota
from gguf_runtime import generate_gguf
import gradio as gr
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

MODEL_ID = "SEE7/Astheral"
MODEL_REVISION = "37adef9b3917e4a3de0bb77f0db71089b7b82202"
settings = json.load(open(hf_hub_download(MODEL_ID, "settings.json", revision=MODEL_REVISION)))
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, revision=MODEL_REVISION)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, revision=MODEL_REVISION, torch_dtype=torch.bfloat16,
    device_map={"": "cuda"}, low_cpu_mem_usage=True,
    attn_implementation="sdpa", use_safetensors=True,
).eval()

@spaces.GPU(duration=120)
def _generate(message, history, max_tokens, temperature, mode, request: gr.Request):
    if not message or not message.strip():
        raise gr.Error("Enter a message first.")
    if len(message) > 32000:
        raise gr.Error("Please keep your message under 32,000 characters.")
    system = next((item["content"] for item in reversed(history) if item.get("role") == "system" and isinstance(item.get("content"), str)), settings["system"])
    if mode == "Thinking":
        system += "\n" + settings["thinking"]
    messages = [{"role": "system", "content": system}]
    # Keep the system instruction and the most recent complete turns within 8K tokens.
    turns = [item for item in history[-12:] if item.get("role") in ("user", "assistant") and isinstance(item.get("content"), str)]
    latest = {"role": "user", "content": message.strip()}
    while turns and len(tokenizer.apply_chat_template(messages + turns + [latest], tokenize=True, add_generation_prompt=True)) > 8192:
        turns = turns[2:]
    inputs = tokenizer.apply_chat_template(
        messages + turns + [latest], tokenize=True, add_generation_prompt=True,
        return_dict=True, return_tensors="pt",
    )
    if inputs["input_ids"].shape[1] > 8192:
        raise gr.Error("That message is too long for this budget-friendly demo. Please shorten it.")
    inputs = inputs.to("cuda")
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=115)
    generation = dict(
        **inputs, streamer=streamer, max_new_tokens=min(max(int(max_tokens), 1), 4096),
        do_sample=float(temperature) > 0, max_time=110,
        pad_token_id=128004, eos_token_id=128009,
    )
    if float(temperature) > 0:
        generation.update(temperature=float(temperature), top_p=0.9)
    errors = []
    def generate():
        try:
            with torch.inference_mode():
                model.generate(**generation)
        except Exception as exc:
            errors.append(exc)
            streamer.end()
    worker = Thread(target=generate, daemon=True)
    worker.start()
    response = ""
    try:
        for text in streamer:
            response += text
            yield response
    finally:
        worker.join(timeout=2)
    if errors:
        raise gr.Error("Generation failed. Please retry with a shorter prompt.") from errors[0]
    if not response.strip():
        raise gr.Error("No text was generated. Please try again.")

quota_lock = Lock()

def respond(message, history, max_tokens, temperature, mode, request: gr.Request):
    # Check before entering the decorated function: even a refused request uses no GPU.
    with quota_lock:
        try:
            verify_free_quota()
        except QuotaBlocked as exc:
            raise gr.Error(str(exc)) from None
        yield from _generate(message, history, max_tokens, temperature, mode, request)

def respond_gguf(model_id, message, history, max_tokens, temperature, request: gr.Request):
    from gguf_models import model_path
    from gguf_runtime import generate_gguf
    with quota_lock:
        try:
            verify_free_quota()
            path = model_path(model_id)
        except (QuotaBlocked, ValueError) as exc:
            raise gr.Error(str(exc)) from None
        if not isinstance(message, str) or not message.strip() or len(message) > 32000:
            raise gr.Error("Enter a message under 32,000 characters.")
        messages = [item for item in history[-12:] if item.get("role") in ("system", "user", "assistant") and isinstance(item.get("content"), str)]
        messages.append({"role": "user", "content": message})
        try:
            yield from generate_gguf(path, messages, min(max(int(max_tokens), 1), 4096), float(temperature))
        except Exception as exc:
            raise gr.Error("GGUF generation failed: " + str(exc)[:300]) from None

with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo"), title="Astheral · Built with Llama") as demo:
    with gr.Accordion("Response settings", open=False):
        max_tokens = gr.Slider(32, 4096, value=1024, step=32, label="Maximum new tokens")
        temperature = gr.Slider(0, 1, value=0.6, step=0.1, label="Temperature (0 = deterministic)")
        mode = gr.Radio(["Instruct", "Thinking"], value="Instruct", label="Mode")
    gr.ChatInterface(
        fn=respond, type="messages", additional_inputs=[max_tokens, temperature, mode],
        examples=[["Explain Python generators in three sentences.", 192, 0.6, "Instruct"], ["Write a short starship announcement.", 192, 0.6, "Instruct"]],
        cache_examples=False, concurrency_limit=1, api_name="chat",
    )
with demo:
    gguf_id = gr.Textbox(visible=False)
    gguf_message = gr.Textbox(visible=False)
    gguf_history = gr.JSON(visible=False)
    gguf_tokens = gr.Number(value=1024, visible=False)
    gguf_temperature = gr.Number(value=0.6, visible=False)
    gguf_output = gr.Textbox(visible=False)
    gguf_trigger = gr.Button(visible=False)
    gguf_trigger.click(respond_gguf, [gguf_id, gguf_message, gguf_history, gguf_tokens, gguf_temperature],
        gguf_output, api_name="gguf_chat", concurrency_limit=1)
demo.queue(max_size=6, default_concurrency_limit=1)
if __name__ == "__main__":
    from llama_ui import install_ui
    demo.launch(prevent_thread_lock=True, show_error=True, ssr_mode=False)
    install_ui(demo.app)
    demo.block_thread()