from typing import Generator

from sortedcontainers import SortedSet


class PathDependencyTracker(object):
    """
    Keep track of paths in a way to make querying if other paths are affected by them optimal.
    Why a path affects another is up for the user code to decide. Upstream, downstream or both ways can be checked
    for tracked dependencies.
    All paths must be POSIX-style normalized (case-corrected and resolved) for this to operate as intended.
    """

    __slots__ = ("_affects_upstream", "_affects_downstream")

    def __init__(self):
        # Using the hierarchical nature of paths we collect changes in a sorted collection so up- and downstream
        # lookups can be done by bisection.
        self._affects_upstream = SortedSet()
        self._affects_downstream = SortedSet()

    def clear(self):
        """
        Clear tracked paths.
        """
        self._affects_upstream.clear()
        self._affects_downstream.clear()

    def track_upstream(self, path: str):
        self._affects_upstream.add(path)

    def track_downstream(self, path: str):
        self._affects_downstream.add(path)

    def track(self, path: str):
        """
        Add path for tracking in both directions.
        """
        self._affects_upstream.add(path)
        self._affects_downstream.add(path)

    def release(self, path: str):
        """
        Remove path from tracking in both directions.
        """
        self._affects_upstream.discard(path)
        self._affects_downstream.discard(path)

    def affects(self, path: str) -> Generator[str, None, None]:
        """
        Get affecting paths of the argument as a generator.
        NOTE: If a path is affected by both upstream and downstream it will be yielded twice!
        """
        # check downstream affecting changes
        for p in self._affects_downstream.irange(maximum=path, reverse=True):
            if path.startswith(p):
                if (pl := len(p)) < len(path):
                    if path[pl] == "/":  # parent match
                        yield p
                else:  # length is implied to be equal, meaning an exact match
                    yield p
            else:  # no more potential matches
                break
        # check upstream affecting changes
        for p in self._affects_upstream.irange(minimum=path):
            if p.startswith(path):
                if (pl := len(path)) < len(p):
                    if p[pl] == "/":  # parent match
                        yield p
                else:  # length is implied to be equal, meaning an exact match
                    yield p
            else:  # no more potential matches
                break
