import os
import pandas as pd
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.text_splitter import CharacterTextSplitter
from langchain.docstore.document import Document
from transformers import pipeline
from langchain.prompts import PromptTemplate
from typing import List, Dict, Any, Optional

class RAGSystem:
    def __init__(self, sql_generator: SQLGenerator, csv_path: str = "apparel.csv"):
        self.sql_generator = sql_generator
        self.setup_system(csv_path)
        self.qa_pipeline = pipeline("question-answering", model="distilbert-base-cased-distilled-squad")
        
    def setup_system(self, csv_path: str):
        if not os.path.exists(csv_path):
            raise FileNotFoundError(f"CSV file not found at {csv_path}")
            
        documents = pd.read_csv(csv_path)
        
        docs = [
            Document(
                page_content=str(row['Title']),
                metadata={'index': idx}
            ) 
            for idx, row in documents.iterrows()
        ]
        
        text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
        split_docs = text_splitter.split_documents(docs)
        
        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
        self.vector_store = FAISS.from_documents(split_docs, embeddings)
        self.retriever = self.vector_store.as_retriever()
        
    def process_query(self, query: str, execute_sql: bool = True) -> Dict[str, Any]:
        """
        Process a query through both RAG and SQL if needed
        """
        # Get relevant documents
        retrieved_docs = self.retriever.get_relevant_documents(query)
        retrieved_text = "\n".join([doc.page_content for doc in retrieved_docs])[:1000]
        
        # Process with QA pipeline
        qa_input = {
            "question": query,
            "context": retrieved_text
        }
        qa_response = self.qa_pipeline(qa_input)
        
        result = {
            "qa_answer": qa_response['answer'],
            "relevant_docs": [doc.page_content for doc in retrieved_docs[:3]],
            "sql_results": None
        }
        
        # If SQL execution is requested and SQL is detected in the query
        if execute_sql and "SELECT" in query.upper():
            if self.sql_generator.validate_query(query):
                sql_results = self.sql_generator.execute_query(query)
                result["sql_results"] = sql_results
            
        return result

    def get_similar_documents(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
        """
        Retrieve similar documents without processing through QA pipeline
        """
        docs = self.retriever.get_relevant_documents(query)
        return [{'content': doc.page_content, 'metadata': doc.metadata} for doc in docs[:k]]

# Example usage
if __name__ == "__main__":
    # Initialize the SQL generator
    sql_gen = SQLGenerator("shopify.db")
    
    # Initialize the RAG system with the SQL generator
    rag = RAGSystem(sql_gen, "apparel.csv")
    
    # Example query that might include SQL
    query = "SELECT * FROM products LIMIT 5"
    results = rag.process_query(query)
    
    # Access different parts of the results
    print("QA Answer:", results["qa_answer"])
    print("Relevant Documents:", results["relevant_docs"])
    print("SQL Results:", results["sql_results"])