import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForCausalLM

class RAG:
    def __init__(self):
        # Multilingual embedding -> مناسب فارسی + انگلیسی
        self.embedder = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")

        # LLM multilingual lightweight
        self.tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B-Instruct")
        self.model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-3B-Instruct")


        # Load dataset
        with open("data.txt", "r", encoding="utf-8") as f:
            self.texts = [line.strip() for line in f if line.strip()]

        # Build FAISS index
        embeddings = self.embedder.encode(self.texts)
        dim = embeddings.shape[1]
        self.index = faiss.IndexFlatL2(dim)
        self.index.add(np.array(embeddings, dtype=np.float32))

    def retrieve(self, query, k=4):
        q_emb = self.embedder.encode([query])
        distances, indices = self.index.search(
            np.array(q_emb, dtype=np.float32), k
        )
        results = [self.texts[i] for i in indices[0]]
        return results

    def generate_answer(self, question):
        retrieved = self.retrieve(question)

        # No relevant data found
        if not retrieved or len("".join(retrieved).strip()) == 0:
            return "در داده‌هایی که دارم پاسخی برای این سؤال وجود ندارد."

        context = "\n".join(retrieved)

        prompt = f"""
شما یک مدل زبانی هستید که فقط و فقط اجازه دارید بر اساس اطلاعات زیر پاسخ بدهید.
هیچ دانشی خارج از متن‌ها استفاده نکنید، و هیچ خلاقیتی نداشته باشید.
پاسخ باید کاملاً فارسی، دقیق و برگرفته از داده‌ها باشد.

اطلاعات:
{context}

سؤال: {question}

پاسخ دقیق فارسی بر اساس اطلاعات بالا:
"""

        inputs = self.tokenizer(prompt, return_tensors="pt")
        output = self.model.generate(
            **inputs,
            max_new_tokens=200,
            do_sample=False,      # no creativity
            temperature=0.0,      # deterministic
        )

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