import random
import re
import shelve
import time
import ai

from configparser import DEFAULTSECT

from pprint import pprint

from statistics import *
from colors import *
from helpers import *
from config import *
from vars import *

from aiogram.exceptions import TelegramBadRequest, TelegramForbiddenError


async def generate(
    prompt: str, bot=None, user_id=None, chat_id=None, message_id=None, images=None, msg=None, process=None, message=None, inline_id=None
) -> str:
    if user_id:
        if user_id not in user_contexts:
            user_contexts[user_id] = []

        config = get_user_config(user_id)
        if "*" in config.get("model", ""):
            model = config.get("model").split("*", 2)
        else:
            print(f"{RED} -- err model - failed to get prov&model in \"{config.get('model','')}\"{RESET}")
            return "произошла ошибка при парсинге вашей выбранной модели. либо хост бота гандон либо тебе надо сменить модель в настройках!!"

        settings = MODELS.get(model[0]).get(model[1])
        
        stream_cf = config.get("stream", "none")
        system_prompt_name = config.get("prompt", "")
        user_prompts = config.get("prompts", {})
        prompts = SYSTEM_PROMPTS | user_prompts
        
        if system_prompt_name in prompts:
            is_user_prompt = False
            if system_prompt_name in SYSTEM_PROMPTS:
                system_prompt = prompts[system_prompt_name]["text"] + ADD_TO_PROMPT
                if config["markdown"] == "new":
                    system_prompt += NEW_MD
                else:
                    system_prompt += OLD_MD
            else:
                is_user_prompt = True
                system_prompt = prompts[system_prompt_name]["text"]
        else:
            text = "❌произошла ошибка, возможные проблемы:\n" \
                f"1. у вас выбран удалённый сист промпт (`{system_prompt_name}`)\n" \
                f"2. конфиг бота настроен неправильно и надо связаться с тем кто его хостит"
            return text

        if stream_cf == "native" and user_id == chat_id and settings.get("nostream", False) == False:
            stream = True
        elif stream_cf == "edit" and chat_id and message_id and settings.get("nostream", False) == False:
            stream = True
        elif stream_cf == "edit" or stream_cf == "native" and inline_id:
            stream = True
        else:
            stream = False

        content = []
        
        if images:
            if settings.get("vision"):
                for image in images:
                    content += [{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image['data']}"}}]
            else:
                for image in images:
                    try:
                        if process and msg and message:
                            process[image["process"]] = "make"
                            await msg.edit_text(make_graph(message, process))
                        image_text = await image2text(image["data"])
                        prompt = f"<фото>{image_text}</фото>\n{prompt}"
                    except Exception as e:
                        if process and msg and message:
                            process[image["process"]] = "error"
                            await msg.edit_text(make_graph(message, process))
                        print(f"{RED} -- err photo - {e} : {config.get('model', '')}{RESET}")
                    else:
                        if process and msg and message:
                            process[image["process"]] = "done"
                            await msg.edit_text(make_graph(message, process))

        content += [{"type": "text", "text": prompt}]
        
        user_contexts[user_id].append({"role": "user", "content": content})
        user_contexts[user_id] = user_contexts[user_id][-MAX_CONTEXT:]

        context = []
        for i, m in enumerate(user_contexts[user_id]):
            if m.get("role") == "user" and i < len(user_contexts[user_id]) - 1:
                m = {**m, "content": [p for p in m["content"] if p.get("type") != "image_url"]}
            context.append(m)

        messages = [{"role": "system", "content": system_prompt}] + context
    else:
        model = DEFAULT_MODEL.split("*", 2)
        settings = MODELS.get(model[0]).get(model[1])
        stream = False
        stream_cf = "none"
        content = []
        if images:
            if settings.get("vision"):
                for image in images:
                    content += [{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image['data']}"}}]
            else:
                for image in images:
                    try:
                        image_text = await image2text(image["data"])
                        prompt = f"<фото>{image_text}</фото>\n{prompt}"
                    except Exception as e:
                        print(f"{RED} -- err photo - {e} : {model}{RESET}")

        content += [{"type": "text", "text": prompt}]
        messages = [{"role": "system", "content": SYSTEM_PROMPTS[DEFAULT_PROMPT]["text"]+ADD_TO_PROMPT+OLD_MD}] + [
            {"role": "user", "content": content}
        ]

    text = ""

    if str(user_id) in temp_benned:
        return temp_benned[str(user_id)]

    add_one_to("stats/gens")

    try:
        provider = PROVIDERS[model[0]]
    except Exception:
        text = "❌произошла ошибка, возможные проблемы:\n" \
            "1. у вас выбрана удалённая модель или провайдер\n" \
            "2. конфиг бота настроен неправильно и надо связаться с тем кто его хостит"
        return text

    if process and msg and message:
        process["text"] = "make"
        await msg.edit_text(make_graph(message, process))
    
    text = await openai_generate_text(
        api=provider,
        bot=bot,
        model=model[1],
        settings=settings,
        messages=messages,
        stream=stream,
        stream_cf=stream_cf,
        chat_id=chat_id,
        draft_id=chat_id,
        message_id=message_id,
        inline_id=inline_id
    )

    if user_id:
        user_contexts[user_id].append({"role": "assistant", "content": [{"type": "text", "text": text}]})
        user_contexts[user_id] = user_contexts[user_id][-MAX_CONTEXT:]

    return text


async def image2text(image) -> str:
    if IMAGE_MODEL == "":
        return ""
    model = IMAGE_MODEL.split("*", 1)
    api = PROVIDERS[model[0]]
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": IMAGE_PROMPT},
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image}"}},
            ],
        }
    ]
    async for completion in ai.completions(
        url=api["url"], api_key=api["key"], model=model[1], messages=messages, stream=False
    ):
        pass
    return completion["choices"][0]["message"]["content"]


async def openai_generate_text(
    api,
    bot,
    model,
    settings,
    messages,
    stream,
    stream_cf,
    chat_id=None,
    draft_id=None,
    message_id=None,
    inline_id=None
) -> str:
    try:
        extra_body = {
            "think": "false"
        }
        if settings.get("noexb"):
            extra_body = None

        alo = {
            "url": api["url"],
            "api_key": api["key"],
            "model": model,
            "messages": messages,
            "stream": stream,  
            "max_tokens": MAX_TOKENS,
            "extra": extra_body
        }

        if stream:
            message = ""
            t_counter = 0
            counter = 0
            counter_2 = 1
            if not draft_id:
                draft_id = chat_id
            async for chunk in ai.completions(**alo):
                # chunk = json.loads(chunk)
                try:
                    content = chunk["choices"][0]["delta"]["content"]
                except:
                    continue
                if content:
                    message += content
                    counter += len(content)
                    t_counter += 1

                    if counter > MAX_SYMBOLS:
                        break
                    
                    if int(counter / STREAMING_SPLIT_SYMBOLS) == counter_2:
                        counter_2 += 1 
                        display_text, s = clean_tail_repeats(message)
                        if s > 0:
                            display_text += f"*... <{s} обрезано>*"
                        if display_text.count("```") % 2 != 0:
                            display_text += "\n```"

                        if counter < 3950:
                            display_text += "▋"
                            
                            try:
                                if stream_cf == "native" and chat_id and draft_id:
                                    await bot.send_message_draft(
                                        chat_id=chat_id,
                                        draft_id=draft_id,
                                        text=display_text,
                                    )
                                elif stream_cf == "edit" and chat_id and message_id:
                                    await bot.edit_message_text(
                                        chat_id=chat_id,
                                        message_id=message_id,
                                        text=display_text,
                                    )
                                elif stream_cf == "edit" or stream_cf == "native" and inline_id:
                                    await bot.edit_message_text(
                                        inline_message_id=inline_id,
                                        text=display_text,
                                    )
                            except TelegramBadRequest as e:
                                if "can't parse entities" in str(e):
                                    if stream_cf == "native" and chat_id and draft_id:
                                        await bot.send_message_draft(
                                            chat_id=chat_id,
                                            draft_id=draft_id,
                                            text=display_text,
                                            parse_mode=None,
                                        )
                                    elif stream_cf == "edit" and chat_id and message_id:
                                        await bot.edit_message_text(
                                            chat_id=chat_id,
                                            message_id=message_id,
                                            text=display_text,
                                            parse_mode=None,
                                        )
                                    elif stream_cf == "edit" or stream_cf == "native" and inline_id:
                                        await bot.edit_message_text(
                                            inline_message_id=inline_id,
                                            text=display_text,
                                            parse_mode=None,
                                        )
                                    else:
                                        print(f"{YELLOW} -- {e}{RESET}")
                            except Exception as e:
                                print(f"{YELLOW} -- {e}{RESET}")

            text = message
        else:
            async for completion in ai.completions(**alo):
                pass
            text = completion["choices"][0]["message"]["content"]
    except Exception as e:
        error = parse_response(str(e).lower())
        print(f"{RED} -- {error['type']} - {e}{RESET}")
        return error["message"]
    else:
        text, s = clean_tail_repeats(text)
        if s > 0:
            text += f"*... <{s} обрезано>*"
        if text == "":
            text = "(пустой ответ)"
        return text