import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from sentence_transformers import SentenceTransformer, util
from pinecone import Pinecone
import gradio as gr

# انتشار متغیرها از Hugging Face secrets
HF_TOKEN = os.getenv("HF_TOKEN")
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
PINECONE_INDEX_NAME = os.getenv("PINECONE_INDEX_NAME")

# مدل کوچک برای embedding (sentence-transformers)
embedding_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")

# مدل زبانی GPT2 فارسی
tokenizer = AutoTokenizer.from_pretrained("HooshvareLab/gpt2-fa")
model = AutoModelForCausalLM.from_pretrained("HooshvareLab/gpt2-fa")

# اتصال به Pinecone
pc = Pinecone(api_key=PINECONE_API_KEY)
index = pc.Index(PINECONE_INDEX_NAME)

# توابع

def retrieve_from_pinecone(query):
    query_embedding = embedding_model.encode(query).tolist()
    search_result = index.query(vector=query_embedding, top_k=1, include_metadata=True)
    try:
        return search_result['matches'][0]['metadata']['answer']
    except:
        return "پاسخی برای این سوال پیدا نشد."

def generate_response(query):
    base_answer = retrieve_from_pinecone(query)
    prompt = f"{query}\n{base_answer}"

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

    output = model.generate(
        inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        max_new_tokens=30,  # کمک به تسریع پاسخگویی
        temperature=0.7,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id
    )

    response = tokenizer.decode(output[0], skip_special_tokens=True)

    # جدا کردن پاسخ تولیدی از prompt
    return response.replace(prompt, "").strip()

# رابط کاربری Gradio
iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(label="question", placeholder="سوال خود را وارد کنید"),
    outputs=gr.Textbox(label="output"),
    title="چتبات هوشمند تیام",
    description="سوالات خود درباره خدمات دیجیتال مارکتینگ تیام را بپرسید"
)

iface.launch()