import torch
from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer


# ==================================================
# CONFIG
# ==================================================

# Either the HF hub id (auto-downloads + caches on first run):
MODEL_PATH = "Qwen/Qwen2.5-0.5B-Instruct"
# ...or a local folder if you already downloaded it yourself
# (e.g. via snapshot_download or git lfs clone), e.g.:
# MODEL_PATH = "./Qwen2.5-0.5B-Instruct"

SYSTEM_PROMPT = "You are a helpful assistant."

MAX_NEW_TOKENS = 512
HISTORY_TURNS = 6  # number of past user/assistant exchanges kept as context

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32


# ==================================================
# LOAD MODEL
# ==================================================

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)

print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, torch_dtype=dtype)
model.to(device)
model.eval()

print("Device:", device)


# ==================================================
# PROMPT BUILDING
# ==================================================

def build_prompt(history, user_input):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    for role, text in history:
        messages.append({"role": role, "content": text})

    messages.append({"role": "user", "content": user_input})

    return tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )


# ==================================================
# GENERATION (streamed)
# ==================================================

def generate_reply(history, user_input):
    prompt = build_prompt(history, user_input)
    inputs = tokenizer(prompt, return_tensors="pt").to(device)

    streamer = TextIteratorStreamer(
        tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
    )

    generation_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=MAX_NEW_TOKENS,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.eos_token_id,
    )

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

    full_text = ""
    for token in streamer:
        print(token, end="", flush=True)
        full_text += token

    thread.join()
    print()

    return full_text


# ==================================================
# CHAT
# ==================================================

print("\n" + "=" * 60)
print("Qwen2.5-0.5B-Instruct Chatbot")
print("Type exit to quit")
print("=" * 60)

history = []  # list of (role, text) tuples, most recent last

while True:
    user_input = input("\nUser: ")

    if user_input.lower() == "exit":
        break

    if not user_input.strip():
        continue

    print("\nAssistant:")
    reply = generate_reply(history, user_input)

    history.append(("user", user_input))
    history.append(("assistant", reply))
    history = history[-HISTORY_TURNS * 2:]