import os
from dotenv import load_dotenv
import logging
import sys
from transformers import pipeline

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    stream=sys.stdout
)
logger = logging.getLogger(__name__)

# Define the RAG prompt template
RAG_PROMPT_TEMPLATE = """
You are an AI assistant analyzing YouTube video transcripts. Your task is to answer questions based on the provided transcript context.

Context from transcript:
{context}

User Question: {question}

Please provide a clear, concise answer based only on the information given in the context. If the context doesn't contain enough information to fully answer the question, acknowledge this in your response.
""".strip()

class RAGSystem:
    def __init__(self, data_processor):
        self.data_processor = data_processor
        self.model = pipeline(
            "text-generation",
            model="google/flan-t5-base",  # Using a smaller model suitable for Spaces
            device=-1  # Use CPU
        )
        logger.info("Initialized RAG system with flan-t5-base model")

    def generate(self, prompt):
        try:
            response = self.model(
                prompt,
                max_length=512,
                min_length=64,
                num_return_sequences=1
            )[0]['generated_text']
            return response
        except Exception as e:
            logger.error(f"Error generating response: {e}")
            return None

    def get_prompt(self, user_query, relevant_docs):
        context = "\n".join([doc['content'] for doc in relevant_docs])
        return RAG_PROMPT_TEMPLATE.format(
            context=context,
            question=user_query
        )

    def query(self, user_query, search_method='hybrid', index_name=None):
        try:
            if not index_name:
                raise ValueError("No index name provided. Please select a video and ensure it has been processed.")

            relevant_docs = self.data_processor.search(
                user_query, 
                num_results=3, 
                method=search_method, 
                index_name=index_name
            )
            
            if not relevant_docs:
                logger.warning("No relevant documents found for the query.")
                return "I couldn't find any relevant information to answer your query.", ""

            prompt = self.get_prompt(user_query, relevant_docs)
            answer = self.generate(prompt)
            
            if not answer:
                return "I encountered an error generating the response.", prompt
                
            return answer, prompt
            
        except Exception as e:
            logger.error(f"An error occurred in the RAG system: {e}")
            return f"An error occurred: {str(e)}", ""
        
    def rewrite_cot(self, query):
        prompt = f"""
        Think through this step by step:
        1. Original query: {query}
        2. What are the key components of this query?
        3. How can we break this down into a clearer question?
        
        Rewritten query:
        """
        response = self.generate(prompt)
        if response:
            return response, prompt
        return query, prompt

    def rewrite_react(self, query):
        prompt = f"""
        Let's approach this step-by-step:
        1. Question: {query}
        2. What information do we need?
        3. What's the best way to structure this query?
        
        Rewritten query:
        """
        response = self.generate(prompt)
        if response:
            return response, prompt
        return query, prompt