import os
from telegram import Update
from telegram.constants import ParseMode
from telegram.ext import ApplicationBuilder, CommandHandler, ContextTypes
from db import get_schedule, init_db

TOKEN = os.getenv('TELEGRAM_TOKEN')


def format_schedule(rows):
    if not rows:
        return 'Розклад на сьогодні недоступний.'
    text = ''
    for num, dep, arr in rows:
        text += f'🚆 {num}: {dep} → {arr}\n'
    return text


async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    await update.message.reply_text(
        'Привіт! Я бот розкладу електричок.\n'
        '/kyiv — Київ→Ніжин\n'
        '/nizhyn — Ніжин→Київ'
    )


async def cmd_kyiv(update: Update, context: ContextTypes.DEFAULT_TYPE):
    rows = get_schedule('kyiv_nizhyn')
    await update.message.reply_text(format_schedule(rows), parse_mode=ParseMode.HTML)


async def cmd_nizhyn(update: Update, context: ContextTypes.DEFAULT_TYPE):
    rows = get_schedule('nizhyn_kyiv')
    await update.message.reply_text(format_schedule(rows), parse_mode=ParseMode.HTML)


def main():
    if not TOKEN:
        print("Будь ласка, встановіть змінну оточення TELEGRAM_TOKEN")
        return

    # Ініціалізуємо БД (створимо таблицю schedules, якщо її ще немає)
    init_db()

    app = ApplicationBuilder().token(TOKEN).build()

    app.add_handler(CommandHandler('start', start))
    app.add_handler(CommandHandler('kyiv', cmd_kyiv))
    app.add_handler(CommandHandler('nizhyn', cmd_nizhyn))

    # Запускаємо бота
    app.run_polling()


if __name__ == '__main__':
    main()