path: root/main.py

import asyncio
import requests
import base64
import io
import sys
import time

import gconfig
import input_msg

# from aiogram.methods import SendMessageDraft
from datetime import datetime
from aiogram import Bot, Dispatcher, F
from aiogram.client.default import DefaultBotProperties
from aiogram.types import Update
from aiogram.methods import AnswerGuestQuery, SendRichMessage
from aiogram.types import InlineQueryResultArticle, InputTextMessageContent
from aiogram import BaseMiddleware


from aiogram.enums import ChatType
from aiogram.exceptions import TelegramForbiddenError
from aiogram.types import (
    ChosenInlineResult,
    InlineKeyboardButton,
    InlineKeyboardMarkup,
    InlineQuery,
    InlineQueryResultArticle,
    InputTextMessageContent,
    InputRichMessageContent,
    InputRichMessage,
    BufferedInputFile,
    Message,
)
from callbacks import *
from statistics import *
from colors import *
from helpers import *
from vars import *
from commands import *
from dotenv import load_dotenv
from random import uniform, shuffle, choices

try:
    from config import *
except:
    gconfig.copy_config()
finally:
    from config import *

from supergenerator import *

load_dotenv()

print(f"{BOLD}{GREEN} -- uzbekgpt - привет как дела. {RESET}")


if not API_TOKEN:
    print(f"{RED}{BOLD} !! suka API_TOKEN поставь в .env либо pashol naxxuy {RESET}")
    sys.exit(1)

try:
    bot = Bot(token=API_TOKEN, default=DefaultBotProperties(parse_mode="Markdown"))
    dp = Dispatcher()
    dp.include_router(router)
except Exception as e:
    print(f"{RED}{BOLD} !! вирус у вас тут ошибка: {e}{RESET}")
    sys.exit(1)


def parse_tools(text):
    ...
    return text


@dp.update.outer_middleware()
async def guest_middleware(handler, event, data):
    # print("e: ", event)
    if event.guest_message:
        message = event.guest_message
    elif event.message:
        if event.message.text and any(event.message.text.startswith(cmd) for cmd in commands):
            return await handler(event, data)
        message = event.message
    else:
        return await handler(event, data)

    bot = data.get("bot")
    chat_type = message.sender_chat.type if message.sender_chat else ""
    user_id = message.sender_chat.id if message.sender_chat else message.from_user.id
    config = get_user_config(user_id)
    prompts = SYSTEM_PROMPTS | config["prompts"]
    p = prompts[config["prompt"]]

    if await input_msg.da(bot, message, user_id) != None:
        return

    if p["filer"]:
        if event.guest_message:
            return await g_answer(message, "сорри, выбранный промпт с ответом как файл не поддерживает гостевой режим.✅🤝", title="ошибка врат")

    is_channel = (
        ChatType.CHANNEL == chat_type
        and message.from_user
        and message.from_user.id == 777000
    )
    current_time = time.time()
    stream = config.get("stream")
    user_text = message.text or message.caption or ""
    
    if not is_replied_bot(message, is_channel, user_text, config) and event.message:
        return

    if user_id:
        await asyncio.sleep(uniform(0,0.3))
            
        text, blocked = is_blocked(user_id)
        if not text and blocked:
            return
        if blocked:
            if event.guest_message:
                await g_answer(message, text, title="ошибка врат")
            elif event.message:
                error_message = await message.reply(text)
                await asyncio.sleep(7)
                try:
                    await error_message.delete()
                except Exception as e:
                    print(f"{YELLOW} -- del err - {e}{RESET}")
        
            return

        ensure_user(user_id)
        
        if not user_id in last_command_time:
            last_command_time[user_id] = [current_time, True, [message.message_id]]
        elif len(last_command_time[user_id][2]) == 0:
            last_command_time[user_id] = [current_time, True, [message.message_id]]
        else:
            last_command_time[user_id][0] = current_time
            last_command_time[user_id][1] = True
            last_command_time[user_id][2] += [message.message_id]
            
        await asyncio.sleep(uniform(0,0.2))

        if event.message:
            try:
                await bot.send_chat_action(chat_id=message.chat.id, action="typing")
            except TelegramForbiddenError:
                return
        
        msg = None
        replied = message.reply_to_message or None
        if replied:
            replied_text = message.reply_to_message.text or message.reply_to_message.caption or None
        else:
            replied_text = None
            
        prompt = ""
        images = []

        voice = None

        model = config["model"].split("*", 1)

        if VOSK_API != "":
            if message.voice:
                voice = message.voice
            elif replied and replied.voice:
                voice = replied.voice
        
        process = {
            "text": "wait",
            "photo_a": "wait" if replied and replied.photo and MODELS[model[0]].get("vision", False) == False else None,
            "photo": "wait" if message.photo and MODELS[model[0]][model[1]].get("vision", False) == False else None,
            "stt": "wait" if voice else None,
        }

        if event.message:
            if (config["stream"] == "edit") or (config["stream"] == "native" and user_id == message.chat.id):
                msg = await message.reply(make_graph(message, process))
        else:
            gmsg = await g_answer(message, make_graph(message, process))

        location = None
        if message.location:
            location = message.location
        elif replied and replied.location:
            location = replied.location
            
        if voice:
            if voice.duration > 60:
                da = "не врат не пойдёт такой длинный голосовой сообщение" \
                    ". \nнадо меньше 60 секунд da✅"
                if event.guest_message:
                    await g_answer(message, da)
                else:
                    await message.reply(da)
                return
            buf = BytesIO()
            try:
                brat = await bot.download(voice, destination=buf)
            except Exception as e:
                print(f"{YELLOW} -- voice download err - {e}{RESET}")
            buf.seek(0)

            try:
                if brat:
                    result = await asyncio.to_thread(
                        send_to_vosk,
                        buf.read()
                    )
            except Exception as e:
                process["stt"] = "error"
                if event.message and msg:
                    await msg.edit_text(make_graph(message, process))
                elif event.guest_message and gmsg:
                    await bot.edit_message_text(text=make_graph(message, process), inline_message_id=gmsg.inline_message_id)
                    
                prompt = f"<голосовое сообщение>ошибка подключения к api stt</голосовое сообщение>{prompt}"
                print(f"{RED} -- voice err - {e}{RESET}")
            else:
                prompt = f"<голосовое сообщение>{result.get('text', 'не распознано')}</голосовое сообщение>{prompt}"
            
        
        if message.document:
            file_bytes = io.BytesIO()
            try:
                await bot.download(message.document, destination=file_bytes)
        
                file_bytes.seek(0)
                if message.document.file_size > 9000:
                    file_content = file_bytes.read().decode("utf-8", errors="ignore")[:9000]
                else:
                    file_content = file_bytes.read().decode("utf-8", errors="ignore")
                file_bytes.close()
        
                prompt = f"<файл>{file_content}</файл>{prompt}"
        
            except Exception as e:
                print(f"{YELLOW} -- {e}{RESET}")
            
        if replied and replied.document:
            file_bytes = io.BytesIO()
            try:
                await bot.download(replied.document, destination=file_bytes)
                file_bytes.seek(0)
                if replied.document.file_size > 9000:
                    file_content = file_bytes.read().decode("utf-8", errors="ignore")[:9000]
                else:
                    file_content = file_bytes.read().decode("utf-8", errors="ignore")
                file_bytes.close()
                prompt = f"<файл>{file_content}</файл>{prompt}"
            except Exception as e:
                print(f"{YELLOW} -- {e}{RESET}")
        
        if location:
            prompt += (
                f"<геолокация>\n"
                f"{location.latitude}, {location.longitude}\n"
                f"</геолокация>\n"
            )
            
        if replied and replied.poll:
            poll = replied.poll
            if poll.allows_multiple_answers:
                prompt += "<множественный "
            else:
                prompt += "<одиночный "
            prompt += f"опрос в ответе: {poll.question}>\n"
            for option in poll.options:
                prompt += f"- {option.text}\n"
            prompt += "</опрос в ответе>"

        if message.poll:
            poll = message.poll
            if poll.allows_multiple_answers:
                prompt += "<множественный "
            else:
                prompt += "<одиночный "
            prompt += f"опрос: {poll.question}>\n"
            for option in poll.options:
                prompt += f"- {option.text}\n"
            prompt += "</опрос>"  
    
        if replied and replied.photo and IMAGE_MODEL != "":
            try:
                photo = replied.photo[-1]
                file_bytes = await bot.download(photo)
                image_bytes = file_bytes.read()
                b64_image = base64.b64encode(image_bytes).decode("utf-8")
                images += [{"data": b64_image, "process": "photo_a"}]
            except Exception as e:
                process["photo_a"] = "error"
                if event.message and msg:
                    await msg.edit_text(make_graph(message, process))
                elif event.guest_message and gmsg:
                    await bot.edit_message_text(text=make_graph(message, process), inline_message_id=gmsg.inline_message_id)
                    
                print(f"{RED} -- photo - {e}{RESET}")
        
        if message.photo and IMAGE_MODEL != "":
            try:
                photo = message.photo[-1]
                file_bytes = await bot.download(photo)
                image_bytes = file_bytes.read()
                b64_image = base64.b64encode(image_bytes).decode("utf-8")
                images += [{"data": b64_image, "process": "photo"}]
            except Exception as e:
                process["photo"] = "error"
                if event.message and msg:
                    await msg.edit_text(make_graph(message, process))
                elif event.guest_message and gmsg:
                    await bot.edit_message_text(text=make_graph(message, process), inline_message_id=gmsg.inline_message_id)
                print(f"{RED} -- photo - {e}{RESET}")        
        prompt = f"{prompt}\n{user_text}"
        if replied_text:
            prompt = f"<ответ на>{replied_text}</ответ> {prompt}"
        prompt = prompt[:MAX_PROMPT]

        if config["markdown"] == "" and message.from_user.id != 777000:
            last_command_time[user_id][1] = False
            last_command_time[user_id][2] = []
            alloh_text = """👆*врат важный вопрос*👆если ты видишь сообщение сверху то нажми "новый стиль",   
а если у тебя тг старый и ты не видишь, то нажми "старый стиль". потому что тг обновил
фукция маркdown.✅✅✅"""
            alloh_one_text = """👇*перед моим ответом ответь на очень важный вопрос*👇
ты хочешь использовать новый стиль выделения текста телеграм 12.8+ (6.9.3+ десктоп) или 
хочешь старый для старых версий телеграм и телеграм х? эту настройку можно будет изменить.✅ """
            reply_markup_url = InlineKeyboardMarkup(inline_keyboard=[
                    [InlineKeyboardButton(text="новый стиль", url=f"https://t.me/{BOT_USERNAME}?start=new_md")],
                    [InlineKeyboardButton(text="старый стиль", url=f"https://t.me/{BOT_USERNAME}?start=old_md")],
                ]
            )
            reply_markup = InlineKeyboardMarkup(inline_keyboard=[
                    [InlineKeyboardButton(text="новый стиль", callback_data=f"new_md%{user_id}")],
                    [InlineKeyboardButton(text="старый стиль", callback_data=f"old_md%{user_id}")],
                ]
            )
            if event.guest_message:
                if gmsg:
                    await bot.edit_message_text(text=alloh_one_text, inline_message_id=gmsg.inline_message_id, reply_markup=reply_markup_url)
                else:
                    await g_answer(message, alloh_one_text, reply_markup=reply_markup_url)
            elif event.message:
                await message.reply_rich(InputRichMessage(markdown="# новый крутой стилёк\n*da*\n```\nтестовое сообщение врат\n```"))
                await message.answer(alloh_text, reply_markup=reply_markup)
            return
        
        if event.message:
            try:
                await bot.send_chat_action(chat_id=message.chat.id, action="typing")
            except TelegramForbiddenError:
                return
    
        if stream == "edit":
            if event.message:
                im_model = IMAGE_MODEL.split("*", 1)
                tx_model = config["model"].split("*", 1)
                if msg:
                    message_id = msg.message_id
                else:
                    msg = await message.reply("💭")
                    message_id = msg.message_id
            elif event.guest_message:
                stream = "none"
        else:
            message_id = None
            
        if len(last_command_time[user_id][2]) == 1:
            if is_channel and event.message:
                result = await generate(
                    bot=bot,
                    prompt=prompt,
                    chat_id=message.chat.id,
                    user_id=None,
                    message_id=message_id,
                    images=images,
                )
            elif event.message:
                result = await generate(
                    bot=bot,
                    prompt=prompt,
                    chat_id=message.chat.id,
                    user_id=user_id,
                    message_id=message_id,
                    images=images,
                    msg=msg,
                    process=process,
                    message=message
                )
            else:
                result = await generate(
                    bot=bot,
                    prompt=prompt,
                    chat_id=None,
                    user_id=user_id,
                    message_id=None,
                    images=images,
                    inline_id=gmsg.inline_message_id
                )
        
            result = parse_tools(result)

            if event.guest_message:
                add_one_to("stats/guest")
        else:
            print(f"{YELLOW} -- err in blocking 2 {RESET}")
            return
    else:
        return

    last_command_time[user_id][1] = False
    last_command_time[user_id][2] = []

    if p["filer"]:
        if msg:
            await msg.delete()
        if result == "":
            result = " "
        file = BufferedInputFile(
            file=result.encode("utf-8"),
            filename=result.replace("\n", "").replace(" ", "_")[:5] + ".txt"
        )
        return await message.reply_document(
            document=file
        )

    if POLLS and not message.chat.id in POLLS_BL:
        try:
            if ":" in result and ("poll" in result or "quiz" in result):
                pw = ["[poll]", "poll", "quiz"]
                rl = result.split("\n")
                for i, line in enumerate(rl):
                    if any(line.startswith(w) for w in pw):
                        rl = rl[i:]
                        break
                poll_type = rl[0].split(":")[0].replace("[", "").replace("]", "")
                poll_name = rl[0].split(":")[1]
                if len(rl) > 1:
                    options = rl[1].split(" =| ")
                else:
                    options = ["варианты не указаны"]
                correct_option = options[0]
                correct_option_id = None
                is_anon = True
                allows_multiple_answers = False
                
                if event.guest_message:
                    mid = user_id        
                elif message.chat.type == "private":
                    mid = user_id
                elif message.chat.type in ("group", "supergroup"):
                    mid = message.chat.id
                else:
                    return

                if poll_type == "quiz":
                    shuffle(options)
                    correct_option_id = options.index(correct_option)
                elif poll_type == "poll":
                    poll_type = "regular"
                    allows_multiple_answers = False
                elif poll_type == "multipoll":
                    poll_type = "regular"
                    allows_multiple_answers = True
                else:
                    poll_type = "regular"

                if msg:
                    await msg.delete()
                
                try:
                    result = await answer_poll(
                        data=data,
                        message=message,
                        question=poll_name,
                        options=options,
                        is_anonymous=is_anon,
                        poll_type=poll_type,
                        allows_multiple_answers=allows_multiple_answers,
                        correct_option_id=correct_option_id,
                        user_id=mid
                    )
                except Exception as e:
                    if event.guest_message:
                        await g_answer(message, "гостевой режим не работает с опросами. либо добавь меня " \
                            "в группу, либо напиши мне в личку чтобы я выдал тебе опрос. ✅"
                        )
                    else:
                        print(f"{RED} -- err poll {e} {RESET}")
                    return

                if result and event.guest_message:
                    await g_answer(message, "отправил опрос в личку, так как гостевой режим не работает с " \
                        "опросами✅✅"
                    )
        except Exception as e:
            print(f"{YELLOW} -- poll parsing - {e}{RESET}")

    kwargs = (
        {"rich_message": InputRichMessage(markdown=result)}
        if config["markdown"] == "new"
        else {"text": result}
    )

    try:
        if event.guest_message:
            if len(result) > 4000:
                result = result[:4000] + "...\n\n__сообщение твой большой и он сокращен__"
            if gmsg:
                await bot.edit_message_text(text=result, inline_message_id=gmsg.inline_message_id)
            else:
                await g_answer(message, result, user_id=user_id)
        elif event.message:
            if len(result) > 4000:
                if msg:
                    await msg.delete()
                file = BufferedInputFile(
                    file=result.encode("utf-8"),
                    filename=result.replace("\n", "").replace(" ", "_")[:5] + ".txt"
                )
                return await message.reply_document(
                    document=file
                )
            if msg:
                if stream == "native":
                    await msg.delete()
                    if config["markdown"] == "new":
                        await message.reply_rich(**kwargs)
                    else:
                        await message.reply(**kwargs)
                else:
                    await msg.edit_text(**kwargs)
            else:
                if config["markdown"] == "new":
                    await message.reply_rich(**kwargs)
                else:
                    await message.reply(**kwargs)
    except Exception as e:
        kwargs["parse_mode"] = None
        try:
            if event.guest_message:
                if len(result) > 4000:
                    result = result[:4000] + "...\n\n__сообщение твой большой и он сокращен__"
                if gmsg:
                    await bot.edit_message_text(text=result, inline_message_id=gmsg.inline_message_id, parse_mode="None")
                else:
                    await g_answer(message, result, parse_mode="None")
            elif event.message:
                if msg:
                    if stream == "native":
                        await msg.delete()
                        if config["markdown"] == "new":
                            await message.reply_rich(**kwargs)
                        else:
                            await message.reply(**kwargs)
                    else:
                        await msg.edit_text(**kwargs)
                else:
                    if config["markdown"] == "new":
                        await message.reply_rich(**kwargs)
                    else:
                        await message.reply(**kwargs)
        except Exception as e:
            print(f"{RED} -- send msg - {e}{RESET}")
        else:
            print(f"{YELLOW} -- markdown - {e}{RESET}")
    return


@router.inline_query()
async def inline_handler(inline_query: InlineQuery):
    user_id = inline_query.from_user.id
    ensure_user(user_id)

    config = get_user_config(user_id)

    prompts = SYSTEM_PROMPTS | config["prompts"]
    p = prompts[config["prompt"]]
    
    button = InlineKeyboardButton(text="ало жди 📞", callback_data="pasholnaxxuy")

    if inline_query.query == "clear":
        result = [
            InlineQueryResultArticle(
                id="1",
                title="очистить контекст",
                description="нажми сюда чтоб узбэкгпт прочистил свои мозги✅",
                input_message_content=InputTextMessageContent(
                    message_text="я промыл себе мозги и контекст очищен✅"
                ),
            )
        ]
    else:
        if p["filer"]:
            result = [
                InlineQueryResultArticle(
                    id="1",
                    title="🚫🚫🚫",
                    description=f'отправка как файл в инлайне не поддерживается',
                    input_message_content=InputTextMessageContent(
                        message_text="🚫отправка как файл в инлайне не поддерживается"
                    ),
                )
            ]
        else:
            result = [
                InlineQueryResultArticle(
                    id="1",
                    title="генерация думат",
                    description=f'думать "{inline_query.query}" ',
                    input_message_content=InputTextMessageContent(
                        message_text="я думаю врат... как дела✅"
                    ),
                    reply_markup=InlineKeyboardMarkup(inline_keyboard=[[button]]),
                )
            ]

    await inline_query.answer(results=result, cache_time=0)


@router.chosen_inline_result()
async def chosen_inline_result_handler(chosen_result: ChosenInlineResult):
    inline_message_id = chosen_result.inline_message_id
    user_id = chosen_result.from_user.id

    config = get_user_config(user_id)
    
    prompts = SYSTEM_PROMPTS | config["prompts"]
    p = prompts[config["prompt"]]
    
    if chosen_result.query == "clear":
        user_contexts[user_id] = []
    elif p["filer"]:
        pass
    else:
        current_time = time.time()
        alo, blocked = is_blocked(user_id)
        if not alo and blocked:
            return
        if alo and blocked:
            await bot.edit_message_text(
                text=alo,
                inline_message_id=inline_message_id
            )
            return

        add_one_to("stats/inline")
        
        if not user_id in last_command_time:
            last_command_time[user_id] = [current_time, True, [inline_message_id]]
        elif len(last_command_time[user_id][2]) == 0:
            last_command_time[user_id] = [current_time, True, [inline_message_id]]
        else:
            last_command_time[user_id][0] = current_time
            last_command_time[user_id][1] = True
            last_command_time[user_id][2] += [inline_message_id]
        
        if inline_message_id:
            try:
                result = await generate(
                    prompt=chosen_result.query[:MAX_PROMPT],
                    user_id=chosen_result.from_user.id,
                )
                last_command_time[user_id][1] = False
                last_command_time[user_id][2] = []
            except Exception as e:
                error = parse_response(str(e))
                await bot.edit_message_text(
                    text=error["message"],
                    inline_message_id=inline_message_id,
                    parse_mode=None
                )
                print(f"{RED} -- {error['type']} -- {error['message']}")
                return
            if result:
                try:
                    await bot.edit_message_text(
                        text=result, inline_message_id=inline_message_id
                    )
                except Exception as e:
                    print(f"{RED} -- {e}{RESET}")
                    if "MESSAGE_TOO_LONG" not in str(e):
                        await bot.edit_message_text(
                            text=result,
                            inline_message_id=inline_message_id,
                            parse_mode="None",
                        )
                    else:
                        await bot.edit_message_text(
                            text=result[:800],
                            inline_message_id=inline_message_id,
                            parse_mode="None",
                        )


async def main():
    ensure_path("stats/gens")
    ensure_path("stats/inline")
    ensure_path("stats/guest")
    ensure_path("stats/ben")
    ensure_path("stats/reporn")
    await bot.delete_webhook(drop_pending_updates=True)
    await dp.start_polling(
        bot,
        skip_updates=True,
        allowed_updates = [
            "message",
            "callback_query",
            "guest_message",
            "inline_query",
            "chosen_inline_result",
        ],
    )


if __name__ == "__main__":
    asyncio.run(main())