"""Various misc. utilities."""

import subprocess
from os import PathLike
from contextlib import contextmanager
import os
from dataclasses import dataclass
from typing import (
    Callable,
    Generic,
    Iterator,
    TypeVar,
    Union,
    overload,
    Literal,
)

# Ts = TypeVarTuple('Ts')
T = TypeVar('T')
V = TypeVar('V')


def find(basedir: str | bytes | PathLike[str] | PathLike[bytes],
         **flags: str | bytes) -> list[bytes]:
    """
    Run the shell command ``find``.

    :param basedir:
        Directory to search under.

    :param flags:
        Key-value pairs passed to find. Each key is prefixed with a
        single dash. Compound searches might be possible with some
        trickery, but extend this function instead of doing that.
    """
    cmdline: list[Union[str, bytes,
                        PathLike[str],
                        PathLike[bytes]]] = ['find', basedir]
    for key, value in flags.items():
        cmdline += [f'-{key}', value]
    cmdline.append('-print0')

    cmd = subprocess.run(cmdline, capture_output=True)
    return cmd.stdout.split(b'\0')[:-1]


@contextmanager
def cwd(path: str) -> Iterator[None]:
    """
    Context manager for changing directories.

    Changes the current working directory for the block. And changes
    it back after, no matter how the block is left.

    .. code-block:: python

        with cwd("/tmp"):
            print(os.getcwd())
            # /tmp
    """
    oldpwd = os.getcwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(oldpwd)


class chain(Generic[T]):
    """
    Chain functions after one another.

    .. code-block:: python

        (chain(range(9))
            @ sum
            @ sqrt).value
        # result: 6.0

    Is equivalent to

    .. code-block:: python

        sqrt(sum(range(9)))
        # result: 6.0
    """

    def __init__(self, value: T) -> None:
        self.value = value

    def __matmul__(self, proc: Callable[[T], V]) -> 'chain[V]':
        return chain(proc(self.value))


def force(x: T | None) -> T:
    """
    Discard the None case from an optional value.

    Only use when you *know* that the value is there.

    :raises AssertionError:
        When value is None after all.
    """
    assert x is not None
    return x


class MutableString:
    """
    A mutatable string.

    Strings are immutable by default in python. This works almost
    exactly like a regular string, but ``+=`` actually changes the
    object in place.
    """

    def __init__(self, init: str = '') -> None:
        self.str = init

    def __add__(self, other: str) -> 'MutableString':
        return MutableString(self.str + other)

    def __iadd__(self, other: str) -> 'MutableString':
        self.str += other
        return self

    def __repr__(self) -> str:
        return f'MutableString("{self.str}")'

    def __str__(self) -> str:
        return self.str


@dataclass
class Lists(Generic[T, V]):
    """
    A group of lists.

    In an ideal world, this would take any number of type parameters,
    and work correctly. Currently it's limited to exactly two lists.

    :param lists:
        The contained lists
    """

    def __init__(self, ts: list[T] = [], vs: list[V] = []) -> None:
        self.lists = (ts, vs)

    def __add__(self, other: 'Lists[T, V]') -> 'Lists[T, V]':
        """Add each sublist of the two groups."""
        return Lists(self[0] + other[0],
                     self[1] + other[1])

    @overload
    def __getitem__(self, idx: Literal[0]) -> list[T]:
        ...  # pragma: no cover

    @overload
    def __getitem__(self, idx: Literal[1]) -> list[V]:
        ...  # pragma: no cover

    def __getitem__(self, idx: int) -> list[T] | list[V]:
        return self.lists[idx]


P = TypeVar('P', str, bytes)


# TODO test this on a system with drive letters
def split_path(path: P) -> list[P]:
    """Split a path into all its components."""
    result: list[P] = []
    dir, item = os.path.split(path)
    while True:
        if not item:
            break
        result.insert(0, item)
        path = dir
        dir, item = os.path.split(path)
    return result