import spaces  # must be imported before torch/transformers on ZeroGPU
import os
import threading

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

MODEL_ID = "vikas117/llama31-8b-support-agent"  # <-- confirm this matches your model repo id

SUPPORT_AGENT_SYSTEM_PROMPT = (
    "You are a customer support agent for an online retail company. Respond to the customer's "
    "message clearly, politely, and concisely. Only state policies, refund amounts, or order "
    "details that are explicitly given to you or that are standard, safe assumptions (e.g. "
    "acknowledging a request and explaining next steps). If the request involves a legal threat, "
    "an unusual or high-value refund, fraud, or anything you are not confident you can resolve "
    "correctly, say so plainly and tell the customer you are escalating them to a human agent "
    "rather than guessing. Do not invent specific dollar amounts, order numbers, or dates that "
    "were not provided to you."
)

print("Loading model, this can take a minute on cold start...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    device_map="auto",
)
model.eval()
print("Model loaded.")


@spaces.GPU(duration=60)  # ZeroGPU: reserves a GPU for up to 60s per call
def respond(message, history, max_new_tokens):
    messages = [{"role": "system", "content": SUPPORT_AGENT_SYSTEM_PROMPT}]
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})

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

    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True
    )
    generation_kwargs = dict(
        **inputs,
        max_new_tokens=int(max_new_tokens),
        do_sample=False,
        temperature=None,
        top_p=None,
        pad_token_id=tokenizer.pad_token_id,
        streamer=streamer,
    )

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

    partial_text = ""
    for new_text in streamer:
        partial_text += new_text
        yield partial_text


with gr.Blocks(title="Support Auto-Responder (Llama-3.1-8B, fine-tuned)") as demo:
    gr.Markdown(
        "# Customer Support Auto-Responder\n"
        "QLoRA fine-tuned Llama-3.1-8B-Instruct. **Draft-suggestion tool** — "
        "review before sending; not evaluated for autonomous customer-facing use."
    )
    max_new_tokens = gr.Slider(
        minimum=32, maximum=512, value=200, step=8, label="Max new tokens"
    )
    gr.ChatInterface(
        fn=respond,
        additional_inputs=[max_new_tokens],
        examples=[
            ["My order #48213 hasn't arrived and it's been 2 weeks.", 200],
            ["I want to return a product I bought a month ago, is that possible?", 200],
            ["I was charged twice for the same order, this is fraud and I'm contacting my lawyer.", 200],
        ],
        title=None,
    )

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