import threading

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

MODEL_ID = "KreatixLabs/Kreatix-K1"
DEFAULT_SYSTEM = "You are Kreatix K1, created by KreatixLabs. You are a helpful coding assistant."

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float16).to("cuda")
model.eval()


def to_text(content):
    if isinstance(content, str):
        return content
    if isinstance(content, dict):
        return content.get("text", "") or ""
    if isinstance(content, (list, tuple)):
        return "".join(to_text(p) for p in content)
    return "" if content is None else str(content)


@spaces.GPU(duration=120)
def chat(message, history, system_prompt, temperature, top_p, max_tokens):
    messages = []
    if system_prompt and system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})
    for m in history or []:
        role = m.get("role") if isinstance(m, dict) else None
        if role in ("user", "assistant"):
            text = to_text(m.get("content"))
            if text:
                messages.append({"role": role, "content": text})
    messages.append({"role": "user", "content": to_text(message)})

    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    temperature = float(temperature)
    kwargs = dict(
        **inputs,
        max_new_tokens=int(max_tokens),
        do_sample=temperature > 0.05,
        temperature=max(temperature, 0.05),
        top_p=float(top_p),
        repetition_penalty=1.05,
        pad_token_id=tokenizer.eos_token_id,
        streamer=streamer,
    )
    worker = threading.Thread(target=model.generate, kwargs=kwargs)
    worker.start()
    out = ""
    for tok in streamer:
        out += tok
        yield out
    worker.join()


demo = gr.ChatInterface(
    fn=chat,
    title="Kreatix K1",
    description="A 7B open coding model by KreatixLabs. Weights and GGUF: [KreatixLabs/Kreatix-K1](https://huggingface.co/KreatixLabs/Kreatix-K1)",
    additional_inputs=[
        gr.Textbox(value=DEFAULT_SYSTEM, label="System prompt", lines=2),
        gr.Slider(0.0, 2.0, value=0.7, step=0.05, label="Temperature"),
        gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p"),
        gr.Slider(64, 2048, value=768, step=64, label="Max new tokens"),
    ],
    examples=[
        ["Who are you?"],
        ["Write a Python function that merges two sorted lists."],
        ["My code won't work."],
        ["Explain async/await in JavaScript like I'm tired."],
    ],
    cache_examples=False,
    fill_height=True,
)

demo.queue(default_concurrency_limit=4).launch()