import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import chromadb
from chromadb.config import Settings

class RAGPipeline:
    def __init__(self):
        self.tokenizer = AutoTokenizer.from_pretrained("aubmindlab/aragpt2-mega")
        self.model = AutoModelForCausalLM.from_pretrained("aubmindlab/aragpt2-mega")
        self.generator = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer, device=0 if torch.cuda.is_available() else -1)
        self.client = chromadb.Client(Settings(allow_reset=True))
        self.collection = None

    def build_index(self, passages):
        self.client.reset()
        self.collection = self.client.create_collection(name="documents")
        documents = [p["text"] for p in passages]
        metadatas = [{"source": p["source"]} for p in passages]
        ids = [str(i) for i in range(len(documents))]
        self.collection.add(documents=documents, metadatas=metadatas, ids=ids)

    def generate_answer(self, question):
        if not self.collection:
            return "⚠️ لا يوجد فهرس معرف.", []
        results = self.collection.query(query_texts=[question], n_results=3)
        retrieved = [d for d in results['documents'][0]]
        metadatas = results['metadatas'][0]
        prompt = question + "\n" + "\n".join(retrieved)
        result = self.generator(prompt, max_new_tokens=150, do_sample=True)[0]['generated_text']
        return result, [f"{md['source']}: {text[:60]}..." for md, text in zip(metadatas, retrieved)]