import io
import json
import time
from queue import Queue
from typing import Dict, List

import numpy as np
import tiktoken
from anyio.from_thread import start_blocking_portal
from django.conf import settings
from langchain.schema import AIMessage, HumanMessage
from openai import OpenAI
from pinecone import QueryResponse

from ayushma.models import ChatMessage
from ayushma.models.chat import Chat
from ayushma.models.document import Document
from ayushma.models.enums import ChatMessageType, ModelType
from ayushma.utils.langchain import LangChainHelper
from ayushma.utils.language_helpers import text_to_speech, translate_text
from core.settings.base import AI_NAME


# https://github.com/openai/openai-python/blob/main/openai/embeddings_utils.py#L65
def cosine_similarity(a, b):
    a = np.array(a).reshape(-1)
    b = np.array(b).reshape(-1)

    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))


def get_embedding(
    text: List[str],
    model: str = "text-embedding-ada-002",
    openai_api_key: str = settings.OPENAI_API_KEY,
) -> List[List[float]]:
    """
    Generates embeddings for the given list of texts using the OpenAI API.

    Args:
        text (List[str]): A list of strings to be embedded.
        model (str, optional): The name of the OpenAI model to use for embedding.
            Defaults to "text-embedding-ada-002".

    Returns:
        A list of embeddings generated by the OpenAI API for the input texts.

    Raises:
        OpenAIError: If there was an error communicating with the OpenAI API.

    Example usage:
        >>> get_embedding(["Hello, world!", "How are you?"])
        [[-0.123, 0.456, 0.789, ...], [0.123, -0.456, 0.789, ...]]

    """
    client = OpenAI(api_key=openai_api_key)

    embedding_args: Dict[str, str | List[str]] = {"input": text}

    if settings.OPENAI_API_TYPE == "azure":
        embedding_args["engine"] = settings.AZURE_EMBEDDING_DEPLOYMENT
    else:
        embedding_args["model"] = model

    res = client.embeddings.create(**embedding_args)

    return [record.embedding for record in res.data]


def get_sanitized_reference(pinecone_references: List[QueryResponse]) -> str:
    """
    Extracts the text from the Pinecone QueryResponse object and sanitizes it.

    Args:
        pinecone_reference (List[QueryResponse]): The similar documents retrieved from
            the Pinecone index.

    Returns:
        A string containing the document id and text from the Pinecone QueryResponse object.

    Example usage:
        >>> get_sanitized_reference([QueryResponse(...), QueryResponse(...)])
        "{'28': 'Hello how are you, I am fine, thank you.', '21': 'How was your day?, Mine was good.'}"
    """
    sanitized_reference = {}

    for reference in pinecone_references:
        for match in reference.matches:
            try:
                document_id = str(match.metadata["document"])
                text = str(match.metadata["text"]).replace("\n", " ") + ","
                if document_id in sanitized_reference:
                    sanitized_reference[document_id] += text
                else:
                    sanitized_reference[document_id] = text
            except Exception as e:
                print(e)
                pass

    return json.dumps(sanitized_reference)


def num_tokens_from_string(string: str, encoding_name: str) -> int:
    """Returns the number of tokens in a text string."""
    encoding = tiktoken.get_encoding(encoding_name)
    num_tokens = len(encoding.encode(string))
    return num_tokens


def split_text(text):
    """Returns one string split into n equal length strings"""
    n = len(text)
    number_of_chars = 8192
    parts = []

    for i in range(0, n, number_of_chars):
        part = text[i : i + number_of_chars]
        parts.append(part)

    return parts


def create_json_response(
    input_text, chat_id, delta, message, stop, error, ayushma_voice
):
    json_data = {
        "chat": str(chat_id),
        "input": input_text,
        "delta": delta,
        "message": message,
        "stop": stop,
        "error": error,
        "ayushma_voice": ayushma_voice,
    }

    return "data: " + json.dumps(json_data) + "\n\n"


def get_reference(text, openai_key, namespace, top_k):
    num_tokens = num_tokens_from_string(text, "cl100k_base")
    embeddings: List[List[List[float]]] = []
    if num_tokens < 8192:
        try:
            embeddings.append(get_embedding(text=[text], openai_api_key=openai_key))
        except Exception as e:
            return Exception(
                e.__str__(),
            )
    else:
        parts = split_text(text)
        for part in parts:
            try:
                embeddings.append(get_embedding(text=[part], openai_api_key=openai_key))
            except Exception as e:
                raise Exception(
                    e.__str__(),
                )
    # find similar embeddings from pinecone index for each embedding
    pinecone_references: List[QueryResponse] = []

    for embedding in embeddings:
        similar: QueryResponse = settings.PINECONE_INDEX_INSTANCE.query(
            vector=embedding,
            top_k=int(top_k),
            namespace=namespace,
            include_metadata=True,
        )
        pinecone_references.append(similar)
    return get_sanitized_reference(pinecone_references=pinecone_references)


def add_reference_documents(chat_message):
    ref_text = "References:"
    chat_text = str(chat_message.original_message)
    ref_start_idx = chat_text.find(ref_text)
    if ref_start_idx == -1:
        return

    try:
        doc_ids = chat_text[ref_start_idx + len(ref_text) :].split(",")
        doc_ids = [doc_id.strip(" .,[]*'\"“”“\n-") for doc_id in doc_ids]
        doc_ids = set([str(doc_id) for doc_id in doc_ids if doc_id != ""])
        for doc_id in doc_ids:
            try:
                doc = Document.objects.get(external_id=doc_id)
                chat_message.reference_documents.add(doc)
            except Document.DoesNotExist:
                pass
    except Exception as e:
        print("Error adding reference documents: ", e)

    chat_message.original_message = chat_text[
        :ref_start_idx
    ].strip()  # Strip to remove empty line at the end \nRefereces:
    chat_message.save()


def handle_post_response(
    chat_response,
    chat,
    match_number,
    user_language,
    temperature,
    stats,
    language,
    tts_engine,
    generate_audio=True,
):
    chat_message: ChatMessage = ChatMessage.objects.create(
        original_message=chat_response,
        chat=chat,
        messageType=ChatMessageType.AYUSHMA,
        top_k=match_number,
        temperature=temperature,
        language=language,
    )
    add_reference_documents(chat_message)
    translated_chat_response = chat_message.original_message
    if user_language != "en-IN":
        stats["response_translation_start_time"] = time.time()
        translated_chat_response = translate_text(
            user_language, chat_message.original_message
        )
    stats["response_translation_end_time"] = time.time()

    ayushma_voice = None
    if generate_audio:
        stats["tts_start_time"] = time.time()
        ayushma_voice = text_to_speech(
            translated_chat_response, user_language, tts_engine
        )
        stats["tts_end_time"] = time.time()

    url = None
    if ayushma_voice:
        stats["upload_start_time"] = time.time()
        chat_message.audio.save(
            f"{chat_message.external_id}.mp3", io.BytesIO(ayushma_voice)
        )
        stats["upload_end_time"] = time.time()

    chat_message.message = translated_chat_response
    chat_message.meta = {
        "translate_start": stats.get("response_translation_start_time"),
        "translate_end": stats.get("response_translation_end_time"),
        "reference_start": stats.get("reference_start_time"),
        "reference_end": stats.get("reference_end_time"),
        "response_start": stats.get("response_start_time"),
        "response_end": stats.get("response_end_time"),
        "tts_start": stats.get("tts_start_time"),
        "tts_end": stats.get("tts_end_time"),
        "upload_start": stats.get("upload_start_time"),
        "upload_end": stats.get("upload_end_time"),
    }
    chat_message.save()
    return translated_chat_response, url, chat_message


def converse(
    english_text,
    local_translated_text,
    openai_key,
    chat,
    match_number,
    user_language,
    temperature,
    stats={},
    stream=True,
    references=None,
    generate_audio=True,
    noonce=None,
    fetch_references=True,
    documents=None,
):
    if not openai_key:
        raise Exception("OpenAI-Key header is required to create a chat or converse")

    english_text = english_text.replace("\n", " ")
    language = user_language.split("-")[0]
    nurse_query = ChatMessage.objects.create(
        message=local_translated_text,
        original_message=english_text,
        chat=chat,
        messageType=ChatMessageType.USER,
        language=language,
        meta={
            "translate_start": stats.get("request_translation_start_time"),
            "translate_end": stats.get("request_translation_end_time"),
        },
        noonce=noonce,
    )

    stats["reference_start_time"] = time.time()

    if references:
        reference = references
    elif fetch_references and chat.project and chat.project.external_id:
        try:
            reference = get_reference(
                english_text, openai_key, str(chat.project.external_id), match_number
            )
        except Exception as e:
            print(e)
            reference = ""
    else:
        reference = ""

    stats["reference_end_time"] = time.time()

    stats["response_start_time"] = time.time()

    prompt = chat.prompt or (chat.project and chat.project.prompt)

    if documents or chat.project.model == ModelType.GPT_4_VISUAL:
        prompt = "Image Capabilities: Enabled\n" + prompt

    # excluding the latest query since it is not a history
    previous_messages = (
        ChatMessage.objects.filter(chat=chat)
        .exclude(id=nurse_query.id)
        .order_by("created_at")
    )
    chat_history = []
    for message in previous_messages:
        if message.messageType == ChatMessageType.USER:
            chat_history.append(HumanMessage(content=f"{message.message}"))
        elif message.messageType == ChatMessageType.AYUSHMA:
            chat_history.append(AIMessage(content=f"Ayushma: {message.message}"))

    tts_engine = chat.project.tts_engine

    if not stream:
        lang_chain_helper = LangChainHelper(
            stream=False,
            openai_api_key=openai_key,
            prompt_template=prompt,
            model=chat.model
            or (chat.project and chat.project.model)
            or ModelType.GPT_3_5,
            temperature=temperature,
        )
        response = lang_chain_helper.get_response(
            english_text, reference, chat_history, documents
        )
        chat_response = response.replace("Ayushma: ", "")
        stats["response_end_time"] = time.time()
        translated_chat_response, url, chat_message = handle_post_response(
            chat_response,
            chat,
            match_number,
            user_language,
            temperature,
            stats,
            language,
            tts_engine,
            generate_audio,
        )

        yield chat_message

    else:
        token_queue = Queue()
        RESPONSE_END = object()
        RESPONSE_ERROR = object()

        try:
            lang_chain_helper = LangChainHelper(
                stream=stream,
                token_queue=token_queue,
                end=RESPONSE_END,
                error=RESPONSE_ERROR,
                openai_api_key=openai_key,
                prompt_template=prompt,
                model=chat.model
                or (chat.project and chat.project.model)
                or ModelType.GPT_3_5,
                temperature=temperature,
            )
            with start_blocking_portal() as portal:
                portal.start_task_soon(
                    lang_chain_helper.get_aresponse,
                    RESPONSE_END,
                    RESPONSE_ERROR,
                    token_queue,
                    english_text,
                    reference,
                    chat_history,
                    documents,
                )
                chat_response = ""
                skip_token = len(f"{AI_NAME}: ")

                while True:
                    if token_queue.empty():
                        continue
                    next_token = token_queue.get(True, timeout=10)
                    if next_token[0] == RESPONSE_ERROR:
                        raise next_token[1]
                    if next_token[0] is RESPONSE_END:
                        stats["response_end_time"] = time.time()
                        (
                            translated_chat_response,
                            url,
                            chat_message,
                        ) = handle_post_response(
                            chat_response,
                            chat,
                            match_number,
                            user_language,
                            temperature,
                            stats,
                            language,
                            tts_engine,
                            generate_audio,
                        )

                        yield create_json_response(
                            local_translated_text,
                            chat.external_id,
                            "",
                            translated_chat_response,
                            True,
                            False,
                            ayushma_voice=url,
                        )
                        break
                    if skip_token > 0:
                        skip_token -= 1
                        continue
                    chat_response += next_token[0]
                    yield create_json_response(
                        local_translated_text,
                        chat.external_id,
                        next_token[0],
                        chat_response,
                        False,
                        False,
                        None,
                    )
        except Exception as e:
            print(e)
            error_text = str(e)
            translated_error_text = error_text
            if user_language != "en-IN":
                translated_error_text = translate_text(user_language, error_text)

            ChatMessage.objects.create(
                message=translated_error_text,
                original_message=error_text,
                chat=chat,
                messageType=ChatMessageType.SYSTEM,
                language=language,
                meta={
                    "translate_start": stats.get("response_translation_start_time"),
                    "translate_end": stats.get("response_translation_end_time"),
                    "reference_start": stats.get("reference_start_time"),
                    "reference_end": stats.get("reference_end_time"),
                    "response_start": stats.get("response_start_time"),
                    "response_end": stats.get("response_end_time"),
                    "tts_start": stats.get("tts_start_time"),
                    "tts_end": stats.get("tts_end_time"),
                    "upload_start": stats.get("upload_start_time"),
                    "upload_end": stats.get("upload_end_time"),
                },
            )
            yield create_json_response(
                local_translated_text, chat.external_id, "", str(e), True, True, None
            )


def converse_thread(
    english_text,
    thread: Chat,
    openai_key,
):
    client = OpenAI(api_key=openai_key)

    if not thread.thread_id:
        thread_id = client.beta.threads.create().id
        thread.thread_id = thread_id
        thread.save()

    client.beta.threads.messages.create(
        thread_id=thread.thread_id, content=english_text, role="user"
    )

    run_id = client.beta.threads.runs.create(
        thread_id=thread.thread_id, assistant_id=thread.project.assistant_id
    ).id

    retires = 60

    while retires > 0:
        retires -= 1
        time.sleep(1)
        status = client.beta.threads.runs.retrieve(
            run_id=run_id, thread_id=thread.thread_id
        ).status
        if status == "completed":
            break
    else:
        raise Exception("Thread timed out")

    response = (
        client.beta.threads.messages.list(thread_id=thread.thread_id)
        .data[0]
        .content[0]
        .text.value
    )

    return response

langchain.schema.AIMessage langchain.schema.HumanMessage
[((1714, 1744), 'openai.OpenAI', 'OpenAI', ([], {'api_key': 'openai_api_key'}), '(api_key=openai_api_key)\n', (1720, 1744), False, 'from openai import OpenAI\n'), ((3257, 3288), 'json.dumps', 'json.dumps', (['sanitized_reference'], {}), '(sanitized_reference)\n', (3267, 3288), False, 'import json\n'), ((3431, 3467), 'tiktoken.get_encoding', 'tiktoken.get_encoding', (['encoding_name'], {}), '(encoding_name)\n', (3452, 3467), False, 'import tiktoken\n'), ((6494, 6673), 'ayushma.models.ChatMessage.objects.create', 'ChatMessage.objects.create', ([], {'original_message': 'chat_response', 'chat': 'chat', 'messageType': 'ChatMessageType.AYUSHMA', 'top_k': 'match_number', 'temperature': 'temperature', 'language': 'language'}), '(original_message=chat_response, chat=chat,\n    messageType=ChatMessageType.AYUSHMA, top_k=match_number, temperature=\n    temperature, language=language)\n', (6520, 6673), False, 'from ayushma.models import ChatMessage\n'), ((7082, 7093), 'time.time', 'time.time', ([], {}), '()\n', (7091, 7093), False, 'import time\n'), ((9284, 9295), 'time.time', 'time.time', ([], {}), '()\n', (9293, 9295), False, 'import time\n'), ((9719, 9730), 'time.time', 'time.time', ([], {}), '()\n', (9728, 9730), False, 'import time\n'), ((9767, 9778), 'time.time', 'time.time', ([], {}), '()\n', (9776, 9778), False, 'import time\n'), ((15876, 15902), 'openai.OpenAI', 'OpenAI', ([], {'api_key': 'openai_key'}), '(api_key=openai_key)\n', (15882, 15902), False, 'from openai import OpenAI\n'), ((860, 872), 'numpy.dot', 'np.dot', (['a', 'b'], {}), '(a, b)\n', (866, 872), True, 'import numpy as np\n'), ((6907, 6918), 'time.time', 'time.time', ([], {}), '()\n', (6916, 6918), False, 'import time\n'), ((6954, 7014), 'ayushma.utils.language_helpers.translate_text', 'translate_text', (['user_language', 'chat_message.original_message'], {}), '(user_language, chat_message.original_message)\n', (6968, 7014), False, 'from ayushma.utils.language_helpers import text_to_speech, translate_text\n'), ((7177, 7188), 'time.time', 'time.time', ([], {}), '()\n', (7186, 7188), False, 'import time\n'), ((7213, 7280), 'ayushma.utils.language_helpers.text_to_speech', 'text_to_speech', (['translated_chat_response', 'user_language', 'tts_engine'], {}), '(translated_chat_response, user_language, tts_engine)\n', (7227, 7280), False, 'from ayushma.utils.language_helpers import text_to_speech, translate_text\n'), ((7335, 7346), 'time.time', 'time.time', ([], {}), '()\n', (7344, 7346), False, 'import time\n'), ((7422, 7433), 'time.time', 'time.time', ([], {}), '()\n', (7431, 7433), False, 'import time\n'), ((7585, 7596), 'time.time', 'time.time', ([], {}), '()\n', (7594, 7596), False, 'import time\n'), ((10603, 10795), 'ayushma.utils.langchain.LangChainHelper', 'LangChainHelper', ([], {'stream': '(False)', 'openai_api_key': 'openai_key', 'prompt_template': 'prompt', 'model': '(chat.model or chat.project and chat.project.model or ModelType.GPT_3_5)', 'temperature': 'temperature'}), '(stream=False, openai_api_key=openai_key, prompt_template=\n    prompt, model=chat.model or chat.project and chat.project.model or\n    ModelType.GPT_3_5, temperature=temperature)\n', (10618, 10795), False, 'from ayushma.utils.langchain import LangChainHelper\n'), ((11101, 11112), 'time.time', 'time.time', ([], {}), '()\n', (11110, 11112), False, 'import time\n'), ((11476, 11483), 'queue.Queue', 'Queue', ([], {}), '()\n', (11481, 11483), False, 'from queue import Queue\n'), ((16366, 16379), 'time.sleep', 'time.sleep', (['(1)'], {}), '(1)\n', (16376, 16379), False, 'import time\n'), ((792, 803), 'numpy.array', 'np.array', (['a'], {}), '(a)\n', (800, 803), True, 'import numpy as np\n'), ((824, 835), 'numpy.array', 'np.array', (['b'], {}), '(b)\n', (832, 835), True, 'import numpy as np\n'), ((876, 893), 'numpy.linalg.norm', 'np.linalg.norm', (['a'], {}), '(a)\n', (890, 893), True, 'import numpy as np\n'), ((896, 913), 'numpy.linalg.norm', 'np.linalg.norm', (['b'], {}), '(b)\n', (910, 913), True, 'import numpy as np\n'), ((4160, 4181), 'json.dumps', 'json.dumps', (['json_data'], {}), '(json_data)\n', (4170, 4181), False, 'import json\n'), ((7514, 7539), 'io.BytesIO', 'io.BytesIO', (['ayushma_voice'], {}), '(ayushma_voice)\n', (7524, 7539), False, 'import io\n'), ((11596, 11858), 'ayushma.utils.langchain.LangChainHelper', 'LangChainHelper', ([], {'stream': 'stream', 'token_queue': 'token_queue', 'end': 'RESPONSE_END', 'error': 'RESPONSE_ERROR', 'openai_api_key': 'openai_key', 'prompt_template': 'prompt', 'model': '(chat.model or chat.project and chat.project.model or ModelType.GPT_3_5)', 'temperature': 'temperature'}), '(stream=stream, token_queue=token_queue, end=RESPONSE_END,\n    error=RESPONSE_ERROR, openai_api_key=openai_key, prompt_template=prompt,\n    model=chat.model or chat.project and chat.project.model or ModelType.\n    GPT_3_5, temperature=temperature)\n', (11611, 11858), False, 'from ayushma.utils.langchain import LangChainHelper\n'), ((5875, 5915), 'ayushma.models.document.Document.objects.get', 'Document.objects.get', ([], {'external_id': 'doc_id'}), '(external_id=doc_id)\n', (5895, 5915), False, 'from ayushma.models.document import Document\n'), ((10326, 10368), 'langchain.schema.HumanMessage', 'HumanMessage', ([], {'content': 'f"""{message.message}"""'}), "(content=f'{message.message}')\n", (10338, 10368), False, 'from langchain.schema import AIMessage, HumanMessage\n'), ((12040, 12063), 'anyio.from_thread.start_blocking_portal', 'start_blocking_portal', ([], {}), '()\n', (12061, 12063), False, 'from anyio.from_thread import start_blocking_portal\n'), ((10066, 10103), 'ayushma.models.ChatMessage.objects.filter', 'ChatMessage.objects.filter', ([], {'chat': 'chat'}), '(chat=chat)\n', (10092, 10103), False, 'from ayushma.models import ChatMessage\n'), ((10463, 10511), 'langchain.schema.AIMessage', 'AIMessage', ([], {'content': 'f"""Ayushma: {message.message}"""'}), "(content=f'Ayushma: {message.message}')\n", (10472, 10511), False, 'from langchain.schema import AIMessage, HumanMessage\n'), ((14598, 14639), 'ayushma.utils.language_helpers.translate_text', 'translate_text', (['user_language', 'error_text'], {}), '(user_language, error_text)\n', (14612, 14639), False, 'from ayushma.utils.language_helpers import text_to_speech, translate_text\n'), ((12883, 12894), 'time.time', 'time.time', ([], {}), '()\n', (12892, 12894), False, 'import time\n')]