import os
import threading
from typing import List, Dict, Optional

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

MODEL_ID = "Amey9766/qwen-0.6b-hospitality-housekeeping"

tokenizer = None
model = None
device = None


def load_model(hf_access_token: Optional[str] = None):
    global tokenizer, model, device

    if tokenizer and model:
        return

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

    use_token = hf_access_token or os.getenv("HF_TOKEN")

    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_ID,
        token=use_token,
        trust_remote_code=True,
        use_fast=True,
    )

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        token=use_token,
        torch_dtype=dtype,
        device_map="auto" if device == "cuda" else None,
        trust_remote_code=True,
    )

    if device == "cpu":
        model.to(device)

    print("✅ Loaded model:", getattr(model.config, "_name_or_path", "unknown"))


def build_prompt(system_message: str, history: List[Dict[str, str]], user_message: str) -> str:
    """
    HARD rules to stop:
    - self questioning
    - rule narration
    - exam-style continuation
    """
    hard_rules = (
        "You are a professional hotel housekeeping assistant.\n"
        "STRICT RULES:\n"
        "- Answer ONLY the user's question.\n"
        "- Do NOT generate follow-up questions.\n"
        "- Do NOT mention rules, instructions, or reasoning.\n"
        "- Do NOT narrate your thinking.\n"
        "- Do NOT continue the conversation on your own.\n"
        "- Respond in English only.\n"
        "- Output ONLY the final answer.\n"
    )

    messages = [{"role": "system", "content": hard_rules}]
    messages.extend(history)
    messages.append({"role": "user", "content": user_message})

    if hasattr(tokenizer, "apply_chat_template"):
        return tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )

    prompt = hard_rules + "\n"
    for m in history:
        prompt += f"{m['role'].capitalize()}: {m['content']}\n"
    prompt += f"User: {user_message}\nAssistant:"
    return prompt


def respond(
    message: str,
    history: List[Dict[str, str]],
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    hf_token: gr.OAuthToken,
):
    load_model(hf_token.token if hf_token else None)

    prompt = build_prompt(system_message, history, message)

    inputs = tokenizer(prompt, return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}

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

    gen_kwargs = dict(
        **inputs,
        max_new_tokens=int(max_tokens),
        do_sample=True,
        temperature=0.3,          # 🔒 low temperature = less roleplay
        top_p=0.85,
        repetition_penalty=1.2,   # 🔒 stops looping
        streamer=streamer,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.eos_token_id,
    )

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

    output = ""
    for text in streamer:
        # HARD STOP if model tries to continue conversation
        if any(bad in text.lower() for bad in [
            "now let's",
            "question:",
            "based on the rules",
            "according to the rules",
            "let us",
        ]):
            break

        output += text
        yield output.strip()


chatbot = gr.ChatInterface(
    respond,
    type="messages",
    title="Hospitality Housekeeping Assistant",
    description=f"Running model: `{MODEL_ID}`",
    additional_inputs=[
        gr.Textbox(
            value="Provide SOP-style answers for hotel housekeeping staff.",
            label="System message",
        ),
        gr.Slider(1, 1024, value=256, step=1, label="Max new tokens"),
        gr.Slider(0.1, 1.0, value=0.3, step=0.05, label="Temperature"),
        gr.Slider(0.5, 1.0, value=0.85, step=0.05, label="Top-p"),
    ],
)

with gr.Blocks() as demo:
    with gr.Sidebar():
        gr.LoginButton()
        gr.Markdown(f"**Model:** `{MODEL_ID}`")

    chatbot.render()

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