schedule.py

import string
from telegram import InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import (
    CallbackQueryHandler,
    ConversationHandler,
    CommandHandler,
    MessageHandler,
    filters,
)
from pycamp_bot.models import Project, Slot, Pycampista, Vote
from pycamp_bot.commands.auth import admin_needed, get_admins_username
from pycamp_bot.scheduler.db_to_json import export_db_2_json
from pycamp_bot.scheduler.schedule_calculator import export_scheduled_result
from pycamp_bot.utils import escape_markdown, get_slot_weekday_name


DAY_SLOT_TIME = {
    'day':[], # Guarda el codigo del dia ej: ['A','B']
    'slot':[], # Guarda la cantidad de slots del dia iterado ej [5] (se sobreescribe)
    'time':[], # Guarda la hora a la que empieza el dia iterado [15] (se sobreescribe)
}


async def cancel(update, context):
    await context.bot.send_message(
        chat_id=update.message.chat_id,
        text="Has cancelado la carga de slots")
    return ConversationHandler.END


@admin_needed
async def define_slot_days(update, context):
    # TODO: filtrar proyectos por pycamp activo.
    if Slot.select().exists():
        await context.bot.send_message(
            chat_id=update.message.chat_id,
            text="El cronograma ya existe."
        )
        return

    if not Project.select().exists():
        await context.bot.send_message(
            chat_id=update.message.chat_id,
            text="No hay proyectos que cronogramear."
        )
        return

    if not Vote.select().exists():
        await context.bot.send_message(
            chat_id=update.message.chat_id,
            text="Todavia no se realizo la votacion."
        )
        return

    await context.bot.send_message(
        chat_id=update.message.chat_id,
        text="Cuantos dias tiene tu cronograma?"
    )
    return 1


async def define_slot_ammount(update, context):
    global DAY_SLOT_TIME
    text = update.message.text
    if text not in ["1", "2", "3", "4", "5", "6", "7"]:
        await context.bot.send_message(
            chat_id=update.message.chat_id,
            text="mmm eso no parece un numero de dias razonable, de nuevo?"
        )
        return 1

    
    DAY_SLOT_TIME['day'] =list(string.ascii_uppercase[0:int(text)])

    await context.bot.send_message(
        chat_id=update.message.chat_id,
        text="Cuantos slots tiene  tu dia {}".format(DAY_SLOT_TIME['day'][0])
    )
    return 2


async def define_slot_times(update, context):
    text = update.message.text
    day = DAY_SLOT_TIME['day'][0]
    await context.bot.send_message(
        chat_id=update.message.chat_id,
        text="A que hora empieza tu dia {}".format(day)
    )
    DAY_SLOT_TIME['slot'] = [text]
    return 3


async def create_slot(update, context):
    username = update.message.from_user.username
    chat_id = update.message.chat_id
    text = update.message.text

    DAY_SLOT_TIME['time'] = [text]
    slot_amount = DAY_SLOT_TIME['slot'][0]
    times = list(range(int(slot_amount)+1))[1:]
    starting_hour = int(text)

    while len(times) > 0:
        new_slot = Slot(code=str(DAY_SLOT_TIME['day'][0]+str(times[0])))
        new_slot.start = starting_hour

        pycampista = Pycampista.get_or_create(username=username, chat_id=chat_id)[0]
        new_slot.current_wizard = pycampista

        new_slot.save()
        times.pop(0)
        starting_hour += 1

    DAY_SLOT_TIME['day'].pop(0)

    if len(DAY_SLOT_TIME['day']) > 0:
        await context.bot.send_message(
            chat_id=update.message.chat_id,
            text="Cuantos slots tiene tu dia {}".format(DAY_SLOT_TIME['day'][0])
        )
        return 2
    else:
        await context.bot.send_message(
            chat_id=update.message.chat_id,
            text="Genial! Slots Asignados"
        )
        make_schedule(update, context)
        return ConversationHandler.END


async def make_schedule(update, context):
    await context.bot.send_message(
        chat_id=update.message.chat_id,
        text="Generando el Cronograma..."
    )

    data_json = export_db_2_json()
    my_schedule = export_scheduled_result(data_json)

    for relationship in my_schedule:
        slot = Slot.get(Slot.code == relationship[1])
        project = Project.get(Project.name == relationship[0])
        project.slot = slot.id
        project.save()

    await context.bot.send_message(
        chat_id=update.message.chat_id,
        text="Cronograma Generado!"
    )


async def check_day_tab(slot, prev_slot, cronograma):
    def append_day_name():
        cronograma.append(f'*{get_slot_weekday_name(slot.code[0])}:*')

    if prev_slot is None:
        append_day_name()
    elif slot.code[0] != prev_slot.code[0]:
        cronograma.append('')
        append_day_name()


async def show_schedule(update, context):
    slots = Slot.select()
    projects = Project.select()
    cronograma = []

    prev_slot = None

    for i, slot in enumerate(slots):
        await check_day_tab(slot, prev_slot, cronograma)

        for project in projects:
            if project.slot_id == slot.id:
                cronograma.append(f'{slot.start}:00 *{escape_markdown(project.name)}*')
                cronograma.append(f'Owner: @{escape_markdown(project.owner.username)}')

        prev_slot = slot

    await context.bot.send_message(
        chat_id=update.message.chat_id,
        text='\n'.join(cronograma),
        parse_mode='MarkdownV2'
    )


BORRAR_CRONO_PATTERN = "borrarcronograma"


@admin_needed
async def borrar_cronograma(update, context):
    if not Slot.select().exists():
        await context.bot.send_message(
            chat_id=update.message.chat_id,
            text="No hay cronograma para borrar."
        )
        return
    keyboard = [
        [
            InlineKeyboardButton("Sí", callback_data=f"{BORRAR_CRONO_PATTERN}:si"),
            InlineKeyboardButton("No", callback_data=f"{BORRAR_CRONO_PATTERN}:no"),
        ]
    ]
    reply_markup = InlineKeyboardMarkup(keyboard)
    await context.bot.send_message(
        chat_id=update.message.chat_id,
        text="¿Borrar el cronograma? Se quitarán todos los slots y asignaciones.",
        reply_markup=reply_markup,
    )


async def borrar_cronograma_confirm(update, context):
    callback_query = update.callback_query
    await callback_query.answer()
    chat_id = callback_query.message.chat_id
    username = callback_query.from_user.username
    if username not in get_admins_username():
        await context.bot.send_message(
            chat_id=chat_id,
            text="No estas Autorizadx para hacer esta acción",
        )
        return
    if callback_query.data.split(":")[1] == "no":
        await context.bot.send_message(
            chat_id=chat_id,
            text="Operación cancelada.",
        )
        return
    Project.update(slot=None).execute()
    Slot.delete().execute()
    await context.bot.send_message(
        chat_id=chat_id,
        text="Cronograma borrado. Podés volver a usar /cronogramear.",
    )


@admin_needed
async def change_slot(update, context):
    projects = Project.select()
    slots = Slot.select()
    text = update.message.text.split(' ')

    if not len(text) >= 3:
        await context.bot.send_message(
            chat_id=update.message.chat_id,
            text="""El formato de este comando es:
                /cambiar_slot NOMBRE_DEL_PROYECTO NUEVO_SLOT
            ej: /cambiar_slot fades AB
        """
        )
        return

    found = False
    project_name = " ".join(text[1:-1])
    for project in projects:
        if project.name == project_name:
            for slot in slots:
                if slot.code == text[-1]:
                    found = True
                    project.slot = slot.id
                    project.save()
    if found:
        await context.bot.send_message(
            chat_id=update.message.chat_id,
            text="Exito"
        )
    else:
        await context.bot.send_message(
            chat_id=update.message.chat_id,
            text="O el slot o el nombre del proyecto no estan en la db"
        )


load_schedule_handler = ConversationHandler(
    entry_points=[CommandHandler('cronogramear', define_slot_days)],
    states={
        1: [MessageHandler(filters.TEXT, define_slot_ammount)],
        2: [MessageHandler(filters.TEXT, define_slot_times)],
        3: [MessageHandler(filters.TEXT, create_slot)]},
    fallbacks=[CommandHandler('cancel', cancel)])


def set_handlers(application):
    application.add_handler(CommandHandler('cronograma', show_schedule))
    application.add_handler(CommandHandler('borrar_cronograma', borrar_cronograma))
    application.add_handler(
        CallbackQueryHandler(
            borrar_cronograma_confirm,
            pattern=f"{BORRAR_CRONO_PATTERN}:",
        )
    )
    application.add_handler(CommandHandler('cambiar_slot', change_slot))
    application.add_handler(load_schedule_handler)