import os
import requests
import aiohttp
import asyncio
import openai
from telegram import (
    Update,
    InlineKeyboardButton,
    InlineKeyboardMarkup,
    InputMediaPhoto
)
from telegram.ext import (
    ApplicationBuilder,
    CommandHandler,
    ContextTypes,
    CallbackQueryHandler,
    MessageHandler,
    filters
)
from collections import deque
import numpy as np
import logging
from dotenv import load_dotenv

# Firebase Imports
import firebase_admin
from firebase_admin import credentials, firestore

# Math Libraries
import sympy as sp
import matplotlib.pyplot as plt
from io import BytesIO

# LangChain Imports
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains import RetrievalQA
from langchain.llms import OpenAI

# YouTube API
from googleapiclient.discovery import build

# Load Environment Variables
load_dotenv()

# Configuration
TELEGRAM_BOT_TOKEN = os.getenv('TELEGRAM_BOT_TOKEN')
FIREBASE_SERVICE_ACCOUNT_PATH = os.getenv('FIREBASE_SERVICE_ACCOUNT_PATH')
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY')
GOOGLE_CX = os.getenv('GOOGLE_CX')
YOUTUBE_API_KEY = os.getenv('YOUTUBE_API_KEY')

# Validate Configuration
if not all([
    TELEGRAM_BOT_TOKEN,
    FIREBASE_SERVICE_ACCOUNT_PATH,
    OPENAI_API_KEY,
    GOOGLE_API_KEY,
    GOOGLE_CX,
    YOUTUBE_API_KEY
]):
    raise ValueError("One or more API keys or configurations are not set. Please check your environment variables.")

# Initialize OpenAI
openai.api_key = OPENAI_API_KEY

# Initialize Firebase
cred = credentials.Certificate(FIREBASE_SERVICE_ACCOUNT_PATH)
firebase_admin.initialize_app(cred)
db = firestore.client()

# Logging Configuration
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)
logger = logging.getLogger(__name__)

# Initialize LangChain Components
def fetch_jamb_questions():
    """
    Fetches all JAMB questions and answers from Firestore.
    
    Returns:
        list: A list of dictionaries containing questions and answers.
    """
    try:
        jamb_qs_ref = db.collection('jamb_questions')
        docs = jamb_qs_ref.stream()
        jamb_qs = [{'question': doc.to_dict()['question'], 'answer': doc.to_dict()['answer']} for doc in docs]
        return jamb_qs
    except Exception as e:
        logger.error(f"Error fetching JAMB questions: {e}")
        return []

def generate_embeddings(jamb_qs):
    """
    Generates embeddings for each JAMB question.
    
    Args:
        jamb_qs (list): List of Q&A dictionaries.
        
    Returns:
        list: List of embeddings corresponding to each question.
    """
    embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
    questions = [item['question'] for item in jamb_qs]
    question_embeddings = embeddings.embed_documents(questions)
    return question_embeddings

def initialize_vector_store(jamb_qs, question_embeddings):
    """
    Initializes a FAISS vector store with JAMB questions.
    
    Args:
        jamb_qs (list): List of Q&A dictionaries.
        question_embeddings (list): List of embedding vectors.
        
    Returns:
        FAISS: Initialized FAISS vector store.
    """
    return FAISS.from_texts(
        texts=[item['question'] for item in jamb_qs],
        embedding=OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY),
        metadatas=jamb_qs
    )

def create_retrieval_chain(vector_store):
    """
    Creates a RetrievalQA chain using LangChain.
    
    Args:
        vector_store (FAISS): FAISS vector store with JAMB questions.
        
    Returns:
        RetrievalQA: LangChain RetrievalQA object.
    """
    llm = OpenAI(temperature=0.2, openai_api_key=OPENAI_API_KEY)
    qa = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vector_store.as_retriever(),
        return_source_documents=True
    )
    return qa

# Initialize Components
jamb_qs = fetch_jamb_questions()
question_embeddings = generate_embeddings(jamb_qs)
vector_store = initialize_vector_store(jamb_qs, question_embeddings)
retrieval_chain = create_retrieval_chain(vector_store)

# Define Retrieval and QA Functions
async def answer_question(query: str):
    """
    Answers the user's question using LangChain's RetrievalQA.
    
    Args:
        query (str): The user's question.
        
    Returns:
        tuple: (answer string, list of source documents)
    """
    try:
        result = retrieval_chain.run(query)
        answer = result['result']
        sources = result['source_documents']
        # Format sources
        sources_formatted = [{'title': doc.metadata['question'], 'link': 'N/A'} for doc in sources]
        return answer, sources_formatted
    except Exception as e:
        logger.error(f"Error in RetrievalQA: {e}")
        return "🔴 An error occurred while processing your request.", []

# Math Solving Functions
def solve_equation(equation: str):
    """
    Solves a mathematical equation symbolically.
    
    Args:
        equation (str): The equation to solve.
    
    Returns:
        str: The solution to the equation in LaTeX format.
    """
    try:
        x = sp.symbols('x')
        eq = sp.sympify(equation)
        solutions = sp.solve(eq, x)
        if not isinstance(solutions, list):
            solutions = [solutions]
        solutions_latex = ", ".join([sp.latex(s) for s in solutions])
        return solutions_latex
    except Exception as e:
        logger.error(f"Error solving equation '{equation}': {e}")
        return "🔴 Unable to solve the equation. Please check the syntax."

def plot_function(function: str, var='x'):
    """
    Plots a mathematical function and returns the plot as a BytesIO object.
    
    Args:
        function (str): The function to plot (e.g., 'sin(x)').
        var (str, optional): The variable used in the function. Defaults to 'x'.
    
    Returns:
        BytesIO: The image of the plot.
    """
    try:
        x = sp.symbols(var)
        func = sp.sympify(function)
        f = sp.lambdify(x, func, 'numpy')

        # Define range for plotting
        x_vals = np.linspace(-10, 10, 400)
        y_vals = f(x_vals)

        plt.figure(figsize=(6, 4))
        plt.plot(x_vals, y_vals, label=sp.pretty(func))
        plt.title(f'Plot of {sp.pretty(func)}')
        plt.xlabel(var)
        plt.ylabel('f(x)')
        plt.legend()
        plt.grid(True)

        # Save plot to BytesIO object
        buf = BytesIO()
        plt.savefig(buf, format='png')
        plt.close()
        buf.seek(0)
        return buf
    except Exception as e:
        logger.error(f"Error plotting function '{function}': {e}")
        return None

# Firestore Interaction Functions
def store_interaction(user_id: int, command: str, response: str, sources: list = None):
    """
    Stores user interaction in Firestore.
    Each user has a collection, and each interaction is a document.
    """
    try:
        user_ref = db.collection('users').document(str(user_id))
        interactions_ref = user_ref.collection('interactions')
        interactions_ref.add({
            'command': command,
            'response': response,
            'sources': sources,
            'timestamp': firestore.SERVER_TIMESTAMP
        })
        logger.info(f"Stored interaction for user {user_id}")
    except Exception as e:
        logger.error(f"Error storing interaction in Firestore: {e}")

def get_last_interactions(user_id: int, limit=5) -> list:
    """
    Retrieves the last 'limit' interactions for a user from Firestore.
    """
    try:
        user_ref = db.collection('users').document(str(user_id))
        interactions_ref = user_ref.collection('interactions')\
                                     .order_by('timestamp', direction=firestore.Query.DESCENDING)\
                                     .limit(limit)
        interactions = [doc.to_dict() for doc in interactions_ref.stream()]
        return interactions
    except Exception as e:
        logger.error(f"Error retrieving interactions from Firestore: {e}")
        return []

def update_user_status(user_id: int, status: str):
    """
    Updates the user's current status in Firestore.
    """
    try:
        user_ref = db.collection('users').document(str(user_id))
        user_ref.set({
            'status': status,
            'last_updated': firestore.SERVER_TIMESTAMP
        }, merge=True)
        logger.info(f"Updated status for user {user_id} to '{status}'")
    except Exception as e:
        logger.error(f"Error updating user status in Firestore: {e}")

def get_sources_from_interaction(user_id: int):
    """
    Retrieves sources from the last interaction of the user.
    """
    try:
        user_ref = db.collection('users').document(str(user_id))
        interactions_ref = user_ref.collection('interactions').order_by('timestamp', direction=firestore.Query.DESCENDING).limit(1)
        interaction = next(interactions_ref.stream(), None)
        if interaction:
            sources = interaction.to_dict().get('sources', [])
            return sources
        return []
    except Exception as e:
        logger.error(f"Error fetching sources: {e}")
        return []

def get_last_question(user_id: int):
    """
    Retrieves the last question asked by the user.
    """
    try:
        user_ref = db.collection('users').document(str(user_id))
        interactions_ref = user_ref.collection('interactions').order_by('timestamp', direction=firestore.Query.DESCENDING).limit(1)
        interaction = next(interactions_ref.stream(), None)
        if interaction:
            command = interaction.to_dict().get('command', '')
            if '/ask' in command:
                return command.replace('/ask', '').strip()
        return ""
    except Exception as e:
        logger.error(f"Error fetching last question: {e}")
        return ""

def fetch_youtube_videos(query: str, max_results: int = 3) -> list:
    """
    Fetches relevant YouTube videos based on the query.
    """
    try:
        youtube = build('youtube', 'v3', developerKey=YOUTUBE_API_KEY)
        search_response = youtube.search().list(
            q=query,
            part='snippet',
            type='video',
            maxResults=max_results
        ).execute()
        
        videos = []
        for item in search_response.get('items', []):
            video_title = item['snippet']['title']
            video_id = item['id']['videoId']
            video_link = f"https://www.youtube.com/watch?v={video_id}"
            videos.append({'title': video_title, 'link': video_link})
        
        return videos
    except Exception as e:
        logger.error(f"Error fetching YouTube videos: {e}")
        return []

def generate_related_questions(question: str, user_id: int, num_questions: int = 3) -> list:
    """
    Generates related questions using OpenAI's Completion API.
    """
    try:
        prompt = (
            f"Generate {num_questions} related and relevant questions based on the following JAMB question:\n\n"
            f"Question: {question}\n"
            f"Related Questions:"
        )
        response = openai.Completion.create(
            engine="text-davinci-003",
            prompt=prompt,
            max_tokens=150,
            n=1,
            stop=None,
            temperature=0.7,
        )
        related_text = response.choices[0].text.strip()
        related_questions = [q.strip() for q in related_text.split('\n') if q.strip()]
        return related_questions[:num_questions]
    except openai.error.OpenAIError as e:
        logger.error(f"Error generating related questions: {e}")
        return []

# Math Solving Functions
def solve_equation(equation: str):
    """
    Solves a mathematical equation symbolically.
    """
    try:
        x = sp.symbols('x')
        eq = sp.sympify(equation)
        solutions = sp.solve(eq, x)
        if not isinstance(solutions, list):
            solutions = [solutions]
        solutions_latex = ", ".join([sp.latex(s) for s in solutions])
        return solutions_latex
    except Exception as e:
        logger.error(f"Error solving equation '{equation}': {e}")
        return "🔴 Unable to solve the equation. Please check the syntax."

def plot_function(function: str, var='x'):
    """
    Plots a mathematical function and returns the plot as a BytesIO object.
    """
    try:
        x = sp.symbols(var)
        func = sp.sympify(function)
        f = sp.lambdify(x, func, 'numpy')

        # Define range for plotting
        x_vals = np.linspace(-10, 10, 400)
        y_vals = f(x_vals)

        plt.figure(figsize=(6, 4))
        plt.plot(x_vals, y_vals, label=sp.pretty(func))
        plt.title(f'Plot of {sp.pretty(func)}')
        plt.xlabel(var)
        plt.ylabel('f(x)')
        plt.legend()
        plt.grid(True)

        # Save plot to BytesIO object
        buf = BytesIO()
        plt.savefig(buf, format='png')
        plt.close()
        buf.seek(0)
        return buf
    except Exception as e:
        logger.error(f"Error plotting function '{function}': {e}")
        return None

# Bot Command Handlers
async def start(update: Update, context: ContextTypes.DEFAULT.TYPE) -> None:
    """
    Handles the /start command. Welcomes the user and initializes their status.
    """
    user = update.effective_user
    user_id = user.id
    user_name = user.first_name

    # Update user status in Firestore
    update_user_status(user_id, 'started')

    welcome_message = (
        f"👋 Hello, {user_name}! I am your JAMB assistant bot.\n\n"
        "You can use the following commands:\n"
        "/ask [your question] - Get answers to your JAMB questions.\n"
        "/search [your query] - Perform a Google search and get the most relevant link.\n"
        "/solve [math equation] - Solve a mathematical equation.\n"
        "/plot [function] - Plot a mathematical function.\n"
        "/memory - Show the last 5 interactions."
    )
    await update.message.reply_text(welcome_message)

async def ask(update: Update, context: ContextTypes.DEFAULT.TYPE) -> None:
    """
    Handles the /ask command. Retrieves relevant JAMB questions and generates answers using LangChain.
    """
    user = update.effective_user
    user_id = user.id
    query = ' '.join(context.args)
    
    if not query:
        await update.message.reply_text(
            "❌ Please enter a question after /ask.\n\nExample: /ask What is the registration process for JAMB?"
        )
        return

    # Update user status in Firestore
    update_user_status(user_id, 'processing /ask command')

    await update.message.reply_text("🔍 Retrieving the best answer for your question...")

    # Fetch relevant answer using LangChain
    answer, sources = await answer_question(query)

    if answer:
        # Store interaction in Firestore
        store_interaction(user_id, f"/ask {query}", answer, sources)
    
        # Update user status in Firestore
        update_user_status(user_id, 'answered /ask command')
    
        # Send the answer with inline buttons
        keyboard = generate_post_answer_keyboard()
        await update.message.reply_text(answer, parse_mode='Markdown', disable_web_page_preview=False, reply_markup=keyboard)
    else:
        await update.message.reply_text("🔴 Unable to generate an answer. Please try again.")

async def search_command(update: Update, context: ContextTypes.DEFAULT.TYPE) -> None:
    """
    Handles the /search command. Performs a Google search and returns relevant links.
    """
    user = update.effective_user
    user_id = user.id
    query = ' '.join(context.args)
    
    if not query:
        await update.message.reply_text(
            "❌ Please enter a search query after /search.\n\nExample: /search JAMB registration process"
        )
        return

    # Update user status in Firestore
    update_user_status(user_id, 'processing /search command')

    await update.message.reply_text("🔍 Performing Google search...")
    search_results = await google_search(query)

    if not search_results or isinstance(search_results[0], str):
        await update.message.reply_text(search_results[0])
        return

    # Select the most relevant link based on embeddings similarity
    user_embedding = await get_embedding(query)
    if not user_embedding:
        await update.message.reply_text("🔴 Failed to generate embedding for your query.")
        return

    max_similarity = -1.0
    selected_result = None

    for item in search_results:
        result_text = f"{item['title']} {item['snippet']}"
        result_embedding = await get_embedding(result_text)
        if not result_embedding:
            continue
        similarity = compute_cosine_similarity(user_embedding, result_embedding)
        if similarity > max_similarity:
            max_similarity = similarity
            selected_result = item

    if selected_result:
        response_text = (
            f"**{selected_result['title']}**\n\n"
            f"{selected_result['snippet']}\n[Read more]({selected_result['link']})"
        )
        await update.message.reply_text(response_text, parse_mode='Markdown', disable_web_page_preview=False)

        # Store search interaction in Firestore
        store_interaction(user_id, f"/search {query}", response_text, sources=search_results)
    
        # Update user status in Firestore
        update_user_status(user_id, 'answered /search command')
    else:
        await update.message.reply_text("🔎 No relevant results found.")

async def solve_math(update: Update, context: ContextTypes.DEFAULT.TYPE) -> None:
    """
    Handles the /solve command. Solves mathematical equations symbolically and returns the solution.
    """
    user = update.effective_user
    user_id = user.id
    equation = ' '.join(context.args)
    
    if not equation:
        await update.message.reply_text(
            "❌ Please enter an equation after /solve.\n\nExample: /solve x^2 - 4 = 0"
        )
        return

    # Update user status in Firestore
    update_user_status(user_id, 'processing /solve command')

    await update.message.reply_text(f"🔍 Solving the equation: `{equation}`")

    solution = solve_equation(equation)

    if solution:
        # Store interaction in Firestore
        store_interaction(user_id, f"/solve {equation}", solution)
    
        # Update user status in Firestore
        update_user_status(user_id, 'answered /solve command')
    
        # Send solution formatted with LaTeX
        formatted_solution = f"**Solution:**\n\\[{solution}\\]"
        await update.message.reply_text(formatted_solution, parse_mode='Markdown', disable_web_page_preview=False)
    else:
        await update.message.reply_text("🔴 Unable to solve the equation. Please check the syntax.")

async def plot_function_command(update: Update, context: ContextTypes.DEFAULT.TYPE) -> None:
    """
    Handles the /plot command. Plots a mathematical function and sends the image.
    """
    user = update.effective_user
    user_id = user.id
    function = ' '.join(context.args)
    
    if not function:
        await update.message.reply_text(
            "❌ Please enter a function to plot after /plot.\n\nExample: /plot sin(x)"
        )
        return

    # Update user status in Firestore
    update_user_status(user_id, 'processing /plot command')

    await update.message.reply_text(f"📈 Plotting the function: `{function}`")

    plot_image = plot_function(function)

    if plot_image:
        # Store interaction in Firestore
        store_interaction(user_id, f"/plot {function}", f"Plotted function: {function}")
    
        # Update user status in Firestore
        update_user_status(user_id, 'answered /plot command')
    
        # Send the plot image
        await update.message.reply_photo(
            photo=plot_image,
            caption=f"**Plot of {function}**",
            parse_mode='Markdown'
        )
    else:
        await update.message.reply_text("🔴 Unable to plot the function. Please check the syntax.")

async def memory_handler(update: Update, context: ContextTypes.DEFAULT.TYPE) -> None:
    """
    Handles the /memory command. Retrieves and displays the user's recent interactions.
    """
    user = update.effective_user
    user_id = user.id
    if memory:
        memory_text = "\n\n".join([
            f"**Q:** {item['command']}\n**A:** {item['response']}" for item in memory
        ])
        await update.message.reply_text(memory_text, parse_mode='Markdown', disable_web_page_preview=False)
    else:
        # Fetch from Firestore if cache is empty
        interactions = get_last_interactions(user_id)
        if interactions:
            memory_text = "\n\n".join([
                f"**Q:** {item['command']}\n**A:** {item['response']}" for item in interactions
            ])
            await update.message.reply_text(memory_text, parse_mode='Markdown', disable_web_page_preview=False)
        else:
            await update.message.reply_text("🧠 No interactions to display in memory.")

async def button_callback(update: Update, context: ContextTypes.DEFAULT.TYPE) -> None:
    """
    Handles callback queries from inline buttons.
    """
    query = update.callback_query
    await query.answer()

    user_id = query.from_user.id
    data = query.data

    if data == 'sources':
        await send_sources(query, user_id)
    elif data == 'videos':
        await send_videos(query, user_id)
    elif data == 'related_questions':
        await send_related_questions(query, user_id)

def generate_post_answer_keyboard():
    """
    Generates an inline keyboard with buttons for sources, videos, and related questions.
    
    Returns:
        InlineKeyboardMarkup: The inline keyboard markup.
    """
    keyboard = [
        [
            InlineKeyboardButton("📄 Sources", callback_data='sources'),
            InlineKeyboardButton("🎥 Videos", callback_data='videos')
        ],
        [
            InlineKeyboardButton("🔗 Related Questions", callback_data='related_questions')
        ]
    ]
    return InlineKeyboardMarkup(keyboard)

# Callback Query Handlers
def get_sources_from_interaction(user_id: int):
    """
    Retrieves sources from the last interaction of the user.
    """
    try:
        user_ref = db.collection('users').document(str(user_id))
        interactions_ref = user_ref.collection('interactions').order_by('timestamp', direction=firestore.Query.DESCENDING).limit(1)
        interaction = next(interactions_ref.stream(), None)
        if interaction:
            sources = interaction.to_dict().get('sources', [])
            return sources
        return []
    except Exception as e:
        logger.error(f"Error fetching sources: {e}")
        return []

async def send_sources(query, user_id):
    """
    Sends the sources related to the last answer to the user.
    """
    sources = get_sources_from_interaction(user_id)
    if sources:
        sources_text = "📄 **Sources:**\n"
        for idx, source in enumerate(sources, start=1):
            # Assuming 'title' and 'link' are available
            sources_text += f"{idx}. [{source['title']}](N/A)\n"  # Update 'link' if available
        await query.edit_message_text(sources_text, parse_mode='Markdown', disable_web_page_preview=False)
    else:
        await query.edit_message_text("🔍 **No sources found for this answer.**")

async def send_videos(query, user_id):
    """
    Sends relevant YouTube video links related to the last question.
    """
    last_question = get_last_question(user_id)
    if last_question:
        video_links = fetch_youtube_videos(last_question)
        if video_links:
            videos_text = "🎥 **Relevant Videos:**\n"
            for idx, video in enumerate(video_links, start=1):
                videos_text += f"{idx}. [{video['title']}]({video['link']})\n"
            await query.edit_message_text(videos_text, parse_mode='Markdown', disable_web_page_preview=False)
        else:
            await query.edit_message_text("🎥 **No relevant videos found.**")
    else:
        await query.edit_message_text("🔍 **No applicable context to fetch videos.**")

async def send_related_questions(query, user_id):
    """
    Sends related JAMB questions to the user.
    """
    interactions = get_last_interactions(user_id, limit=1)
    if interactions:
        last_command = interactions[0].get('command', '')
        if '/ask' in last_command:
            user_question = last_command.replace('/ask', '').strip()
            related_questions = await generate_related_questions(user_question, user_id)
            if related_questions:
                related_text = "🔗 **Related Questions:**\n"
                for idx, ques in enumerate(related_questions, start=1):
                    related_text += f"{idx}. {ques}\n"
                await query.edit_message_text(related_text, parse_mode='Markdown')
            else:
                await query.edit_message_text("🔗 **No related questions found.**")
        else:
            await query.edit_message_text("❌ **No applicable context to fetch related questions.**")
    else:
        await query.edit_message_text("🔍 **No interactions found.**")

async def echo(update: Update, context: ContextTypes.DEFAULT.TYPE) -> None:
    """
    Echoes unrecognized text messages back to the user.
    """
    user = update.effective_user
    user_id = user.id
    message_text = update.message.text

    # Update user status in Firestore
    update_user_status(user_id, 'received unknown message')

    response = "🤖 Sorry, I didn't understand that command. Please use /start to see available commands."
    await update.message.reply_text(response)

# Main Function
def main():
    """
    Initializes and starts the Telegram bot application.
    """
    application = (
        ApplicationBuilder()
        .token(TELEGRAM_BOT_TOKEN)
        .connect_timeout(60)
        .read_timeout(60)
        .build()
    )
    application.add_handler(CommandHandler("start", start))
    application.add_handler(CommandHandler("ask", ask))
    application.add_handler(CommandHandler("search", search_command))
    application.add_handler(CommandHandler("solve", solve_math))
    application.add_handler(CommandHandler("plot", plot_function_command))
    application.add_handler(CommandHandler("memory", memory_handler))
    application.add_handler(CallbackQueryHandler(button_callback))
    application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, echo))

    # Start polling with error handling
    try:
        logger.info("Bot is starting...")
        application.run_polling()
    except Exception as e:
        logger.error(f"Error running bot: {e}", exc_info=True)

if __name__ == "__main__":
    main()