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__)


DEFAULT_SEGMENT_TABLE_PATH = ypath_join(DATACLOUD_CONFIG_FOLDER, 'audience/segment')


class SegmentTable(yt_config_table.ConfigTable):
    def __init__(self, yt_client, table_path=DEFAULT_SEGMENT_TABLE_PATH):
        schema = [
            {'name': 'audience_name', 'type': 'string', 'sort_order': 'ascending'},
            {'name': 'segment_number', 'type': 'int64', 'sort_order': 'ascending'},
            {'name': 'segment_id', 'type': 'uint64'},
            {'name': 'creation_date', 'type': 'string'},
            {'name': 'update_date', 'type': 'string'},
            {'name': 'additional', 'type': 'any'},
        ]
        super(SegmentTable, self).__init__(
            table_path, schema, yt_client
        )

    def add_segment(self, audience_name, segment_number,
                    segment_id, creation_date,
                    update_date, additional=None):
        additional = additional or {}
        record = {
            'audience_name': audience_name,
            'segment_number': segment_number,
            'segment_id': segment_id,
            'creation_date': creation_date,
            'update_date': update_date,
            'additional': additional,
        }
        self.insert_records([record])

    def list_segments(self):
        for record in self.list_records():
            yield record

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

    def list_with_audience_name(self, audience_name):
        for rec in self.get_records_by_params({'audience_name': audience_name}):
            yield rec

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