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

from transformers import TextIteratorStreamer
import threading


@spaces.GPU
def chat_with_model(messages):
    global current_model, current_tokenizer
    if current_model is None or current_tokenizer is None:
        yield messages + [{"role": "assistant", "content": "⚠️ No model loaded."}]
        return

    current_model.to("cuda").half()

    pad_id = current_tokenizer.pad_token_id
    if pad_id is None:
        pad_id = current_tokenizer.unk_token_id or 0

    prompt = format_prompt(messages)
    inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)

    streamer = TextIteratorStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=False)

    generation_kwargs = dict(
        **inputs,
        max_new_tokens=256,
        do_sample=True,
        streamer=streamer,
        eos_token_id=current_tokenizer.eos_token_id,
        pad_token_id=pad_id
    )

    thread = threading.Thread(target=current_model.generate, kwargs=generation_kwargs)
    thread.start()

    output_text = ""
    messages = messages.copy()
    messages.append({"role": "assistant", "content": ""})

    for new_text in streamer:
        output_text += new_text
        messages[-1]["content"] = output_text
        yield messages

    current_model.to("cpu")
    torch.cuda.empty_cache()



# Globals
current_model = None
current_tokenizer = None

def load_model_on_selection(model_name, progress=gr.Progress(track_tqdm=False)):
    global current_model, current_tokenizer
    token = os.getenv("HF_TOKEN")

    progress(0, desc="Loading tokenizer...")
    current_tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token)

    progress(0.5, desc="Loading model...")
    current_model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map="cpu",  # loaded to CPU initially
        use_auth_token=token
    )

    progress(1, desc="Model ready.")
    return f"{model_name} loaded and ready!"

# Format conversation as plain text
def format_prompt(messages):
    prompt = ""
    for msg in messages:
        role = msg["role"]
        if role == "user":
            prompt += f"User: {msg['content'].strip()}\n"
        elif role == "assistant":
            prompt += f"Assistant: {msg['content'].strip()}\n"
    prompt += "Assistant:"
    return prompt

def add_user_message(user_input, history):
    return "", history + [{"role": "user", "content": user_input}]

# Available models
model_choices = [
    "meta-llama/Llama-3.2-3B-Instruct",
    "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
    "google/gemma-7b"
]

# UI
with gr.Blocks() as demo:
    gr.Markdown("## Clinical Chatbot (Streaming) — LLaMA, DeepSeek, Gemma")

    default_model = gr.State("meta-llama/Llama-3.2-3B-Instruct")

    # @spaces.GPU
    # def chat_with_model(messages):
    #     global current_model, current_tokenizer
    #     if current_model is None or current_tokenizer is None:
    #         yield messages + [{"role": "assistant", "content": "⚠️ No model loaded."}]
    #         return

    #     current_model = current_model.to("cuda").half()

    #     prompt = format_prompt(messages)
    #     inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)

    #     output_ids = []
    #     messages = messages.copy()
    #     messages.append({"role": "assistant", "content": ""})

    #     for token_id in current_model.generate(
    #         **inputs,
    #         max_new_tokens=256,
    #         do_sample=True,
    #         return_dict_in_generate=True,
    #         output_scores=False
    #     ).sequences[0][inputs['input_ids'].shape[-1]:]:  # skip input tokens
    #         output_ids.append(token_id.item())
    #         decoded = current_tokenizer.decode(output_ids, skip_special_tokens=False)
    #         if output_ids[-1] == current_tokenizer.eos_token_id:
    #             current_model.to("cpu")
    #             torch.cuda.empty_cache()
    #             return
    #         messages[-1]["content"] = decoded
    #         yield messages

    #     current_model.to("cpu")
    #     torch.cuda.empty_cache()
    #     return

    with gr.Row():
        model_selector = gr.Dropdown(choices=model_choices, label="Select Model")
        model_status = gr.Textbox(label="Model Status", interactive=False)

    chatbot = gr.Chatbot(label="Chat", type="messages")
    msg = gr.Textbox(label="Your message", placeholder="Enter clinical input...", show_label=False)
    clear = gr.Button("Clear")

    # Load default model on startup
    demo.load(fn=load_model_on_selection, inputs=default_model, outputs=model_status)

    # Load selected model manually
    model_selector.change(fn=load_model_on_selection, inputs=model_selector, outputs=model_status)

    # Submit message + stream model response
    msg.submit(add_user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
        chat_with_model, chatbot, chatbot
    )

    # Clear chat
    clear.click(lambda: [], None, chatbot, queue=False)

demo.launch()