from datacloud.dev_utils.yt import yt_utils
from datacloud.batch.lib.tables.batch_description_table import (
    BatchDescriptionTable, TaskType, StorageType, SubscriptionType)
from datacloud.batch.lib.tables.batch_status_table import (
    BatchStatusTable)


def add_batch_description(yt_client):
    table = BatchDescriptionTable(yt_client=yt_client)
    table.create_table(force=True)
    table.add(
        task_id='sample-r9-task',
        partner_id='r9',
        task_type=TaskType.SCORING,
        subscription_type=SubscriptionType.BATCH,
        storage=StorageType.S3,
        path='sample-r9-task/',
        data={
            'score_partner_id': 'hcb',
            'score_name': 'hcb_xprod_969',
        }
    )


def add_batch_status(yt_client):
    table = BatchStatusTable(yt_client=yt_client)  # noqa
    table.create_table(force=True)
    # table.add('just-test', 'some-input-path/')


def main():
    yt_client = yt_utils.get_yt_client()
    add_batch_description(yt_client)
    add_batch_status(yt_client)


if __name__ == '__main__':
    main()
