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

model_id = "h2oai/h2o-danube2-1.8b-chat"

tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float32,
    device_map="cpu"
)

torch.set_num_threads(2)

COMPANY_NAME = "TSF"

SYSTEM_PROMPT = f"""You are the official customer support assistant for {COMPANY_NAME}.

Rules:
- Be professional and friendly
- Give clear, concise answers
- If unsure, offer to connect with a human agent
- Never invent information about products or policies
- Keep responses short and helpful
- Only reply as the assistant, do not generate user messages
- Stop after your single response
- Never output JSON, lists of dicts, or code formatting in your reply
- Reply in plain natural language only"""


def clean_response(response):
    # Remove any dict/list formatting like [{'text': '...', 'type': '...'}]
    pattern = r"\[\{'text':\s*'(.*?)',\s*'type':\s*'text'\}\]"
    match = re.search(pattern, response)
    if match:
        response = match.group(1)

    # Remove fake conversation continuations
    cut_markers = [
        "User:", "user:", "Human:", "human:",
        "Assistant:", "assistant:",
        "\nA:", "\nQ:",
    ]
    for marker in cut_markers:
        if marker in response:
            response = response[:response.index(marker)]

    # Remove escape characters
    response = response.replace("\\n", "\n")
    response = response.replace("\\'", "'")
    response = response.replace('\\"', '"')

    # Clean up whitespace
    response = response.strip()

    return response


def generate_response(message, history):
    message = message.strip()

    if not message:
        return "How can I assist you today?"

    greetings = ["hi", "hello", "hey", "good morning", "good afternoon"]
    if message.lower() in greetings:
        return f"Hello! Welcome to {COMPANY_NAME}. How can I help you today?"

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    for entry in history:
        if isinstance(entry, dict) and "role" in entry and "content" in entry:
            messages.append({"role": entry["role"], "content": entry["content"]})

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

    try:
        input_text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
    except Exception:
        prompt = ""
        for m in messages:
            if m["role"] == "system":
                prompt += f"System: {m['content']}\n"
            elif m["role"] == "user":
                prompt += f"User: {m['content']}\n"
            elif m["role"] == "assistant":
                prompt += f"Assistant: {m['content']}\n"
        prompt += "Assistant:"
        input_text = prompt

    inputs = tokenizer(input_text, return_tensors="pt")

    # Get the input length to extract only new tokens
    input_length = inputs["input_ids"].shape[1]

    outputs = model.generate(
        **inputs,
        max_new_tokens=150,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id,
        repetition_penalty=1.2
    )

    # Decode only the NEW tokens (not the input)
    new_tokens = outputs[0][input_length:]
    response = tokenizer.decode(new_tokens, skip_special_tokens=True)

    # Clean the response
    response = clean_response(response)

    if not response:
        response = "I apologize, could you please rephrase your question?"

    return response


with gr.Blocks(title=f"{COMPANY_NAME} Support") as demo:

    gr.Markdown(f"""
    ## 💬 {COMPANY_NAME} Customer Support
    Welcome! Ask me anything about our services.
    """)

    chatbot = gr.Chatbot(height=400)
    msg = gr.Textbox(placeholder="Type your question...", show_label=False)
    send_btn = gr.Button("Send", variant="primary")

    def respond(message, chat_history):
        if chat_history is None:
            chat_history = []
        if not message.strip():
            return "", chat_history

        bot_reply = generate_response(message, chat_history)

        chat_history.append({"role": "user", "content": message})
        chat_history.append({"role": "assistant", "content": bot_reply})

        return "", chat_history

    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    send_btn.click(respond, [msg, chatbot], [msg, chatbot])

    gr.Markdown("---\n*AI-powered assistant. For urgent matters, contact our support team.*")

demo.launch()