from .base import Column


class TupleColumn(Column):
    py_types = (list, tuple)

    def __init__(self, nested_columns, **kwargs):
        self.nested_columns = nested_columns
        super().__init__(**kwargs)
        self.null_value = tuple(x.null_value for x in nested_columns)

    async def write_data(
        self,
        items,
    ):
        if not items:
            return
        items = self.prepare_items(items)
        items = list(zip(*items))

        for i, x in enumerate(self.nested_columns):
            await x.write_data(
                list(items[i]),
            )

    async def write_items(
        self,
        items,
    ):
        return await self.write_data(
            items,
        )

    async def read_data(
        self,
        n_items,
    ):
        rv = [
            await x.read_data(
                n_items,
            )
            for x in self.nested_columns
        ]
        return list(zip(*rv))

    async def read_items(
        self,
        n_items,
    ):
        return await self.read_data(
            n_items,
        )


def create_tuple_column(spec, column_by_spec_getter, column_options):
    brackets = 0
    column_begin = 0

    inner_spec = get_inner_spec(spec)
    nested_columns = []
    for i, x in enumerate(inner_spec + ","):
        if x == ",":
            if brackets == 0:
                nested_columns.append(inner_spec[column_begin:i])
                column_begin = i + 1
        elif x == "(":
            brackets += 1
        elif x == ")":
            brackets -= 1
        elif x == " ":
            if brackets == 0:
                column_begin = i + 1

    return TupleColumn([column_by_spec_getter(x) for x in nested_columns], **column_options)


def get_inner_spec(spec):
    brackets = 1
    offset = len("Tuple(")
    i = offset
    for i, ch in enumerate(spec[offset:], offset):
        if brackets == 0:
            break

        if ch == "(":
            brackets += 1

        elif ch == ")":
            brackets -= 1

    return spec[offset:i]