duty-teller/duty_teller/handlers/import_duty_schedule.py

"""Import duty-schedule: /import_duty_schedule (admin only). Two steps: handover time -> JSON file."""

import asyncio

import duty_teller.config as config
from telegram import Update
from telegram.ext import CommandHandler, ContextTypes, MessageHandler, filters

from duty_teller.db.session import session_scope
from duty_teller.handlers.common import is_admin_async
from duty_teller.i18n import get_lang, t
from duty_teller.importers.duty_schedule import (
    DutyScheduleParseError,
    parse_duty_schedule,
)
from duty_teller.services.import_service import run_import
from duty_teller.utils.handover import parse_handover_time


async def import_duty_schedule_cmd(
    update: Update, context: ContextTypes.DEFAULT_TYPE
) -> None:
    """Handle /import_duty_schedule: start two-step import (admin only); asks for handover time."""
    if not update.message or not update.effective_user:
        return
    lang = get_lang(update.effective_user)
    if not await is_admin_async(update.effective_user.id):
        await update.message.reply_text(t(lang, "import.admin_only"))
        return
    context.user_data["awaiting_handover_time"] = True
    await update.message.reply_text(t(lang, "import.handover_format"))


async def handle_handover_time_text(
    update: Update, context: ContextTypes.DEFAULT_TYPE
) -> None:
    """Handle text message when awaiting handover time (e.g. 09:00 Europe/Moscow)."""
    if not update.message or not update.effective_user or not update.message.text:
        return
    if not context.user_data.get("awaiting_handover_time"):
        return
    if not await is_admin_async(update.effective_user.id):
        return
    lang = get_lang(update.effective_user)
    text = update.message.text.strip()
    parsed = parse_handover_time(text)
    if parsed is None:
        await update.message.reply_text(t(lang, "import.parse_time_error"))
        return
    hour_utc, minute_utc = parsed
    context.user_data["handover_utc_time"] = (hour_utc, minute_utc)
    context.user_data["awaiting_handover_time"] = False
    context.user_data["awaiting_duty_schedule_file"] = True
    await update.message.reply_text(t(lang, "import.send_json"))


async def handle_duty_schedule_document(
    update: Update, context: ContextTypes.DEFAULT_TYPE
) -> None:
    """Handle uploaded JSON file: parse duty-schedule and run import."""
    if not update.message or not update.message.document or not update.effective_user:
        return
    if not context.user_data.get("awaiting_duty_schedule_file"):
        return
    lang = get_lang(update.effective_user)
    handover = context.user_data.get("handover_utc_time")
    if not handover:
        return
    if not await is_admin_async(update.effective_user.id):
        return
    if not (update.message.document.file_name or "").lower().endswith(".json"):
        await update.message.reply_text(t(lang, "import.need_json"))
        return

    hour_utc, minute_utc = handover
    file_id = update.message.document.file_id

    file = await context.bot.get_file(file_id)
    raw = bytes(await file.download_as_bytearray())
    try:
        result = parse_duty_schedule(raw)
    except DutyScheduleParseError as e:
        context.user_data.pop("awaiting_duty_schedule_file", None)
        context.user_data.pop("handover_utc_time", None)
        await update.message.reply_text(t(lang, "import.parse_error", error=str(e)))
        return

    def run_import_with_scope():
        with session_scope(config.DATABASE_URL) as session:
            return run_import(session, result, hour_utc, minute_utc)

    loop = asyncio.get_running_loop()
    try:
        num_users, num_duty, num_unavailable, num_vacation = await loop.run_in_executor(
            None, run_import_with_scope
        )
    except Exception as e:
        await update.message.reply_text(t(lang, "import.import_error", error=str(e)))
    else:
        total = num_duty + num_unavailable + num_vacation
        unavailable_suffix = (
            t(lang, "import.done_unavailable", count=str(num_unavailable))
            if num_unavailable
            else ""
        )
        vacation_suffix = (
            t(lang, "import.done_vacation", count=str(num_vacation))
            if num_vacation
            else ""
        )
        await update.message.reply_text(
            t(
                lang,
                "import.done",
                users=str(num_users),
                duties=str(num_duty),
                unavailable=unavailable_suffix,
                vacation=vacation_suffix,
                total=str(total),
            )
        )
    finally:
        context.user_data.pop("awaiting_duty_schedule_file", None)
        context.user_data.pop("handover_utc_time", None)


import_duty_schedule_handler = CommandHandler(
    "import_duty_schedule", import_duty_schedule_cmd
)
handover_time_handler = MessageHandler(
    filters.TEXT & ~filters.COMMAND,
    handle_handover_time_text,
)
duty_schedule_document_handler = MessageHandler(
    filters.Document.FileExtension("json"),
    handle_duty_schedule_document,
)