import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from woocommerce import API

# اتصال به ووکامرس
wcapi = API(
    url="https://vashnia.com",
    consumer_key="ck_f284e213686b919d3f4552dab91a336543274b04",
    consumer_secret="cs_15fd20967c669efa724f7a0c683a11910ea441e6",
    timeout=50
)

# جستجوی محصولات و ساخت جدول HTML
def search_products(query):
    try:
        res = wcapi.get("products", params={"search": query})
        products = res.json()
        if not isinstance(products, list) or len(products) == 0:
            return "❌ محصولی با این نام پیدا نشد."

        table = "<div style='display:flex; justify-content:center;'><table style='border-collapse:collapse; text-align:center; direction:rtl;'>"
        table += "<thead><tr><th style='padding:10px; border:1px solid #ccc;'>نام محصول</th><th style='padding:10px; border:1px solid #ccc;'>قیمت و وزن</th><th style='padding:10px; border:1px solid #ccc;'>خرید</th></tr></thead><tbody>"

        for product in products:
            name = product.get("name", "نامشخص")
            permalink = product.get("permalink", "#")
            price_section = ""

            # گرفتن قیمت و وزن از وارییشن‌ها
            variations = product.get("variations", [])
            if variations:
                for var_id in variations:
                    var_res = wcapi.get(f"products/variations/{var_id}")
                    var = var_res.json()
                    weight = "-"
                    for attr in var.get("attributes", []):
                        if "وزن" in attr.get("name", ""):
                            weight = attr.get("option", "")
                    price = var.get("price", "نامشخص")
                    price_section += f"{weight} : {price} تومان<br>"
            else:
                price = product.get("price", "نامشخص")
                price_section = f"{price} تومان"

            table += f"""
            <tr>
                <td style='padding:10px; border:1px solid #ccc;'>{name}</td>
                <td style='padding:10px; border:1px solid #ccc;'>{price_section}</td>
                <td style='padding:10px; border:1px solid #ccc;'>
                    <a href='{permalink}' target='_blank'>
                        <button style='background:#f97316; color:#fff; padding:8px 12px; border:none; border-radius:5px;'>خرید محصول</button>
                    </a>
                </td>
            </tr>
            """

        table += "</tbody></table></div>"
        return table
    except Exception as e:
        return f"❌ خطا در جستجوی محصولات: {str(e)}"

# بارگذاری مدل جدید GPT2 فارسی
model_name = "HooshvareLab/gpt2-fa"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

# تابع اصلی چت
def chat_with_agent(message, history):
    message_lower = message.strip().lower()

    # بررسی اگر پیام درباره قیمت محصول باشد
    if any(x in message_lower for x in ["قیمت", "برنج", "پسته", "فندق", "بادام", "تخمه", "کشمش", "هندی"]):
        product_name = message_lower.replace("قیمت", "").strip()
        return search_products(product_name)

    # پاسخ عمومی با GPT2 فارسی
    prompt = f"پرسش: {message}\nپاسخ:"
    result = generator(prompt, max_length=80, num_return_sequences=1, do_sample=True)
    response = result[0]['generated_text'].replace(prompt, "").strip()
    return response

# رابط چت Gradio
chat = gr.ChatInterface(
    fn=chat_with_agent,
    title="🛒 ایجنت چت فروشگاه وش نیا",
    description="با ما چت کن، قیمت بپرس، سوال کن 😊",
    chatbot=gr.Chatbot(height=450),
    textbox=gr.Textbox(placeholder="مثلاً: قیمت بادام شور یا سلام 👋", label="پیام شما"),
    theme="soft"
)

chat.launch()