import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  — MUST be imported before torch

import threading  # noqa: E402

import gradio as gr  # noqa: E402
import torch  # noqa: E402
from transformers import (  # noqa: E402
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TextIteratorStreamer,
)

MODEL_ID = "OBLITERATUS/Qwen3.8-27B-OBLITERATED"

bnb = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=bnb,
    device_map="cuda",
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
).eval()


def _duration(message, history, system_prompt, max_new_tokens, enable_thinking, *args, **kwargs):
    """Estimate GPU seconds: cold-start overhead + greedy decode time."""
    try:
        return min(300, 60 + int(float(max_new_tokens) * 0.06))
    except (TypeError, ValueError):
        return 180


@spaces.GPU(duration=_duration)
def chat(message, history, system_prompt, max_new_tokens, enable_thinking):
    """Chat with Qwen3.8-27B-OBLITERATED (4-bit NF4, greedy decoding, rep penalty 1.15).

    Args:
        message: The user's message.
        history: Conversation history (list of {role, content} dicts).
        system_prompt: Optional system prompt. Empty by default — the model
            card recommends no system prompt (system prompts can reintroduce
            refusals).
        max_new_tokens: Generation budget. Card recommends >= 2048.
        enable_thinking: Qwen3 thinking mode. Off by default — thinking
            chains consume token budget and shorten final answers.

    Returns:
        Streamed assistant reply.
    """
    messages = []
    if system_prompt and system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})
    for m in history:
        content = m["content"]
        if isinstance(content, list):
            content = " ".join(
                part.get("text", "") for part in content if isinstance(part, dict)
            )
        messages.append({"role": m["role"], "content": content})
    messages.append({"role": "user", "content": message})

    try:
        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=bool(enable_thinking),
        )
    except TypeError:
        text = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )

    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True
    )
    thread = threading.Thread(
        target=model.generate,
        kwargs=dict(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=False,
            repetition_penalty=1.15,
            streamer=streamer,
        ),
    )
    thread.start()

    partial = ""
    for token in streamer:
        partial += token
        yield partial
    thread.join()


DESCRIPTION = """
# 💥 Qwen3.8-27B — OBLITERATED

[OBLITERATUS/Qwen3.8-27B-OBLITERATED](https://huggingface.co/OBLITERATUS/Qwen3.8-27B-OBLITERATED) · 27B · 4-bit NF4 · ZeroGPU

0% hard refusal across an 842-prompt corpus. Greedy decoding (temp 0) + repetition penalty 1.15 per the
[model card](https://huggingface.co/OBLITERATUS/Qwen3.8-27B-OBLITERATED#%E2%9A%99%EF%B8%8F-optimal-settings--these-matter) — these settings are fixed in the backend for best quality.

> ⚠️ **Research model**: safety guardrails have been surgically removed. Intended for alignment
> research, red-teaming, and AI safety evaluation. You are solely responsible for the content it generates.
"""

with gr.Blocks(theme=gr.themes.Base(), fill_height=True) as demo:
    gr.Markdown(DESCRIPTION)
    gr.ChatInterface(
        fn=chat,
        type="messages",
        title=None,
        additional_inputs=[
            gr.Textbox(
                label="System prompt (optional — card recommends leaving it empty)",
                placeholder="Leave empty: system prompts can reintroduce refusals",
                lines=1,
                value="",
            ),
            gr.Slider(
                label="Max new tokens",
                minimum=256,
                maximum=8192,
                step=256,
                value=2048,
                info="Card recommends ≥ 2048 — complex outputs need room",
            ),
            gr.Checkbox(
                label="Enable thinking mode",
                value=False,
                info="Off by default: thinking chains consume token budget → shorter answers",
            ),
        ],
        examples=[
            ["Write a Python script that monitors a directory and uploads new files to S3."],
            ["Explain how attention works in transformers, with the math."],
            ["Draft a challenging red-team evaluation plan for a customer-support chatbot."],
            ["What are the main failure modes of RAG pipelines and how do you detect them?"],
        ],
        cache_examples=True,
        cache_mode="lazy",
        fill_height=True,
    )

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