from nile.api.v1 import (
    stream as nstream,
    Record,
)
from typing import (
    List
)
from maps.wikimap.stat.arm.common.lib import helpers


def modify_obj_type(objects: nstream.Stream) -> nstream.Stream:
    def type_to_tree(
        object_type: str,
        category: str,
        template_id: int,
        fake: bool,
    ) -> List[str]:
        ret_val = ['_total_']
        if object_type == 'closure':
            ret_val.append('closure')
            if template_id is not None:
                ret_val.append('template-no-transit')
            elif fake:
                ret_val.append('fake-no-transit')
            elif category == 'closed_maneuver':
                ret_val.append('maneuver')
            else:
                ret_val.append('no-transit')
        else:
            ret_val.append('template')

        return helpers.list_to_tree(ret_val)

    def object_fd_mapper(rows):
        for row in rows:
            for obj_type in type_to_tree(
                row.object_type.decode(),
                row.category.decode(),
                row.template_id,
                row.fake,
            ):
                yield Record(row, object_type=obj_type)

    return objects.map(object_fd_mapper)
