from pathlib import Path
from typing import Optional
from typing import Iterator


def build_stems(
    release_path: Path,
    annotations_dir: Path,
    annotation_type: str,
    split: str,
    partition: Optional[str] = None,
    split_type: str = "random",
) -> Iterator[str]:
    """
    Builds the stems for the given release with the given annotations as base.

    Parameters
    ----------
    release_path : Path
        The path of the ``Release`` saved locally.
    annotations_dir : Path
        The path for a directory where annotations.
    annotation_type : str
        The type of the annotations.
    split : str
        The split name.
    partition : Optional[str], default: None
        How to partition files. If no partition is specified, then it takes all the json files in
        the annotations directory.
        The resulting generator prepends parent directories relative to the main annotation
        directory.

        E.g.: ``["annotations/test/1.json", "annotations/2.json", "annotations/test/2/3.json"]``:

            - annotations/test/1
            - annotations/2
            - annotations/test/2/3
    split_type str, default: "random"
        The type of split. Can be ``"random"`` or ``"stratified"``.

    Returns
    -------
    Iterator[str]
        An iterator with the path for the stem files.

    Raises
    ------
    ValueError
        If the provided ``split_type`` is invalid.

    FileNotFoundError
        If no dataset partitions are found.
    """

    if partition is None:
        return (str(e.relative_to(annotations_dir).parent / e.stem) for e in sorted(annotations_dir.glob("**/*.json")))

    if split_type == "random":
        split_filename = f"{split_type}_{partition}.txt"
    elif split_type == "stratified":
        split_filename = f"{split_type}_{annotation_type}_{partition}.txt"
    else:
        raise ValueError(f'Unknown split type "{split_type}"')

    split_path = release_path / "lists" / split / split_filename
    if split_path.is_file():
        return (e.strip("\n\r") for e in split_path.open())

    raise FileNotFoundError(
        f"could not find a dataset partition. "
        f"Split the dataset using `split_dataset()` from `darwin.dataset.split_manager`"
    )