from .yt_yql_interop import ensure_str

import nile.api.v1.extractors as ne
from nile.api.v1.stream import BatchStream
from qb2.api.v1 import (
    extractors as qe,
    filters as qf,
    typing,
)
from yt.wrapper.ypath import ypath_join


def get_tables_with_names(job, yt_dir, name_column, names):
    '''
    Return concatenation of tables from `yt_dir` listed in `names`.
    `name_column` with respective `name` value is added to each record.
    '''
    return job.concat(*[
        job.table(
            ypath_join(yt_dir, name)
        ).project(
            ne.all(),
            qe.const(name_column, name).with_type(typing.Unicode),
        ) for name in names
    ])


def _startswith_filter(column, prefix):  # yt/yql interoperability workaround
    return qf.and_(
        qf.defined(column),
        qf.custom(lambda value: ensure_str(value).startswith(prefix), column),
    )


def _split_by_prefixes(stream, column, prefixes):
    return BatchStream(
        stream.split(
            *[_startswith_filter(column, prefix) for prefix in prefixes],
            multisplit=True,
            strategy='stop_if_true',
        )[:-1]
    )


def put_tables_by_prefixes(stream, yt_dir, column, prefixes, schema=None):
    '''
    Split stream by `column` value `prefixes`. Put resulting streams into respective `yt_dir/prefix`.
    If `column` doesn't have any of `prefixes`, record is omitted.
    '''
    for prefix, result in zip(prefixes, _split_by_prefixes(stream, column, prefixes)):
        result.put(
            ypath_join(yt_dir, prefix),
            schema=schema,
            ensure_optional=False,
        )
