bot.py

import asyncio
import logging
import os
import ssl
import sys

import aiohttp
import certifi
from aiohttp import web
from aiogram import Bot, Dispatcher
from aiogram.enums import ParseMode
from aiogram.client.default import DefaultBotProperties
from aiogram.webhook.aiohttp_server import SimpleRequestHandler, setup_application

from config import config
from database import db
from handlers.commands import router as commands_router
from handlers.chat import router as chat_router
from handlers.callbacks import router as callbacks_router
from middlewares.owner import OwnerMiddleware
from middlewares.rate_limit import RateLimitMiddleware


# ═══════════════════════════════════════════════════════════════════
# BLOCK 0: SSL-патчи для HF Spaces (certifi + proxy)
# ═══════════════════════════════════════════════════════════════════
custom_ssl = ssl.create_default_context(cafile=certifi.where())
custom_ssl.check_hostname = True
custom_ssl.verify_mode = ssl.CERT_REQUIRED

_orig_tcp_init = aiohttp.TCPConnector.__init__


def _patched_tcp_init(self, *args, **kwargs):
    if kwargs.get("ssl") is not False:
        kwargs["ssl"] = custom_ssl
    _orig_tcp_init(self, *args, **kwargs)


aiohttp.TCPConnector.__init__ = _patched_tcp_init

_orig_session_init = aiohttp.ClientSession.__init__


def _patched_session_init(self, *args, **kwargs):
    kwargs["trust_env"] = True
    _orig_session_init(self, *args, **kwargs)


aiohttp.ClientSession.__init__ = _patched_session_init

_orig_request = aiohttp.ClientSession._request


async def _patched_request(self, method, url, *args, **kwargs):
    proxy_server = os.getenv("TELEGRAM_API_SERVER")
    if proxy_server and "api.telegram.org" in str(url):
        proxy_server = proxy_server.strip().rstrip("/")
        str_url = str(url).replace("https://api.telegram.org", proxy_server)
        logging.info("🔀 Переадресация aiogram через прокси ➡️ %s", str_url)
        url = str_url
    return await _orig_request(self, method, url, *args, **kwargs)


aiohttp.ClientSession._request = _patched_request


# ═══════════════════════════════════════════════════════════════════
# BLOCK 1: Логирование
# ═══════════════════════════════════════════════════════════════════
def setup_logging() -> None:
    level = getattr(logging, config.LOG_LEVEL.upper(), logging.INFO)
    logging.basicConfig(
        level=level,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        stream=sys.stdout,
    )
    # Уменьшаем шум от библиотек
    logging.getLogger("httpx").setLevel(logging.WARNING)
    logging.getLogger("httpcore").setLevel(logging.WARNING)
    logging.getLogger("aiogram").setLevel(logging.INFO)


# ═══════════════════════════════════════════════════════════════════
# BLOCK 2: Инициализация бота и диспетчера
# ═══════════════════════════════════════════════════════════════════
bot = Bot(
    token=config.BOT_TOKEN,
    default=DefaultBotProperties(parse_mode=ParseMode.HTML)
)
dp = Dispatcher()

# Middlewares (порядок важен!)
dp.message.middleware(OwnerMiddleware())
dp.message.middleware(RateLimitMiddleware())

# Logging middleware
@dp.update.middleware()
async def log_updates(handler, event, data):
    logger = logging.getLogger(__name__)
    uid = event.update_id if hasattr(event, 'update_id') else 'N/A'
    logger.info("→ Update received: %s", uid)
    
    if hasattr(event, 'message') and event.message:
        logger.info("→ MSG user_id=%s text=%s", event.message.from_user.id, event.message.text)
    elif hasattr(event, 'callback_query') and event.callback_query:
        logger.info("→ CB user_id=%s data=%s", event.callback_query.from_user.id, event.callback_query.data)
    
    try:
        result = await handler(event, data)
        logger.info("← Handler OK for update %s", uid)
        return result
    except Exception as e:
        logger.error("❌ Handler FAILED for update %s: %s", uid, e, exc_info=True)
        raise
        
# Routers (порядок важен: commands → callbacks → chat)
dp.include_router(commands_router)
dp.include_router(callbacks_router)
dp.include_router(chat_router)


# ═══════════════════════════════════════════════════════════════════
# BLOCK 3: HTTP Handlers
# ═══════════════════════════════════════════════════════════════════
@web.middleware
async def hf_logging_middleware(request, handler):
    return await handler(request)


async def health_check(request: web.Request) -> web.Response:
    """HF Spaces Health Check — обязательно отвечать 200 на '/'."""
    return web.Response(text="🚀 GLM Bot РАБОТАЕТ!")


# ═══════════════════════════════════════════════════════════════════
# BLOCK 4: Lifecycle hooks
# ═══════════════════════════════════════════════════════════════════
async def on_startup(app: web.Application) -> None:
    logger = logging.getLogger(__name__)
    await db.connect()
    logger.info("✅ Database connected")

    space_host = os.getenv("SPACE_HOST", "")
    if space_host:
        full_webhook_link = f"https://{space_host.strip()}{config.WEBHOOK_PATH}"
        for attempt in range(5):
            try:
                await bot.set_webhook(
                    url=full_webhook_link,
                    drop_pending_updates=True,
                    request_timeout=30,
                )
                logger.info("✅ Webhook установлен: %s", full_webhook_link)
                break
            except Exception as e:
                logger.warning(
                    "⚠️ Попытка %d/5 установки webhook: %s", attempt + 1, e
                )
                await asyncio.sleep(5)
    else:
        logger.warning("⚠️ SPACE_HOST не задан, webhook не установлен!")


async def on_shutdown(app: web.Application) -> None:
    logger = logging.getLogger(__name__)
    logger.info("🛑 Shutdown начат...")

    try:
        await bot.delete_webhook(drop_pending_updates=True)
        logger.info("Webhook удалён")
    except Exception as e:
        logger.warning("Ошибка при удалении webhook: %s", e)

    try:
        await dp.storage.close()
        await bot.session.close()
        logger.info("Bot session закрыт")
    except Exception as e:
        logger.warning("Ошибка при закрытии сессии бота: %s", e)

    try:
        await db.disconnect()
        logger.info("Database disconnected")
    except Exception as e:
        logger.warning("Ошибка при отключении БД: %s", e)


# ═══════════════════════════════════════════════════════════════════
# BLOCK 5: Main
# ═══════════════════════════════════════════════════════════════════
def main() -> None:
    setup_logging()
    logger = logging.getLogger(__name__)
    logger.info("🚀 Starting HF Spaces bot...")
    logger.info("Config: model=%s, fallback=%s, streaming=%s, rate_limit=%s",
                config.PRIMARY_MODEL, config.FALLBACK_MODEL,
                config.STREAMING_ENABLED, config.RATE_LIMIT_ENABLED)

    app = web.Application(middlewares=[hf_logging_middleware])
    app.router.add_get("/", health_check)

    webhook_requests_handler = SimpleRequestHandler(dispatcher=dp, bot=bot)
    webhook_requests_handler.register(app, path=config.WEBHOOK_PATH)
    setup_application(app, dp, bot=bot)
    
    app.on_startup.append(on_startup)
    app.on_shutdown.append(on_shutdown)

    port = int(os.environ.get("PORT", 7860))
    logger.info("🚀 Запуск сервера на %s:%d...", "0.0.0.0", port)
    web.run_app(app, host="0.0.0.0", port=port)


if __name__ == "__main__":
    main()