import threading

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

# ---------------------------------------------------------------------------
# Fix a real bug in gradio_client 1.5.2 (bundled with ZeroGPU's forced
# gradio==5.9.1): its JSON-schema-to-python-type parser assumes every schema
# node is a dict, but valid JSON Schema allows a bare boolean (e.g.
# additionalProperties: true) — that crashes EVERY API call to this Space
# with "TypeError: argument of type 'bool' is not iterable". Patched here
# rather than pinning an old pydantic (which broke docling's own
# serialization when this Space briefly tried to carry a Documents/RAG tab —
# reverted, see README: the 72B model's 145GB raw download leaves no real
# room in the 150GB per-Space storage cap for anything else).
import gradio_client.utils as _gc_utils

_orig_json_schema_to_python_type = _gc_utils._json_schema_to_python_type


def _patched_json_schema_to_python_type(schema, defs):
    if not isinstance(schema, dict):
        return "bool" if isinstance(schema, bool) else "Any"
    return _orig_json_schema_to_python_type(schema, defs)


_gc_utils._json_schema_to_python_type = _patched_json_schema_to_python_type

MODEL_ID = "Qwen/Qwen2.5-72B-Instruct"
MAX_HISTORY_TURNS = 12

SYSTEM_PROMPTS = {
    "General": "You are a helpful assistant.",
    "Code": (
        "You are an expert software engineer. Answer with correct, working code. "
        "Prefer minimal, direct solutions. Explain tricky parts briefly, in comments "
        "only where the reasoning isn't obvious from the code itself."
    ),
    "Writing": (
        "You are a skilled editor and writer. Write clearly and concisely. "
        "Match the requested tone and length exactly. Avoid filler and cliches."
    ),
    "Structured JSON output": (
        "You output ONLY valid JSON matching what the user asks for — no prose, "
        "no markdown code fences, no explanation before or after. If the request "
        "is ambiguous, make a reasonable schema choice and stay consistent."
    ),
}

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
    ),
    device_map="auto",
    attn_implementation="sdpa",
)


def _duration(message, history, system_prompt_name, max_new_tokens, temperature):
    return min(120, 40 + int(max_new_tokens) // 10)


@spaces.GPU(duration=_duration)
def generate(message, history, system_prompt_name, max_new_tokens, temperature):
    history = history[-(MAX_HISTORY_TURNS * 2):]
    system_prompt = SYSTEM_PROMPTS.get(system_prompt_name, SYSTEM_PROMPTS["General"])

    messages = [{"role": "system", "content": system_prompt}]
    for turn in history:
        messages.append({"role": turn["role"], "content": turn["content"]})
    messages.append({"role": "user", "content": message})

    inputs = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
    ).to(model.device)

    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True
    )
    do_sample = temperature > 0
    gen_kwargs = dict(
        **inputs,
        max_new_tokens=int(max_new_tokens),
        do_sample=do_sample,
        repetition_penalty=1.05,
        pad_token_id=tokenizer.eos_token_id,
        streamer=streamer,
    )
    if do_sample:
        gen_kwargs["temperature"] = float(temperature)
        gen_kwargs["top_p"] = 0.9

    thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
    thread.start()

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


chatbot = gr.Chatbot(type="messages", show_copy_button=True, label="Chat")

demo = gr.ChatInterface(
    fn=generate,
    type="messages",
    chatbot=chatbot,
    additional_inputs=[
        gr.Dropdown(
            list(SYSTEM_PROMPTS.keys()), value="General", label="System prompt"
        ),
        gr.Slider(64, 1024, value=512, step=64, label="Max new tokens"),
        gr.Slider(0.0, 1.5, value=0.7, step=0.1, label="Temperature"),
    ],
    title="Beexly Studio Chat (ZeroGPU)",
    description=(
        f"Personal free-GPU chat — `{MODEL_ID}` (4-bit), streaming, on ZeroGPU. "
        "Hover any reply for a copy button. Not part of any Galaxy Sports Edge product."
    ),
)

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