import gradio as gr
from sentence_transformers import SentenceTransformer
import pinecone
from transformers import GPT2Tokenizer, GPT2LMHeadModel
import torch
import os

# Load secrets and environment variables
PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY")
PINECONE_INDEX_NAME = os.environ.get("PINECONE_INDEX_NAME")
HF_TOKEN = os.environ.get("HF_TOKEN")

# Step 1: Load embedding model and Pinecone
embedding_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
pinecone.init(api_key=PINECONE_API_KEY)
pc = pinecone.Pinecone(api_key=PINECONE_API_KEY)
index = pc.Index(PINECONE_INDEX_NAME)

# Step 2: Load GPT-2 language model
model_name = "HooshvareLab/gpt2-fa"
tokenizer = GPT2Tokenizer.from_pretrained(model_name, use_auth_token=HF_TOKEN)
model = GPT2LMHeadModel.from_pretrained(model_name, use_auth_token=HF_TOKEN)
model.eval()

# Function: Embed input and search in Pinecone
def retrieve_context(query, top_k=1):
    xq = embedding_model.encode(query).tolist()
    res = index.query(vector=xq, top_k=top_k, include_metadata=True)
    if res.matches:
        return res.matches[0].metadata['text']
    return ""

# Function: Generate response using GPT-2
def generate_response(query, context):
    prompt = f"پرسش: {query}\nپاسخ با توجه به اطلاعات زیر: {context}\nپاسخ:"
    input_ids = tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=512)
    output_ids = model.generate(input_ids, max_length=256, num_beams=4, no_repeat_ngram_size=2, early_stopping=True)
    output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return output.split("پاسخ:")[-1].strip()

# Gradio interface
def chat(query):
    context = retrieve_context(query)
    response = generate_response(query, context)
    return response

# UI
with gr.Blocks() as demo:
    gr.Markdown("## چت‌بات هوشمند تیام\nسوالات خود درباره خدمات دیجیتال مارکتینگ تیام را بپرسید.")
    with gr.Row():
        inp = gr.Textbox(label="question", placeholder="سوال خود را وارد کنید")
        out = gr.Textbox(label="output")
    submit = gr.Button("Submit")
    submit.click(chat, inputs=inp, outputs=out)

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