class MetaReplay:
    def __init__(self, transitions):
        cached_meta = {0: {}}

        running_state = {}
        for transition_point, transition in transitions.items():
            running_state[transition.metatarget] = transition.assignments
            cached_meta[transition_point.start] = running_state.copy()

        self.cached_meta = sorted(cached_meta.items(), key=lambda x: x[0])
        self.keys = [x[0] for x in self.cached_meta]

    def __getitem__(self, key):
        if isinstance(key, int):
            assert key >= 0

            relevant_meta_index = bisect.bisect_right(self.keys, key)
            if relevant_meta_index:
                return self.cached_meta[relevant_meta_index - 1][1]

        raise IndexError

    def get_value(self, index, segprop, segvalue):
        result = {"id": segvalue}
        relevant_meta = self[index]

        meta_target = MetaTarget(segprop, segvalue)
        if meta_target in relevant_meta:
            for assignment in relevant_meta[meta_target]:
                result[assignment.sprop] = assignment.value

        return result


class SegmentTopology:
    OVERLAPPING = "overlapping"
    COVERING = "covering"

    def __init__(self, segments, story_length):
        self.by_start = sorted(segments, key=lambda x: x.start)
        self.start_keys = [x.start for x in self.by_start]
        self.by_end = sorted(segments, key=lambda x: x.end)
        self.end_keys = [x.end for x in self.by_end]
        self.story_length = story_length

    def _end_after_inclusive(self, index):
        first = bisect.bisect_left(self.end_keys, index)
        return self.by_end[first:]

    def _start_strictly_before(self, index):
        post_last = bisect.bisect_left(self.start_keys, index)
        if post_last == None:
            return []
        return self.by_start[:post_last]

    def _start_before_inclusive(self, index):
        post_first = bisect.bisect_right(self.start_keys, index)
        if post_first == None:
            return []
        return self.by_start[:post_first]

    def _end_strictly_after(self, index):
        last = bisect.bisect_right(self.end_keys, index)
        return self.by_end[last:]

    def get_overlapping(self, start, end):
        # find overlapping segments that end after this one starts and start before
        # this one ends

        result = set()

        for candidate in self._end_after_inclusive(start):
            if candidate.start >= end:
                break
            result.add(candidate)

        for candidate in self._start_strictly_before(end)[::-1]:
            if candidate.end < start:
                break
            result.add(candidate)

        return result

    def get_covers(self, start, end):
        # TODO: use an rtree to make get_covers() more efficient

        # find covering segments that start before this one start and end after
        # this one ends

        start_before = set(self._start_before_inclusive(start))
        end_after = set(self._end_strictly_after(end))

        print(start_before)
        print(end_after)

        return start_before.intersection(end_after)

    def __getitem__(self, key):
        if isinstance(key, int):
            key = slice(key, key + 1, SegmentTopology.COVERING)

        start, end, _ = slice(key.start, key.stop).indices(self.story_length)
        match_type = key.step

        if match_type == SegmentTopology.OVERLAPPING:
            return self.get_overlapping(start, end)
        elif match_type == SegmentTopology.COVERING:
            return self.get_covers(start, end)

        raise TypeError(
            'match type ("stride" value) must be one of OVERLAPPING or COVERING'
        )


class SegmentDict:
    def __init__(self, items_by_seg, story_length):
        self._dict = items_by_seg.copy()
        self.segment_db = SegmentTopology(self._dict.keys(), story_length)

    def __getitem__(self, key):
        result = []
        for matching_segment in self.segment_db[key]:
            result.append((matching_segment, self._dict[matching_segment]))

        return result

class MetaSegmentDict:
    # TODO: finish this next

class CookieSynthNavigator:
    def __init__(self, marked_story):
        tokenizer = CookieSynthTokenizer()
        parse_tree = COOKIESYNTH_PARSER.parse(marked_story)
        tokenizer.visit(parse_tree)

        self.story = ''.join(tokenizer.story_pieces)
        self.story_length = len(self.story)
        self.meta = MetaReplay(tokenizer.meta_transitions)
        self.types = SegmentDict(tokenizer.stypes, self.story_length)
        self.assignments = SegmentDict(tokenizer.assignments, self.story_length)