import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch
import spaces

# Custom inline CSS for dark/black theme
custom_css = """
:root {
  --color-bg-light: #000000;
  --color-accent: #35a387;
  --color-primary-dark: #0e3229;
  --color-text-dark: #f3f3f3;
}
"""

# Force dark mode + black theme overrides
force_dark_css = """
<style>
html {
  color-scheme: dark;
}

body {
  background-color: #000000 !important;
  color: #f3f3f3 !important;
}

.gr-box,
.gr-panel,
.gr-chatbox,
input,
textarea {
  background-color: #000000 !important;
  border-color: #333 !important;
  color: #f3f3f3 !important;
}

button.gr-button {
  background-color: #35a387 !important;
  color: white !important;
}
</style>
"""

# Load model and tokenizer
model_name = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    low_cpu_mem_usage=True
)

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=275,
    temperature=0.7,
    top_p=0.95,
    do_sample=True
)

@spaces.GPU(duration=60)
def sales_agent(message, chat_history=[]):
    # Manually map prompt prefix to flag
    if message.startswith("🇫🇷"):
        lang_flag = "🇫🇷"
        prompt_lang = "en français"
    elif message.startswith("🇳🇱"):
        lang_flag = "🇳🇱"
        prompt_lang = "in het Nederlands"
    else:
        lang_flag = "🇬🇧"
        prompt_lang = "in English"

    prompt = f"""
You are a helpful pre-sales agent at Cotubex, a tech store in Brussels. Answer in {prompt_lang}.

Product List:
- Samsung 990 PRO 1TB NVME – EUR 119.00 – In Stock
- Travel adapter Europe to Switzerland + Italy + Brazil – EUR 15.79 – In Stock
- Be Quiet! Pure Loop 2 240 Watercooling – EUR 109.90 – In Stock
- Zotac 5060 TI 16GB OC – EUR 535.00 – In Stock
- Logitech G502 HERO – EUR 49.99 – In Stock

Question: {message}

Answer ({lang_flag}):
"""

    response = pipe(prompt)
    answer = response[0]['generated_text'].replace(prompt, "").strip()
    return f"{lang_flag} {answer}"

# Example prompts
examples = [
    ["🇫🇷 Quel est le prix de la carte graphique Zotac 5060 TI ?"],
    ["🇳🇱 Wat kost de Zotac 5060 TI videokaart?"],
    ["🇬🇧 Is the Be Quiet water cooler available?"]
]

# Build Gradio interface
with gr.Blocks(css=custom_css) as demo:
    gr.HTML(force_dark_css)
    gr.Markdown("### Cotubex Pre-Sales Assistant 🇫🇷🇧🇪🇳🇱\nAsk us anything about products, pricing, availability.")
    gr.ChatInterface(
        fn=sales_agent,
        chatbot=gr.Chatbot(height=400, type="messages"),
        examples=examples
    )

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