from yt.wrapper import ypath_join
from datacloud.config.yt import DATACLOUD_CONFIG_FOLDER
from datacloud.dev_utils.logging.logger import get_basic_logger
from datacloud.dev_utils.yt import yt_config_table


logger = get_basic_logger(__name__)
AUDIENCE_CONFIG_TABLE_PATH = ypath_join(
    DATACLOUD_CONFIG_FOLDER, 'audience/audience_config'
)


class AudienceConfigTable(yt_config_table.ConfigTable):
    def __init__(self, yt_client, table_path=AUDIENCE_CONFIG_TABLE_PATH):
        schema = [
            {'name': 'audience_name', 'type': 'string', 'sort_order': 'ascending'},
            {'name': 'interval_start', 'type': 'double'},
            {'name': 'interval_end', 'type': 'double'},
            {'name': 'n_segments', 'type': 'int64'},
            {'name': 'partner_id', 'type': 'string'},
            {'name': 'score_name', 'type': 'string'},
            {'name': 'is_active', 'type': 'boolean'},
            {'name': 'additional', 'type': 'any'},
        ]
        super(AudienceConfigTable, self).__init__(
            table_path, schema, yt_client
        )

    def add_audience(self, audience_name, interval, n_segments, score,
                     is_active=True, additional=None):
        additional = additional or {}
        record = {
            'audience_name': audience_name,
            'interval_start': interval.ratio_from,
            'interval_end': interval.ratio_to,
            'n_segments': n_segments,
            'partner_id': score.partner_id,
            'score_name': score.name,
            'is_active': is_active,
            'additional': additional
        }
        self.insert_records([record])

    def delete_audience(self, audience_name):
        self.remove_records([{'audience_name': audience_name}])

    def list_active(self):
        for record in self.list_records():
            if record.get('is_active') is True:
                yield record

    def get_audience(self, audience_name):
        return self.get_record_by_params({'audience_name': audience_name})
