import json
import ai

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

from aiogram.exceptions import TelegramBadRequest


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, system_prompt=None
) -> str:
    if user_id:
        if user_id not in user_contexts:
            user_contexts[user_id] = []

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

        is_auto = model_str == "all*auto" and AUTO_MODEL and AUTO_MODELS
        model = model_str.split("*")

        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

        content = []

        if images and not is_auto:
            provider_name = model[0]
            model_name = model[1]
            settings = MODELS.get(provider_name, {}).get(model_name, {})
            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:
        is_auto = True
        model = ["all", "auto"]
        stream_cf = "none"
        content = []

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

    text = ""

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

    add_one_to("stats/gens")

    if is_auto:
        photo_summary = ""
        if images:
            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"])
                    photo_summary += f"<фото>{image_text}</фото>\n"
                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} : all*auto{RESET}")
                else:
                    if process and msg and message:
                        process[image["process"]] = "done"
                        await msg.edit_text(make_graph(message, process))

        expanded = expand_auto_models()
        last_type = ""
        for entry in expanded:
            if last_type == "da":
                break
            parts = entry.split("*")
            provider_name = parts[0]
            model_name = parts[1]
            provider = PROVIDERS.get(provider_name)
            if not provider:
                continue
            settings = MODELS.get(provider_name, {}).get(model_name, {})
            key = get_key(provider_name, entry)

            auto_content = [{"type": "text", "text": photo_summary + prompt if images and not settings.get("vision") else prompt}]
            if images and settings.get("vision"):
                for image in images:
                    auto_content += [{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image['data']}"}}]

            auto_messages = [{"role": "system", "content": messages[0]["content"]}, {"role": "user", "content": auto_content}] + messages[2:]

            text = await openai_generate_text(
                api=provider,
                api_key=key,
                bot=bot,
                model=model_name,
                settings=settings,
                messages=auto_messages,
                stream=False,
                stream_cf=stream_cf,
                chat_id=chat_id,
                draft_id=chat_id,
                message_id=message_id,
                inline_id=inline_id
            )
            last_type = text["type"]
        if text == "":
            text = {"type": "error", "text": "все модели в автовыборе вернули ошибку"}
    else:
        provider_name = model[0]
        model_name = model[1]
        provider = PROVIDERS.get(provider_name)
        if not provider:
            return "❌произошла ошибка, возможные проблемы:\n" \
                "1. у вас выбрана удалённая модель или провайдер\n" \
                "2. конфиг бота настроен неправильно и надо связаться с тем кто его хостит"
        settings = MODELS.get(provider_name, {}).get(model_name, {})
        key = get_key(provider_name, model_str)

        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

        if process and msg and message:
            process["text"] = "make"
            await msg.edit_text(make_graph(message, process))

        text = await openai_generate_text(
            api=provider,
            api_key=key,
            bot=bot,
            model=model_name,
            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["text"]}})
        user_contexts[user_id] = user_contexts[user_id][-MAX_CONTEXT:]

    return text["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=get_key(model[0], IMAGE_MODEL), model=model[1], messages=messages, stream=False
    ):
        pass
    return completion["choices"][0]["message"]["content"]


def extract_error(resp):
    if isinstance(resp, dict):
        err = resp.get("error")
        if isinstance(err, dict):
            return str(err.get("message") or err)
        elif err:
            return str(err)
        return str(resp.get("message") or "") if "message" in resp else ""
    return ""


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,
    api_key=None
) -> str:
    try:
        extra_body = {
            "think": "false"
        }
        if settings.get("noexb"):
            extra_body = None

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

        if stream:
            message = ""
            stream_error = ""
            t_counter = 0
            counter = 0
            counter_2 = 1
            if not draft_id:
                draft_id = chat_id
            async for chunk in ai.completions(**alo):
                try:
                    content = chunk["choices"][0]["delta"]["content"]
                except:
                    err_text = extract_error(chunk)
                    if err_text:
                        error = parse_response(err_text.lower())
                        print(f"{RED} -- {error['type']} - {err_text}{RESET}")
                        return {
                            "type": "error",
                            "text": error["message"]
                        }
                    if isinstance(chunk, dict) and "choices" not in chunk:
                        stream_error = str(chunk)[:300]
                    elif isinstance(chunk, str) and chunk.strip():
                        stream_error = chunk.strip()[:300]
                    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
            if not text and stream_error:
                error = parse_response(stream_error.lower())
                print(f"{RED} -- {error['type']} - {stream_error}{RESET}")
                return {
                    "type": "error",
                    "text": error["message"]
                }
        else:
            async for completion in ai.completions(**alo):
                pass
            try:
                text = completion["choices"][0]["message"]["content"]
            except Exception as e:
                err_text = extract_error(completion)
                if err_text:
                    error = parse_response(err_text.lower())
                    print(f"{RED} -- {error['type']} - {err_text}{RESET}")
                    return {
                        "type": "error",
                        "text": error["message"]
                    }
                error = parse_response(str(completion).lower())
                print(f"{RED} -- {error['type']} - {completion}{RESET}")
                return {
                    "type": "error",
                    "text": error["message"]
                }
    except Exception as e:
        error = parse_response(str(e).lower())
        print(f"{RED} -- {error['type']} - {e}{RESET}")
        err_str = str(e)
        try:
            status_code = int(err_str.split(" ", 1)[0])
        except ValueError:
            status_code = 0
        if 400 <= status_code < 500:
            debug_alo = {k: v for k, v in alo.items() if k != "api_key"}
            return {
                "type": "error",
                "text": error["message"] + "\n\n" \
                    f"модель: `{model}`\n" \
                    f"статус: `{status_code}`\n" \
                    f"raw: `{err_str[:500]}`\n" \
                    f"запрос:\n```{json.dumps(debug_alo, ensure_ascii=False)[:1200]}\n```"
            }
        return {
            "type": "error",
            "text": error["message"]
        }
    else:
        text, s = clean_tail_repeats(text)
        if s > 0:
            text += f"*... <{s} обрезано>*"
        if text == "":
            return {
                "type": "error",
                "text": "модель вернула пустой ответ"
            }
        return {
            "type": "da",
            "text": text
        }