duty-teller/handlers/import_duty_schedule.py

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

import asyncio
import re
from datetime import date, datetime, timedelta, timezone

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

from db.session import get_session
from db.repository import (
    get_or_create_user_by_full_name,
    delete_duties_in_range,
    insert_duty,
)
from importers.duty_schedule import (
    DutyScheduleParseError,
    DutyScheduleResult,
    parse_duty_schedule,
)

# HH:MM or HH:MM:SS, optional space + timezone (IANA or "UTC")
HANDOVER_TIME_RE = re.compile(
    r"^\s*(\d{1,2}):(\d{2})(?::(\d{2}))?\s*(?:\s+(\S+))?\s*$", re.IGNORECASE
)


def _parse_handover_time(text: str) -> tuple[int, int] | None:
    """Parse handover time string to (hour_utc, minute_utc). Returns None on failure."""
    m = HANDOVER_TIME_RE.match(text)
    if not m:
        return None
    hour = int(m.group(1))
    minute = int(m.group(2))
    # second = m.group(3) ignored
    tz_str = (m.group(4) or "").strip()
    if not tz_str or tz_str.upper() == "UTC":
        return (hour % 24, minute)
    try:
        from zoneinfo import ZoneInfo
    except ImportError:
        try:
            from backports.zoneinfo import ZoneInfo  # type: ignore
        except ImportError:
            return None
    try:
        tz = ZoneInfo(tz_str)
    except Exception:
        return None
    # Build datetime in that tz and convert to UTC
    dt = datetime(2000, 1, 1, hour, minute, 0, tzinfo=tz)
    utc = dt.astimezone(timezone.utc)
    return (utc.hour, utc.minute)


def _duty_to_iso(d: date, hour_utc: int, minute_utc: int) -> str:
    """ISO 8601 with Z for start of duty on date d at given UTC time."""
    dt = datetime(d.year, d.month, d.day, hour_utc, minute_utc, 0, tzinfo=timezone.utc)
    return dt.strftime("%Y-%m-%dT%H:%M:%SZ")


async def import_duty_schedule_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if not update.message or not update.effective_user:
        return
    if not config.is_admin(update.effective_user.username or ""):
        await update.message.reply_text("Доступ только для администраторов.")
        return
    context.user_data["awaiting_handover_time"] = True
    await update.message.reply_text(
        "Укажите время пересменки в формате ЧЧ:ММ и часовой пояс, "
        "например 09:00 Europe/Moscow или 06:00 UTC."
    )


async def handle_handover_time_text(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    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 config.is_admin(update.effective_user.username or ""):
        return
    text = update.message.text.strip()
    parsed = _parse_handover_time(text)
    if parsed is None:
        await update.message.reply_text(
            "Не удалось разобрать время. Укажите, например: 09:00 Europe/Moscow"
        )
        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(
        "Отправьте файл в формате duty-schedule (JSON)."
    )


def _run_import(
    database_url: str,
    result: DutyScheduleResult,
    hour_utc: int,
    minute_utc: int,
) -> tuple[int, int]:
    session = get_session(database_url)
    try:
        from_date_str = result.start_date.isoformat()
        to_date_str = result.end_date.isoformat()
        total_duties = 0
        for full_name, duty_dates in result.entries:
            user = get_or_create_user_by_full_name(session, full_name)
            delete_duties_in_range(session, user.id, from_date_str, to_date_str)
            for d in duty_dates:
                start_at = _duty_to_iso(d, hour_utc, minute_utc)
                d_next = d + timedelta(days=1)
                end_at = _duty_to_iso(d_next, hour_utc, minute_utc)
                insert_duty(session, user.id, start_at, end_at)
                total_duties += 1
        return (len(result.entries), total_duties)
    finally:
        session.close()


async def handle_duty_schedule_document(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    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
    handover = context.user_data.get("handover_utc_time")
    if not handover or not config.is_admin(update.effective_user.username or ""):
        return
    if not (update.message.document.file_name or "").lower().endswith(".json"):
        await update.message.reply_text("Нужен файл с расширением .json")
        return

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

    # Download and parse in async context
    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(f"Ошибка разбора файла: {e}")
        return

    loop = asyncio.get_running_loop()
    try:
        num_users, num_duties = await loop.run_in_executor(
            None,
            lambda: _run_import(config.DATABASE_URL, result, hour_utc, minute_utc),
        )
    except Exception as e:
        await update.message.reply_text(f"Ошибка импорта: {e}")
    else:
        await update.message.reply_text(
            f"Импорт выполнен: {num_users} пользователей, {num_duties} дежурств."
        )
    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,
)