path: root/main.py

import asyncio
import base64
import io
import sys
import time

import gconfig

# 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
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,
    Message,
)
from callbacks import *
from statistics import *
from colors import *
from helpers import *
from commands import *
from dotenv import load_dotenv
from random import uniform

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)

    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
    is_channel = (
        ChatType.CHANNEL == chat_type
        and message.from_user
        and message.from_user.id == 777000
    )
    current_time = time.time()
    config = get_user_config(user_id)
    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(data, 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.5))

        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 = ""

        if message.document:
            if message.document.file_size > 4096:
                if event.guest_message:
                    await g_answer(data, message, "слишком много данных не хочу отвечать☝️☝️")
                elif event.message:
                    await message.reply("слишком много данных не хочу отвечать☝️☝️")
                return
        
            file_bytes = io.BytesIO()
        
            try:
                await bot.download(message.document, destination=file_bytes)
        
                file_bytes.seek(0)
                file_content = file_bytes.read().decode("utf-8", errors="ignore")
                file_bytes.close()
        
                prompt = f"<файл>{file_content}</файл>"
        
            except UnicodeDecodeError:
                if event.guest_message:
                    await g_answer(data, message, "файл не вошёл брат❌❌")
                elif event.message:
                    await message.reply("файл не вошёл брат❌❌")
                return
        
            except Exception:
                if event.guest_message:
                    await g_answer(data, message, "файл не вошёл ало ошибка❌❌")
                elif event.message:
                    await message.reply("файл не вошёл ало ошибка❌❌")
                return
        
        if replied and replied.document:
            if replied.document.file_size > 4096:
                if event.guest_message:
                    await g_answer(data, message, "слишком много данных не хочу отвечать☝️☝️")
                elif event.message:
                    await message.reply("слишком много данных не хочу отвечать☝️☝️")
                return
        
            file_bytes = io.BytesIO()
        
            try:
                await bot.download(replied.document, destination=file_bytes)
                file_bytes.seek(0)
                file_content = file_bytes.read().decode("utf-8", errors="ignore")
                prompt = f"<файл>{file_content}</файл>{user_text}"
            except UnicodeDecodeError:
                if event.guest_message:
                    await g_answer(data, message, "файл не вошёл❌")
                elif event.message:
                    await message.reply("файл не вошёл❌")
                return
            except Exception as e:
                print(f"{RED} -- {e}{RESET}")
                return
            
        if replied and replied.location:
            loc = replied.location
            prompt += (
                f"<ответ на локацию>\n"
                f"{loc.latitude}, {loc.longitude}\n"
                f"</ответ на локацию>\n"
            )
        
        if message.location:
            loc = message.location
            prompt += (
                f"<локация>\n"
                f"{loc.latitude}, {loc.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")
                image_text = await image2text(image=b64_image)
            except Exception as e:
                image_text = "произошла ошибка при анализе фото"
                print(f"{RED} -- photo - {e}{RESET}")
            prompt = f"{prompt}<фото в ответе>{image_text}</фото в ответе>\n"
        
        if message.photo and IMAGE_MODEL != "":
            try:
                if event.message:
                    if not msg:
                        msg = await message.reply("👀")
                photo = message.photo[-1]
                file_bytes = await bot.download(photo)
                image_bytes = file_bytes.read()
                b64_image = base64.b64encode(image_bytes).decode("utf-8")
                image_text = await image2text(image=b64_image)
            except Exception as e:
                image_text = "произошла ошибка при анализе фото"
                print(f"{RED} -- photo - {e}{RESET}")
            prompt = f"{prompt}<фото>{image_text}</фото>\n"

        if event.message:
            if msg:
                await msg.edit_text("💬")
            if not msg and stream == "edit":
                msg = await message.reply("💬")
        
        prompt = f"{prompt}\n{user_text}"
        if replied_text:
            prompt = f"<ответ на>{replied_text}</ответ> {prompt}"
        prompt = prompt[:MAX_PROMPT]

        if event.message:
            try:
                await bot.send_chat_action(chat_id=message.chat.id, action="typing")
            except TelegramForbiddenError:
                return
    
        if stream == "edit":
            message_id = msg.message_id
        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,
                )
            elif event.message:
                result = await generate(
                    bot=bot,
                    prompt=prompt,
                    chat_id=message.chat.id,
                    user_id=user_id,
                    message_id=message_id,
                )
            else:
                result = await generate(
                    bot=bot,
                    prompt=prompt,
                    chat_id=None,
                    user_id=user_id,
                    message_id=None,
                )
        
            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] = []
    
    try:
        if event.guest_message:
            await g_answer(data, message, result)
        elif event.message:
            if msg:
                if stream == "native":
                    await msg.delete()
                    await message.reply(result)
                else:
                    await msg.edit_text(result)
            else:
                await message.reply(result)
    except Exception as e:
        try:
            if event.guest_message:
                await g_answer(data, message, result, parse_mode="None")
            elif event.message:
                if msg:
                    if stream == "native":
                        await msg.delete()
                        await message.reply(result, parse_mode="None")
                    else:
                        await msg.edit_text(result, parse_mode="None")
                else:
                    await message.reply(result, parse_mode="None")
        except Exception as e:
            print(f"{RED} -- send guest - {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)
    
    button = InlineKeyboardButton(text="ало жди 📞", callback_data="pasholnaxxuy")

    if inline_query.query == "clear":
        result = [
            InlineQueryResultArticle(
                id="1",
                title="очистить контекст",
                description="нажми сюда чтоб узбэкгпт прочистил свои мозги✅",
                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

    add_one_to("stats/inline")
    
    if chosen_result.query == "clear":
        user_contexts[user_id] = []
    else:
        current_time = time.time()
        alo, blocked = is_blocked(user_id)
        if not alo:
            return
        if alo and blocked:
            await bot.edit_message_text(
                text=alo,
                inline_message_id=inline_message_id
            )
            return
        
        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,
                )
                print(f"{RED} -- {error['type']} -- {error['message']}")
                return
            if result:
                if result.startswith("<!block>") or result.endswith("<!block>"):
                    result = result.replace("<!block>", "")
                    last_block_time[user_id] = current_time
                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,
        allowed_updates = [
            "message",
            "callback_query",
            "guest_message",
            "inline_query",
            "chosen_inline_result",
        ],
    )


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