from functools import lru_cache
from pathlib import Path
from typing import Iterable

import bidsschematools.schema as bst
import pandas as pd
import pandera.pandas as pa
from typer import BadParameter

from bagel import bids_table_model, models
from bagel.logger import log_error, logger
from bagel.utilities import file_utils

# NOTE: A copy of the imaging modality vocab will likely end up in all community config directories,
# but since the contents will be the same, we always pull it from the Neurobagel config for now for simplicity.
IMAGING_MODALITIES_URL = "https://raw.githubusercontent.com/neurobagel/communities/refs/heads/main/configs/Neurobagel/imaging_modalities.json"
IMAGING_MODALITIES_PATH = (
    Path(__file__).parents[1]
    / "communities/configs/Neurobagel/imaging_modalities.json"
)

bids_schema = bst.load_schema()


def get_bids_suffix_to_std_term_mapping() -> dict[str, str]:
    """
    Fetch the standardized imaging modality vocabulary from the neurobagel/communities repository
    and return a mapping of supported BIDS suffixes to prefixed standardized terms.

    Returns:
        dict[str, str]: A mapping where keys are BIDS suffixes (e.g., "T1w", "bold")
                        and values are namespaced standardized terms (e.g., "nidm:T1Weighted").

    Raises:
        typer.Exit: If the vocabulary cannot be fetched from either the remote URL or local backup.
    """
    # TODO: Revisit once we cache community config files locally as part of https://github.com/neurobagel/bagel-cli/issues/493.
    # For now we revert to a local submodule backup if the request fails, but this means we might not have the latest vocab.
    bids_terms_vocab, err = file_utils.request_file(
        url=IMAGING_MODALITIES_URL, backup_path=IMAGING_MODALITIES_PATH
    )
    if bids_terms_vocab == []:
        log_error(
            logger,
            f"Failed to fetch the standardized imaging modality vocabulary required to validate your BIDS metadata. Error: {err} "
            "Please check that you have an internet connection and try again, or open an issue in https://github.com/neurobagel/bagel-cli/issues if the problem persists.",
        )
    # We only expect one term namespace in the bids_term_vocab list
    bids_terms_vocab = bids_terms_vocab[0]
    namespace_prefix = bids_terms_vocab["namespace_prefix"]
    bids_suffix_to_std_term_mapping = {}
    for term in bids_terms_vocab["terms"]:
        bids_suffix_to_std_term_mapping[term["abbreviation"]] = (
            f"{namespace_prefix}:{term['id']}"
        )

    return bids_suffix_to_std_term_mapping


@lru_cache()
def get_all_bids_suffixes() -> set[str]:
    """Return all file suffixes that are recognized by BIDS."""
    return {
        bids_suffix["value"]
        for bids_suffix in bids_schema.objects.suffixes.values()
    }


@lru_cache()
def get_bids_raw_data_suffixes() -> set[str]:
    """Return all BIDS recognized suffixes corresponding to raw data files."""
    bids_raw_data_suffixes = set()
    for file_datatype in bids_schema.rules.files.raw.values():
        for file_subtype in file_datatype.values():
            bids_raw_data_suffixes.update(file_subtype["suffixes"])

    return bids_raw_data_suffixes


def partition_suffixes(
    suffixes: Iterable[str], reference_suffixes: Iterable[str]
) -> tuple[set[str], set[str]]:
    """
    Partition suffixes into those found in a reference collection and those not found.

    Parameters
    ----------
    suffixes : Iterable[str]
        File suffixes to partition.
    reference_suffixes : Iterable[str]
        Suffixes to compare the input list against.

    Returns
    -------
    tuple[set[str], set[str]]
        A tuple containing two sets:
        - The first set contains suffixes found in the reference list.
        - The second set contains suffixes not found in the reference list.
    """
    reference_set = set(reference_suffixes)
    in_reference = set()
    not_in_reference = set()
    for suffix in suffixes:
        if suffix in reference_set:
            in_reference.add(suffix)
        else:
            not_in_reference.add(suffix)
    return in_reference, not_in_reference


def check_absolute_path(dir_path: Path | None) -> Path | None:
    """
    Raise an error if the input path does not look like an absolute path.
    This is a workaround for --dataset-source-path not requiring the path to exist on the host machine
    (and thus not being able to resolve the path automatically).
    """
    # Allow POSIX-style absolute paths (e.g., "/data/...") across OSes - useful for referencing paths on remote Unix servers.
    if dir_path is not None and not (
        dir_path.is_absolute() or dir_path.as_posix().startswith("/")
    ):
        raise BadParameter(
            "Dataset source directory must be an absolute path."
        )
    return dir_path


def validate_bids_table(bids_table: pd.DataFrame):
    """Error and exit if the provided BIDS table is empty or fails schema validation."""
    try:
        bids_table_model.model.validate(bids_table)
    except pa.errors.SchemaError as err:
        rows_with_errs_msg = ""
        # When validation fails due to a column value check (e.g., as opposed to a missing column),
        # printing the row indices helps with debugging, especially for invalid empty values.
        if isinstance(err.failure_cases, pd.DataFrame):
            rows_with_errs_msg = f"Rows with error (0 = first non-header row): {err.failure_cases['index'].tolist()}. "
        log_error(
            logger,
            f"Invalid BIDS table. Error: {err}. {rows_with_errs_msg}",
        )
    if bids_table.empty:
        log_error(
            logger,
            "BIDS table is empty (only a header row was found). No imaging metadata to add.",
        )


def create_acquisitions(
    session_df: pd.DataFrame,
    bids_suffix_term_map: dict,
) -> list:
    """Parses BIDS image file suffixes for a specified session to create a list of Acquisition objects."""
    image_list = []

    for bids_file_suffix in session_df["suffix"]:
        if (
            mapped_term := bids_suffix_term_map.get(bids_file_suffix)
        ) is not None:
            image_list.append(
                models.Acquisition(
                    hasContrastType=models.Image(identifier=mapped_term)
                )
            )

    return image_list


def get_session_path(
    dataset_root: Path | None,
    bids_sub_id: str,
    session_id: str,
) -> str:
    """
    Construct the session directory or subject directory path (when there is no session ID) from the source BIDS directory path, subject ID, and session ID.
    If no source BIDS directory is available, return a relative path.
    """
    subject_path = (
        Path(dataset_root / bids_sub_id) if dataset_root else Path(bids_sub_id)
    )
    session_path = (
        subject_path / session_id if session_id.strip() != "" else subject_path
    )
    return session_path.as_posix()
