import os
import requests
import azure.functions as func
import openai
from pinecone import Pinecone, ServerlessSpec
import logging
import json

openai_key = os.getenv("openai_key")
openAI_client = openai.OpenAI(api_key=openai_key)
pckey =  os.getenv("pinecone_key")
pc = Pinecone(api_key=pckey)

index_name = os.getenv('index_name') or 'id-pc1-01'

embed_model = "text-embedding-ada-002"

# UTILITY FUNCTIONS 👇
def get_index():
    try:
        if not pc.has_index(index_name):
            pc.create_index(
                name=index_name,
                dimension=1536,
                metric='cosine',
                spec=ServerlessSpec(cloud='aws', region='us-east-1')
            )
        
        return pc.Index(index_name)
    
    except Exception as e:
        logging.error("Error while getting index.")

def embed_text(text):   # convert text to embeddings
    try:
        res = openAI_client.embeddings.create(
            input=[text],
            model=embed_model
        )
        return res.data[0].embedding
    except Exception as e:
        logging.error("Error while embedding text.")

def retrieve_data(query, index, ns):
    """
    Performs similarity search on pinecone and retrieves relevant data.

    Args:
        query (string): Users query.
        index (index): Pinecone index.
        ns (string): Users device id (used as namespace).
    
    Returns:
        list: A list of query-relevant data from pinecone.
    """
    try:
        query_embeds = embed_text(query)
        results = index.query(
            namespace=ns,
            vector=query_embeds,
            top_k=20,
            include_metadata=True
        )

        return [m['metadata']['text'] for m in results['matches']]
    except Exception as e:
        logging.error("Error while retrieving data from pinecone.")

def reply(prompt):  # get reply from LLM
    try:
        response = openAI_client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": "You are a helpful assistant. "
                        "When answering questions, only use information from the context if it directly relates to the question. "
                        "Do not include irrelevant context information. "
                        "If the context contains relevant information, use it to answer the question, even if it may be incorrect. "
                        "If you notice the information from the context is incorrect, you may politely correct it after providing the context-based answer. "
                        "If the context does not contain information relevant to the question, answer based on your own knowledge without mentioning the context. "
                        "Do not say anything about the context lacking information."
                },
                {"role": "user", "content": f"{prompt}"},
            ]
        )

        return response.choices[0].message.content.strip()
    
    except openai.error.OpenAIError as e:
        logging.error(f"Error while getting reply from LLM: {str(e)}")
        return f"An error occurred: {e}"

# MAIN FUNCTION 👇
def main(req: func.HttpRequest) -> func.HttpResponse:   # query LLM based on context similar to the query received from pinecone
    """
    Generates an output based on user query

    Args(req: func.HttpRequest):
        device_id(string): Users device id (used as namespace)
        chat_session_id(Int): Chat session ID on users account
        query(string): Users query that needs to be resolved
    
    Returns(func.HttpResponse):
        string: Solution to users query
    """
    try:
        query = req.get_json()
        device_id = query.get('device_id')
        chat_session_id = query.get('chat_session_id')
        if not device_id:
            logging.error("No user ID found.")
            return func.HttpResponse(json.dumps({"error": "Device ID is required"}), status_code=400, mimetype="application/json")
        if not chat_session_id:
            logging.error("No chat session ID found.")
            return func.HttpResponse(json.dumps({"error": "Chat session ID is required"}), status_code=400, mimetype="application/json")
        
        try:
            index = get_index()
        except Exception as e:
            logging.error(f"{str(e)}")
        query_text = query.get('query')

        if not query_text:
            logging.error("No query text.")
            return func.HttpResponse(json.dumps({"error": "Query text is required"}), status_code=400, mimetype="application/json")

        contexts = retrieve_data(query_text, index, ns=device_id)

        chat_retrieval_URL = os.getenv('chat_retrieval_url') or "https://memory-assist-twt.azurewebsites.net/api/retrieve_chat?"
        payload = {
            'device_id': device_id,
            'chat_session_id': chat_session_id
        }
        chat_retrieval_response = requests.post(chat_retrieval_URL, json=payload)
        chat_retrieval_response.raise_for_status()

        chat_data = chat_retrieval_response.json()
        previous_chat = chat_data['chats']

        prompt = (
            "Context: \n" + "\n".join(contexts) +
            "Previous chat for context: \n" + previous_chat +
            f"\nQuestion: {query_text}"
        )

        answer = reply(prompt)
        payload = {
            'device_id': device_id,
            'chat_session_id': chat_session_id,
            'user_query': query_text,
            'bot_reply': answer
        }

        chat_storage_URL = os.getenv('chat_storage_url') or "https://memory-assist-twt.azurewebsites.net/api/store_chat?"
        chat_storage = requests.post(chat_storage_URL, json=payload)
        chat_storage.raise_for_status()

        return func.HttpResponse(json.dumps({"answer": answer}), status_code=200, mimetype="application/json")
    
    except Exception as e:
        logging.error(f"Error while querying: {str(e)}, (main)")
        return func.HttpResponse(json.dumps({"error": str(e)}), status_code=500, mimetype="application/json")

i have created a chat history management system using mongodb from where I retrieve the whole current chat and use it as context for further generation using an LLM, now using the whole chat session exceeds token limit after like 2-3 messages in a chat session, how do I deal with that?