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

MODEL_ID = "OBLITERATUS/Llama-3.1-8B-Lexi-Uncensored-V2-OBLITERATED"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
).to("cuda")

model.eval()
print("Model ready.")


def extract_text(content):
    if content is None:
        return ""

    if isinstance(content, str):
        return content

    if isinstance(content, list):
        parts = []

        for block in content:
            text = extract_text(block)

            if text:
                parts.append(text)

        return "\n".join(parts)

    if isinstance(content, dict):
        if content.get("type") == "text":
            return extract_text(content.get("text"))

        if "text" in content:
            return extract_text(content["text"])

        if "content" in content:
            return extract_text(content["content"])

    return ""


@spaces.GPU(duration=120)
def chat(message, history):

    # Lexi V2 expects a system turn to exist.
    # "." is the model author's recommended effectively-empty value.
    messages = [
        {
            "role": "system",
            "content": "."
        }
    ]

    for item in history or []:
        if not isinstance(item, dict):
            continue

        role = item.get("role")

        if role not in ("user", "assistant"):
            continue

        text = extract_text(item.get("content"))

        if text:
            messages.append({
                "role": role,
                "content": text
            })

    user_text = extract_text(message)

    messages.append({
        "role": "user",
        "content": user_text
    })

    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    )

    inputs = {
        key: value.to("cuda")
        for key, value in inputs.items()
    }

    with torch.inference_mode():
        output = model.generate(
            **inputs,
            max_new_tokens=4000,
            do_sample=True,
            temperature=1.0,
            top_p=1.0,
            repetition_penalty=1.0,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
            use_cache=True,
        )

    prompt_length = inputs["input_ids"].shape[-1]

    new_tokens = output[0, prompt_length:]

    response = tokenizer.decode(
        new_tokens,
        skip_special_tokens=True
    )

    return response.strip()


demo = gr.ChatInterface(
    fn=chat,
    title="Lexi 8B — OBLITERATED",
    description="Native Lexi framing • temp 1.0 • top-p 1.0 • 4000 max tokens",
)

demo.launch()