import os
from threading import Thread
import gradio as gr
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TextIteratorStreamer,
)

# Support spaces.GPU if available on ZeroGPU
try:
    import spaces
    has_spaces = True
except ImportError:
    has_spaces = False

def gpu_decorator(fn):
    if has_spaces:
        return spaces.GPU(duration=120)(fn)
    return fn

MODEL_ID = "Hatsunama/Qwen3-4B-OBLITERATED-repaired"
device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Loading tokenizer from {MODEL_ID}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id or 151643

print(f"Loading model {MODEL_ID} on {device}...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
    trust_remote_code=True
).to(device)

if hasattr(model, "generation_config") and model.generation_config is not None:
    model.generation_config.max_length = None

def strip_thought(text):
    """Strip reasoning/thought block from historical turns for cleaner multi-turn context."""
    if "</think>" in text:
        return text.split("</think>")[-1].strip()
    return text

def format_display_text(raw_text):
    """Format <think>...</think> block into a collapsible thought accordion for Gradio chat."""
    if "<think>" in raw_text:
        parts = raw_text.split("</think>")
        if len(parts) > 1:
            thought = parts[0].replace("<think>", "").strip()
            answer = parts[1].strip()
            return f"<details open><summary>💭 <i>Thought Process</i></summary>\n\n{thought}\n\n</details>\n\n{answer}"
        else:
            thought = parts[0].replace("<think>", "").strip()
            return f"<details open><summary>💭 <i>Thinking...</i></summary>\n\n{thought}\n\n</details>"
    return raw_text

@gpu_decorator
def predict(message, history):
    # Format messages
    conversation = [
        {"role": "system", "content": "You are a helpful, respectful, and honest assistant."}
    ]
    
    if history:
        for item in history:
            if isinstance(item, dict):
                if item.get("role") == "system":
                    continue
                role = item.get("role", "user")
                content = item.get("content", "")
                if role == "assistant":
                    content = strip_thought(content)
                conversation.append({"role": role, "content": content})
            elif isinstance(item, (list, tuple)) and len(item) == 2:
                conversation.append({"role": "user", "content": str(item[0])})
                conversation.append({"role": "assistant", "content": strip_thought(str(item[1]))})
                
    conversation.append({"role": "user", "content": str(message)})
    
    text = tokenizer.apply_chat_template(
        conversation,
        tokenize=False,
        add_generation_prompt=True
    )
    
    model_inputs = tokenizer([text], return_tensors="pt").to(device)
    
    streamer = TextIteratorStreamer(
        tokenizer,
        timeout=120.0,
        skip_prompt=True,
        skip_special_tokens=True
    )
    
    pad_id = tokenizer.pad_token_id or 151643
    eos_tokens = [151645, 151643]
    if tokenizer.eos_token_id is not None and tokenizer.eos_token_id not in eos_tokens:
        eos_tokens.append(tokenizer.eos_token_id)
        
    generate_kwargs = dict(
        model_inputs,
        streamer=streamer,
        max_new_tokens=2048,
        do_sample=True,
        top_p=0.95,
        top_k=20,
        temperature=0.6,
        repetition_penalty=1.15,
        pad_token_id=pad_id,
        eos_token_id=eos_tokens,
    )
    
    errors = []
    def worker():
        try:
            model.generate(**generate_kwargs)
        except Exception as e:
            errors.append(e)
            streamer.end()
            
    thread = Thread(target=worker)
    thread.daemon = True
    thread.start()
    
    try:
        partial_text = ""
        for new_text in streamer:
            partial_text += new_text
            yield format_display_text(partial_text)
    except Exception as e:
        if errors:
            yield f"⚠️ Generation error: {errors[0]}"
        else:
            yield f"⚠️ Stream error: {e}"

demo = gr.ChatInterface(
    fn=predict,
    title="🧠 Qwen3-4B-OBLITERATED Chat",
    description="Live un-censored / abliterated Qwen 4B model running on Hugging Face ZeroGPU. Supports dynamic reasoning (`/think` and `/no_think`)."
)

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