# flake8: noqa
import asyncio
import json

import aiohttp
import requests

from ocab_core.logger import log

from ...standard.config.config import *
from ...standard.database import *


class YandexGPT:
    token = None
    catalog_id = None
    languages = {
        "ru": "русский язык",
        "en": "английский язык",
        "de": "немецкий язык",
        "uk": "украинский язык",
        "es": "испанский язык",
        "be": "белорусский язык",
    }

    def __init__(self, token, catalog_id):
        self.token = token
        self.catalog_id = catalog_id

    async def async_request(self, url, headers, prompt) -> dict:
        async with aiohttp.ClientSession() as session:
            async with session.post(url, headers=headers, json=prompt) as response:
                return await response.json()

    async def async_token_check(
        self, messages, gpt, max_tokens, stream, temperature, del_msg_id=1
    ):
        url = "https://llm.api.cloud.yandex.net/foundationModels/v1/tokenizeCompletion"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Api-Key {self.token}",
        }
        answer_token = get_yandexgpt_token_for_answer()
        while True:
            try:
                request = {
                    "modelUri": gpt,
                    "completionOptions": {
                        "stream": stream,
                        "temperature": temperature,
                        "maxTokens": max_tokens,
                    },
                    "messages": messages,
                }
                response = await self.async_request(
                    url=url, headers=headers, prompt=request
                )
            except Exception as e:  # TODO: Переделать обработку ошибок
                # print(e)
                await log(f"Error: {e}")

                continue
            if int(len(response["tokens"])) < (max_tokens - answer_token):
                break
            else:
                try:
                    messages.pop(del_msg_id)
                except IndexError:
                    Exception("IndexError: list index out of range")
        return messages

    async def async_yandexgpt_lite(
        self,
        system_prompt,
        input_messages,
        stream=False,
        temperature=0.6,
        max_tokens=8000,
    ):
        url = "https://llm.api.cloud.yandex.net/foundationModels/v1/completion"
        gpt = f"gpt://{self.catalog_id}/yandexgpt-lite/latest"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Api-Key {self.token}",
        }

        messages = [{"role": "system", "text": system_prompt}]
        for message in input_messages:
            messages.append(message)
        messages = await self.async_token_check(messages, gpt, max_tokens)

        prompt = {
            "modelUri": gpt,
            "completionOptions": {
                "stream": stream,
                "temperature": temperature,
                "maxTokens": max_tokens,
            },
            "messages": messages,
        }

        response = requests.post(url, headers=headers, json=prompt).text  # nosec
        return json.loads(response)["result"]["alternatives"][0]["message"]["text"]

    async def async_yandexgpt(
        self,
        system_prompt,
        input_messages,
        stream=False,
        temperature=0.6,
        max_tokens=get_yandexgpt_token_for_request(),
    ):
        url = "https://llm.api.cloud.yandex.net/foundationModels/v1/completion"
        gpt = f"gpt://{self.catalog_id}/yandexgpt/latest"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Api-Key {self.token}",
        }

        messages = []
        messages.append({"role": "system", "text": system_prompt})
        for message in input_messages:
            messages.append(message)

        messages = await self.async_token_check(
            messages, gpt, max_tokens, stream, temperature
        )

        request = {
            "modelUri": gpt,
            "completionOptions": {
                "stream": stream,
                "temperature": temperature,
                "maxTokens": max_tokens,
            },
            "messages": messages,
        }
        response = await self.async_request(
            url=url, headers=headers, prompt=request
        )  # nosec
        return response["result"]["alternatives"][0]["message"]["text"]

    async def async_yandexgpt_translate(self, input_language, output_language, text):
        input_language = self.languages[input_language]
        output_language = self.languages[output_language]

        return await self.async_yandexgpt(
            f"Переведи на {output_language} сохранив оригинальный смысл текста. Верни только результат:",
            [{"role": "user", "text": text}],
            stream=False,
            temperature=0.6,
            max_tokens=8000,
        )

    async def async_yandexgpt_spelling_check(self, input_language, text):
        input_language = self.languages[input_language]

        return await self.async_yandexgpt(
            f"Проверьте орфографию и пунктуацию текста на {input_language}. Верни исправленный текст "
            f"без смысловых искажений:",
            [{"role": "user", "text": text}],
            stream=False,
            temperature=0.6,
            max_tokens=8000,
        )

    async def async_yandexgpt_text_history(
        self, input_messages, stream=False, temperature=0.6, max_tokens=8000
    ):
        url = "https://llm.api.cloud.yandex.net/foundationModels/v1/completion"
        gpt = f"gpt://{self.catalog_id}/summarization/latest"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Api-Key {self.token}",
        }

        messages = []
        for message in input_messages:
            messages.append(message)
        messages = await self.async_token_check(messages, gpt, max_tokens, del_msg_id=0)

        prompt = {
            "modelUri": gpt,
            "completionOptions": {
                "stream": stream,
                "temperature": temperature,
                "maxTokens": max_tokens,
            },
            "messages": messages,
        }

        response = requests.post(url, headers=headers, json=prompt).text  # nosec
        return json.loads(response)["result"]["alternatives"][0]["message"]["text"]

    async def async_yandex_cloud_text_to_speech(
        self, text, voice, emotion, speed, format, quality
    ):
        tts = "tts.api.cloud.yandex.net/speech/v1/tts:synthesize"
        # TODO: Сделать функцию TTS
        return 0

    async def async_yandex_cloud_vision(self, image, features, language):
        # TODO: Сделать функцию Vision
        return 0

    async def collect_messages(self, message_id, chat_id):
        messages = []
        # Собираем цепочку сообщений в формате: [{"role": "user", "text": "<Имя_пользователя>: Привет!"},
        # {"role": "assistant", "text": "Привет!"}]
        while True:
            message = db_api.get_message_text(chat_id, message_id)
            if db_api.get_message_ai_model(chat_id, message_id) != None:
                messages.append({"role": "assistant", "text": message})
            else:
                sender_name = db_api.get_user_name(
                    db_api.get_message_sender_id(chat_id, message_id)
                )
                messages.append({"role": "user", "text": sender_name + ": " + message})
            message_id = db_api.get_answer_to_message_id(chat_id, message_id)
            if message_id is None:
                break
        return list(reversed(messages))

    async def collecting_messages_for_history(
        self, start_message_id, end_message_id, chat_id
    ):
        messages = []
        # Собираем цепочку сообщений в формате: [{"role": "user", "text": "<Имя_пользователя>: Привет!"},
        # {"role": "assistant", "text": "Привет!"}]
        while True:
            message = db_api.get_message_text(chat_id, start_message_id)
            if db_api.get_message_ai_model(chat_id, start_message_id) != None:
                messages.append({"role": "assistant", "text": message})
            else:
                sender_name = db_api.get_user_name(
                    db_api.get_message_sender_id(chat_id, start_message_id)
                )
                messages.append({"role": "user", "text": sender_name + ": " + message})
            start_message_id -= 1
            if start_message_id <= end_message_id:
                break
        return messages.reverse()

    async def yandexgpt_request(
        self,
        message_id=None,
        type="yandexgpt-lite",
        chat_id=None,
        message_id_end=None,
        input_language=None,
        output_language=None,
        text=None,
    ):
        if type == "yandexgpt-lite":
            messages = await self.collect_messages(message_id, chat_id)
            return await self.async_yandexgpt_lite(
                system_prompt=get_yandexgpt_prompt(),
                input_messages=messages,
                stream=False,
                temperature=0.6,
                max_tokens=8000,
            )
        elif type == "yandexgpt":
            # print("yandexgpt_request")
            await log("yandexgpt_request")
            messages = await self.collect_messages(message_id, chat_id)
            return await self.async_yandexgpt(
                system_prompt=get_yandexgpt_prompt(),
                input_messages=messages,
                stream=False,
                temperature=0.6,
                max_tokens=get_yandexgpt_token_for_request(),
            )
        elif type == "yandexgpt-translate":
            return await self.async_yandexgpt_translate(
                input_language,
                output_language,
                text=db_api.get_message_text(chat_id, message_id),
            )
        elif type == "yandexgpt-spelling-check":
            return await self.async_yandexgpt_spelling_check(
                input_language, text=db_api.get_message_text(chat_id, message_id)
            )
        elif type == "yandexgpt-text-history":
            messages = await self.collect_messages_for_history(
                message_id, message_id_end, chat_id
            )
            return await self.async_yandexgpt_text_history(
                messages=messages, stream=False, temperature=0.6, max_tokens=8000
            )
        else:
            return "Ошибка: Неизвестный тип запроса | Error: Unknown request type"