import click
from datacloud.dev_utils.logging.logger import get_basic_logger
from datacloud.dev_utils.yt import yt_utils
from datacloud.dev_utils.status_db.db import StatusDB
from datacloud.dev_utils.status_db.task import Status


logger = get_basic_logger(__name__)


@click.command()
@click.argument('program')
@click.argument('key')
@click.argument('status', type=click.Choice(['READY', 'DONE', 'SKIPPED']))
def main(program, key, status):
    yt_client = yt_utils.get_yt_client()
    db = StatusDB(yt_client=yt_client)
    task = db.get_task(program, key)
    if status == 'READY':
        set_status = Status.READY
    elif status == 'DONE':
        set_status = Status.DONE
    elif status == 'SKIPPED':
        set_status = Status.SKIPPED

    logger.info('Task: {}'.format(task))
    task.make_done(new_status=set_status)
    db.add_task(task)
    logger.info('Status changed to: {}'.format(set_status))


if __name__ == '__main__':
    main()
