import sys
import click
from yt.wrapper import ypath_join
from datacloud.dev_utils.yt import yt_utils
from datacloud.dev_utils.logging.logger import get_basic_logger
from datacloud.dev_utils.yt.take_part import Interval
from datacloud.audience.lib.segment_table import SegmentTable
from datacloud.audience.lib.score_info import Score
from datacloud.audience.lib.config_table import (
    AudienceConfigTable, AUDIENCE_CONFIG_TABLE_PATH)

logger = get_basic_logger(__name__)


@click.group()
def cli():
    pass


@cli.command('add-audience')
@click.option('--audience-name', required=True)
@click.option('--partner-id', required=True)
@click.option('--score-name', required=True)
@click.option('--low', type=click.FLOAT, required=True)
@click.option('--up', type=click.FLOAT, required=True)
@click.option('--parts', type=click.INT, required=True)
@click.option('--export', is_flag=True, default=False)
@click.option('--export-folder', required=False)
def add_audience(audience_name, partner_id, score_name, low, up, parts, export, export_folder):
    yt_client = yt_utils.get_yt_client()
    table = AudienceConfigTable(yt_client, AUDIENCE_CONFIG_TABLE_PATH)
    score = Score(partner_id=partner_id, name=score_name)
    audience_name = audience_name + '_{:d}-{:d}'.format(int(low*100), int(up*100))
    interval = Interval(low, up)
    additional = {}
    if export:
        additional['is_export'] = True
    if export_folder:
        additional['export_folder_name'] = export_folder
    if click.confirm('Add audience {} : [{}/{}] {}-{} with {} parts?'.format(
            audience_name, partner_id, score_name, low, up, parts)):
        score_path_table = ypath_join('//home/x-products/production/datacloud/models', partner_id, score_name)
        if not yt_client.exists(score_path_table):
            logger.info('Fail: Score not exits: {}'.format(score_path_table))

        table.add_audience(audience_name, interval, parts, score, additional=additional)
        logger.info('Done')


@cli.command('delete-audience')
@click.option('--audience-name', required=True)
def delete_audience(audience_name):
    yt_client = yt_utils.get_yt_client()
    audience_table = AudienceConfigTable(yt_client, AUDIENCE_CONFIG_TABLE_PATH)
    audience = audience_table.get_audience(audience_name)
    if audience is None:
        print('Audience {} not found'.format(audience_name))
        sys.exit(1)
    segment_table = SegmentTable(yt_client)
    segments = list(segment_table.list_with_audience_name(audience_name))
    # segments_table = ...

    if click.confirm('Remove audience {} and {} segments?'.format(audience, len(segments))):
        for s in segments:
            segment_table.delete_segment(audience_name, s['segment_number'])
        audience_table.delete_audience(audience_name)
        print('Done')
    else:
        print('Abort')


@cli.command('list-segments')
@click.option('--audience-name', required=True)
@click.option('--id-only', is_flag=True, default=False)
def list_segments(audience_name, id_only):
    yt_client = yt_utils.get_yt_client()
    table = SegmentTable(yt_client)
    for row in table.list_with_audience_name(audience_name):
        if id_only:
            print(row['segment_id'])
        else:
            print(row)


@cli.command('list-audience')
@click.option('--active', is_flag=True, default=False)
@click.option('--name-only', is_flag=True, default=False)
def list_audience(active, name_only):
    yt_client = yt_utils.get_yt_client()
    table = AudienceConfigTable(yt_client)
    source = None
    if active:
        source = table.list_active
    else:
        source = table.list_records
    for rec in source():
        if name_only:
            print(rec['audience_name'])
        else:
            print(rec)


# TODO: Update to use click
def add_segment_record(yt_client):
    table = SegmentTable(yt_client)
    date = '2019-12-03'
    table.add_segment('cloud_cs_bot_0-10', 1, 11155450, date, date)


if __name__ == '__main__':
    cli()
