import os
import json

def fetch_block_indices(
    base_dir,
    model_name,
    task_name,
    sample_name,
    config,
):
    """
    Returns a list of generation step indices to block based on POS tags.

    Args:
        base_dir (str): Base directory containing tagged outputs
        model_name (str)
        task_name (str)
        sample_name (str)
        config (dict): blocking config with keys:
            - pos_tag_to_block (str)
            - only_last (bool)
            - only_first (bool)

    Returns:
        List[int]: generation step indices
    """
    file_path = os.path.join(
        base_dir,
        model_name,
        task_name,
        f"{sample_name}.json"
    )

    if not os.path.exists(file_path):
        raise FileNotFoundError(f"Block indices file not found: {file_path}")

    with open(file_path, "r") as f:
        tagged_tokens = json.load(f)

    pos_tags = config.get("pos_tag_to_block")
    only_last = config.get("only_last", False)
    only_first = config.get("only_first", False)

    if only_last and only_first:
        raise ValueError("Cannot set both `only_last` and `only_first` to True")

    if pos_tags is None:
        raise ValueError("`pos_tag_to_block` must be specified in config")

    if isinstance(pos_tags, str):
        pos_tags = {pos_tags}
    else:
        pos_tags = set(pos_tags)
    # collect all indices where tag matches


    indices = [
        i+1 for i, tok in enumerate(tagged_tokens)
        if tok.get("tag") in pos_tags
    ]

    if not indices:
        return []

    if only_last:
        return [indices[-1]]

    if only_first:
        return [indices[0]]

    return indices
