import asyncio
import logging
import os
from contextlib import asynccontextmanager
from aiogram import Bot, Dispatcher, Router, F
from aiogram.client.default import DefaultBotProperties
from aiogram.enums import ParseMode
from aiogram.filters import CommandStart, Command
from aiogram.types import Message
from aiogram.fsm.storage.mongo import MongoStorage
from dotenv import load_dotenv
from fastapi import FastAPI
from motor.motor_asyncio import AsyncIOMotorClient
from prometheus_client import Info
from starlette.middleware.cors import CORSMiddleware
from prometheus_fastapi_instrumentator import Instrumentator
# --- ИМПОРТЫ ПРОЕКТА ---
from adapters.google_adapter import GoogleAdapter
from adapters.kling_adapter import KlingAdapter
from adapters.s3_adapter import S3Adapter
from api.service.generation_service import GenerationService
from api.service.album_service import AlbumService
from middlewares.album import AlbumMiddleware
from middlewares.auth import AuthMiddleware
from middlewares.dao import DaoMiddleware
# Репозитории и DAO
from repos.char_repo import CharacterRepo
from repos.user_repo import UsersRepo
from repos.dao import DAO
# Роутеры
from routers.auth_router import router as auth_router
from routers.gen_router import router as gen_router
from routers.char_router import router as char_router
from routers.assets_router import router as assets_router # Роутер бота для ассетов
from api.endpoints.assets_router import router as api_assets_router # Роутер FastAPI
from api.endpoints.character_router import router as api_char_router # Роутер FastAPI
from api.endpoints.generation_router import router as api_gen_router
from api.endpoints.auth import router as api_auth_router
from api.endpoints.admin import router as api_admin_router
from api.endpoints.album_router import router as api_album_router
from api.endpoints.project_router import router as project_api_router
load_dotenv()
logger = logging.getLogger(__name__)
# --- КОНФИГУРАЦИЯ ---
BOT_TOKEN = os.getenv("BOT_TOKEN")
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
MONGO_HOST = os.getenv("MONGO_HOST") # Например: mongodb://localhost:27017
DB_NAME = os.getenv("DB_NAME", "my_bot_db") # Имя базы данных
ADMIN_ID = int(os.getenv("ADMIN_ID", 0))
def setup_logging():
logging.basicConfig(level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(name)s (%(filename)s:%(lineno)d): %(message)s",
force=True)
# --- ИНИЦИАЛИЗАЦИЯ ЗАВИСИМОСТЕЙ ---
bot = Bot(token=BOT_TOKEN, default=DefaultBotProperties(parse_mode=ParseMode.HTML))
# Клиент БД создаем глобально, чтобы он был доступен и боту (Storage), и API
mongo_client = AsyncIOMotorClient(MONGO_HOST)
# Репозитории
# Репозитории
users_repo = UsersRepo(mongo_client)
char_repo = CharacterRepo(mongo_client)
# S3 Adapter
s3_adapter = S3Adapter(
endpoint_url=os.getenv("MINIO_ENDPOINT", "http://31.59.58.220:9000"),
aws_access_key_id=os.getenv("MINIO_ACCESS_KEY", "minioadmin"),
aws_secret_access_key=os.getenv("MINIO_SECRET_KEY", "minioadmin"),
bucket_name=os.getenv("MINIO_BUCKET", "ai-char")
)
dao = DAO(mongo_client, s3_adapter) # Главный DAO для бота
gemini = GoogleAdapter(api_key=GEMINI_API_KEY)
# Kling Adapter (optional, for video generation)
kling_access_key = os.getenv("KLING_ACCESS_KEY", "")
kling_secret_key = os.getenv("KLING_SECRET_KEY", "")
kling_adapter = None
if kling_access_key and kling_secret_key:
kling_adapter = KlingAdapter(access_key=kling_access_key, secret_key=kling_secret_key)
logger.info("Kling adapter initialized")
else:
logger.warning("KLING_ACCESS_KEY / KLING_SECRET_KEY not set — video generation disabled")
generation_service = GenerationService(dao, gemini, s3_adapter, bot, kling_adapter)
album_service = AlbumService(dao)
# Dispatcher
dp = Dispatcher(storage=MongoStorage(mongo_client, db_name=DB_NAME))
# Внедрение зависимостей (глобально для бота)
dp["repo"] = users_repo
dp["admin_id"] = ADMIN_ID
dp["gemini"] = gemini
# --- НАСТРОЙКА РОУТЕРОВ БОТА ---
# 1. Роутеры без мидлварей (например, auth)
dp.include_router(auth_router)
# 2. Основные роутеры
main_router = Router()
dp.include_router(main_router)
dp.include_router(assets_router)
dp.include_router(char_router)
dp.include_router(gen_router)
# --- НАСТРОЙКА MIDDLEWARES БОТА ---
# DaoMiddleware прокидывает объект 'dao' во все хендлеры
dp.update.middleware(DaoMiddleware(dao=dao))
# AuthMiddleware проверяет права доступа
main_router.message.middleware(AuthMiddleware(repo=users_repo, admin_id=ADMIN_ID))
gen_router.message.middleware(AuthMiddleware(repo=users_repo, admin_id=ADMIN_ID))
assets_router.message.middleware(AuthMiddleware(repo=users_repo, admin_id=ADMIN_ID))
# AlbumMiddleware для обработки групп фото
gen_router.message.middleware(AlbumMiddleware(latency=0.8))
# --- LIFESPAN (Запуск FastAPI + Bot) ---
@asynccontextmanager
async def lifespan(app: FastAPI):
# --- STARTUP ---
setup_logging()
print("🚀 Starting up...")
# 1. Настройка DAO для FastAPI
# Используем уже созданный mongo_client
db = mongo_client[DB_NAME]
# Инициализируем DAO для ассетов и кладем в state приложения
# Теперь в эндпоинтах можно делать request.app.state.assets_dao
app.state.mongo_client = mongo_client
app.state.gemini_client = gemini
app.state.bot = bot
app.state.s3_adapter = s3_adapter
app.state.kling_adapter = kling_adapter
app.state.album_service = album_service
app.state.users_repo = users_repo # Добавляем репозиторий в state
print("✅ DB & DAO initialized")
# 2. ЗАПУСК БОТА (в фоне)
# Важно: handle_signals=False, чтобы бот не перехватывал сигналы остановки у uvicorn
# Мы НЕ передаем сюда dao=..., так как он уже подключен через Middleware выше
# polling_task = asyncio.create_task(
# dp.start_polling(bot, handle_signals=False)
# )
# print("🤖 Bot polling started")
yield
# --- SHUTDOWN ---
print("🛑 Shutting down...")
# 3. Остановка бота
polling_task.cancel()
try:
await polling_task
except asyncio.CancelledError:
print("🤖 Bot polling stopped")
# 4. Отключение БД
# Обычно Motor закрывать не обязательно при выходе, но хорошим тоном считается
# mongo_client.close()
print("🛑 DB Connection closed")
# --- НАСТРОЙКА FASTAPI ---
app = FastAPI(title="Assets API", lifespan=lifespan)
# CORS
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
# Подключаем роутеры API
app.include_router(api_auth_router)
app.include_router(api_admin_router)
app.include_router(api_assets_router)
app.include_router(api_char_router)
app.include_router(api_gen_router)
app.include_router(api_album_router)
app.include_router(project_api_router)
# Prometheus Metrics (Instrument after all routers are added)
Instrumentator(
should_group_status_codes=False, # 200/201/204 отдельно (по желанию)
should_ignore_untemplated=False, # НЕ игнорировать "сырые" пути
# should_group_untemplated=False, # (опционально) не схлопывать untemplated в "none"
).instrument(
app,
metric_namespace="ai_bot",
).expose(app, endpoint="/metrics", include_in_schema=False)
app_info = Info("fastapi_app_info", "FastAPI application info")
app_info.info({"app_name": "ai-bot"})
# --- ХЕНДЛЕРЫ БОТА (Main Router) ---
@main_router.message(Command("help"))
async def show_help(message: Message) -> None:
await message.answer("ℹ️ Справка:\n\n"
"📝 Текст: Просто отправь промпт.\n"
"🎨 Фото: /image {промпт} (или прикрепи фото с подписью).\n\n"
"⚠️ Диалоги не сохраняются (каждое сообщение — новый запрос).")
@main_router.message(CommandStart())
async def cmd_start(message: Message):
await message.answer("👋 Привет! Я готов к работе.\n\n"
"Напиши мне, что нужно сгенерировать, или используй /help.")
# --- ЗАПУСК ---
if __name__ == "__main__":
import uvicorn
setup_logging()
async def main():
# Создаем конфигурацию uvicorn вручную
# loop="asyncio" заставляет использовать стандартный цикл
config = uvicorn.Config(app, host="0.0.0.0", port=8090, loop="asyncio", timeout_keep_alive=120, env_file=".env.development")
server = uvicorn.Server(config)
# Запускаем сервер (lifespan запустится внутри)
await server.serve()
try:
# Сами запускаем цикл, контролируя аргументы
asyncio.run(main())
except KeyboardInterrupt:
# Корректно обрабатываем выход
pass