bot/bot/admin_schedule.py

from __future__ import annotations

import html
import json
import secrets
import tempfile
from dataclasses import dataclass
from io import BytesIO
from pathlib import Path
from typing import Any

import requests
from telebot.types import CallbackQuery, InputFile, Message
from telebot.util import quick_markup

import config
from models import logger
from models.bot import bot, send
from models.schedule import clear_schedule_cache


@dataclass
class PendingArchive:
    admin_id: int
    filename: str
    path: Path


PENDING_ARCHIVES: dict[str, PendingArchive] = {}
UPLOAD_CALLBACK_PREFIX = "$schedule_upload:"
CANCEL_CALLBACK_PREFIX = "$schedule_cancel:"


def _is_admin(user_id: int) -> bool:
    return user_id in config.ADMINS


def _api_url(path: str) -> str:
    return f"{config.SCHEDULE_BASE_URL.rstrip('/')}/{path.lstrip('/')}"


def _download_archive(message: Message) -> PendingArchive:
    file_info = bot.get_file(message.document.file_id)
    content = bot.download_file(file_info.file_path)
    filename = message.document.file_name or "schedule.zip"
    with tempfile.NamedTemporaryFile(prefix="zatups-schedule-", suffix=".zip", delete=False) as file:
        file.write(content)
        path = Path(file.name)
    return PendingArchive(message.from_user.id, filename, path)


def _upload_archive(archive: PendingArchive) -> dict[str, Any]:
    with archive.path.open("rb") as file:
        response = requests.post(
            _api_url("schedule/default/upload"),
            files={"file": (archive.filename, file, "application/zip")},
            timeout=(10, 300),
        )
    response.raise_for_status()
    return response.json()


def _clear_database() -> dict[str, Any]:
    response = requests.delete(
        _api_url("admin/database"),
        params={"confirm": "DELETE_ALL_DATA"},
        timeout=30,
    )
    response.raise_for_status()
    return response.json()


def _send_upload_result(admin_id: int, payload: dict[str, Any]):
    content = json.dumps(payload, ensure_ascii=False, indent=2).encode("utf-8")
    caption = "База применена" if payload.get("ok") else "Архив обработан с ошибками"
    return bot.send_document(
        admin_id,
        InputFile(BytesIO(content), "schedule-upload-result.json"),
        caption=caption,
    )


def _format_request_error(error: requests.RequestException) -> str:
    response = error.response
    if response is None:
        return str(error)
    try:
        detail = response.json()
    except ValueError:
        detail = response.text
    return json.dumps(detail, ensure_ascii=False) if isinstance(detail, dict) else str(detail)


@bot.message_handler(
    content_types=["document", "file"],
    func=lambda message: _is_admin(message.from_user.id)
    and bool(message.document)
    and (message.document.file_name or "").lower().endswith(".zip"),
)
def schedule_archive(message: Message):
    try:
        archive = _download_archive(message)
    except Exception as error:
        logger.error("Schedule", f"Archive download failed: {error}")
        return send(message, f"Не удалось скачать архив: <code>{html.escape(str(error))}</code>")

    token = secrets.token_urlsafe(8)
    PENDING_ARCHIVES[token] = archive
    return send(
        message,
        "Вы уверены, что хотите отправить архив в API?\n"
        "Сбросить БД в API можно командой /reset_schedule_db",
        reply_markup=quick_markup(
            {
                "Отправить": {"callback_data": UPLOAD_CALLBACK_PREFIX + token},
                "Отменить": {"callback_data": CANCEL_CALLBACK_PREFIX + token},
            },
            row_width=1,
        ),
    )


@bot.callback_query_handler(func=lambda call: str(call.data).startswith(CANCEL_CALLBACK_PREFIX))
def cancel_schedule_upload(call: CallbackQuery):
    if not _is_admin(call.from_user.id):
        return
    token = str(call.data).removeprefix(CANCEL_CALLBACK_PREFIX)
    archive = PENDING_ARCHIVES.get(token)
    if archive and archive.admin_id == call.from_user.id:
        PENDING_ARCHIVES.pop(token, None)
        archive.path.unlink(missing_ok=True)
    bot.answer_callback_query(call.id, "Отменено")


@bot.callback_query_handler(func=lambda call: str(call.data).startswith(UPLOAD_CALLBACK_PREFIX))
def upload_schedule_archive(call: CallbackQuery):
    if not _is_admin(call.from_user.id):
        return
    token = str(call.data).removeprefix(UPLOAD_CALLBACK_PREFIX)
    archive = PENDING_ARCHIVES.get(token)
    if not archive or archive.admin_id != call.from_user.id:
        return bot.answer_callback_query(call.id, "Архив не найден или уже отправлен")

    bot.answer_callback_query(call.id, "Отправляю архив")
    try:
        payload = _upload_archive(archive)
    except requests.RequestException as error:
        logger.error("Schedule", f"Archive upload failed: {error}")
        return send(
            call.from_user.id,
            f"API отклонил архив: <code>{html.escape(_format_request_error(error))}</code>",
        )
    except (OSError, ValueError) as error:
        logger.error("Schedule", f"Archive upload failed: {error}")
        return send(call.from_user.id, f"Не удалось отправить архив: <code>{html.escape(str(error))}</code>")

    PENDING_ARCHIVES.pop(token, None)
    archive.path.unlink(missing_ok=True)
    clear_schedule_cache()
    return _send_upload_result(call.from_user.id, payload)


@bot.message_handler(commands=["reset_schedule_db"], func=lambda message: _is_admin(message.from_user.id))
def reset_schedule_database(message: Message):
    try:
        _clear_database()
    except requests.RequestException as error:
        logger.error("Schedule", f"Database reset failed: {error}")
        return send(
            message,
            f"Не удалось очистить базу: <code>{html.escape(_format_request_error(error))}</code>",
        )
    except ValueError as error:
        logger.error("Schedule", f"Database reset failed: {error}")
        return send(message, f"API вернул некорректный ответ: <code>{html.escape(str(error))}</code>")

    clear_schedule_cache()
    return send(message, "База очищена")