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

# The original WhiteRabbitNeo/WhiteRabbitNeo-V3-7B repository redirects here.
MODEL_ID = "DeepHat/DeepHat-V1-7B"
SYSTEM_PROMPT = (
    "You are DeepHat, a helpful assistant specializing in cybersecurity, "
    "programming, and DevOps. Answer in the user's language. "
    "Be precise, explain assumptions, and acknowledge uncertainty."
)
MAX_INPUT_TOKENS = 4096

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
# ZeroGPU requires CUDA placement at module scope, outside the GPU function.
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map="cuda",
    low_cpu_mem_usage=True,
    attn_implementation="sdpa",
).eval()


def text_content(content):
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return "\n".join(
            item.get("text", "")
            for item in content
            if isinstance(item, dict) and item.get("type") == "text"
        )
    return ""


def prepare_inputs(message, history):
    if not message.strip():
        raise gr.Error("Введите сообщение.")
    recent = []
    for item in (history or [])[-12:]:
        if item.get("role") in ("user", "assistant"):
            content = text_content(item.get("content", ""))
            if content:
                recent.append({"role": item["role"], "content": content})
    while recent and recent[0]["role"] != "user":
        recent.pop(0)
    while True:
        messages = (
            [{"role": "system", "content": SYSTEM_PROMPT}]
            + recent
            + [{"role": "user", "content": message}]
        )
        prompt = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
        if inputs.input_ids.shape[1] <= MAX_INPUT_TOKENS:
            return inputs
        if not recent:
            raise gr.Error("Сообщение слишком длинное. Сократите его и попробуйте ещё раз.")
        recent.pop(0)
        while recent and recent[0]["role"] != "user":
            recent.pop(0)


@spaces.GPU(duration=60)
def generate(inputs, max_tokens):
    inputs = inputs.to("cuda")
    with torch.inference_mode():
        output = model.generate(
            **inputs,
            max_new_tokens=int(max_tokens),
            do_sample=True,
            temperature=0.6,
            top_p=0.9,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.eos_token_id,
        )
    return tokenizer.decode(
        output[0, inputs.input_ids.shape[1]:], skip_special_tokens=True
    ).strip()


def respond(message, history, max_tokens):
    return generate(prepare_inputs(message, history), max_tokens)


demo = gr.ChatInterface(
    fn=respond,
    type="messages",
    title="WhiteRabbitNeo / DeepHat 7B",
    description=(
        "Помощник по кибербезопасности, программированию и DevOps. "
        "Ответ появится после завершения генерации. "
        "ZeroGPU работает с очередью и дневной квотой. "
        "Длинная история сокращается до последних сообщений."
    ),
    additional_inputs=[
        gr.Slider(64, 512, value=256, step=64, label="Максимальная длина ответа в токенах")
    ],
    examples=[
        ["Объясни разницу между XSS и CSRF и способы защиты.", 256],
        ["Помоги составить план проверки безопасности моего веб-приложения.", 256],
    ],
    cache_examples=False,
    concurrency_limit=1,
)

if __name__ == "__main__":
    demo.queue(max_size=8).launch()