"""Internal support module.

Conceptually, this module groups the lightweight session-discovery helpers and
minimal unit-table accessors used by the session-aware analysis stack.

It exists as a separate unit so session listing and area discovery remain
usable without pulling in the heavier context-analysis logic.

It connects session roots and session-like payloads to the unit tables consumed
by `session_trials` and `session_context`.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, List, Mapping, Optional, Union

import pandas as pd

PathLike = Union[str, Path]

__all__ = [
    "PathLike",
    "get_available_areas_for_session",
    "list_available_sessions",
]


def list_available_sessions(data_root: PathLike, pattern: str = "*.nwb") -> List[Path]:
    """Return sorted NWB session paths from a dataset root."""

    data_root = Path(data_root)
    return sorted(data_root.glob(pattern))


def _session_label(session_id: Optional[str], session_path: Optional[PathLike]) -> str:
    """Return a readable label for logs and fallback messages."""

    if session_id:
        return str(session_id)
    if session_path is not None:
        return Path(session_path).name
    return "session"


def _units_df_from_session_or_path(session_or_path: Union[Mapping[str, Any], PathLike]) -> pd.DataFrame:
    """Return a units table from a session mapping or an NWB path.

    Parameters
    ----------
    session_or_path:
        Either a session-like mapping containing `units_df` or a path to an
        NWB file.

    Returns
    -------
    pd.DataFrame
        Units table for the session, or an empty frame when no units are
        available.
    """

    if isinstance(session_or_path, Mapping):
        units_df = session_or_path.get("units_df", None)
        if isinstance(units_df, pd.DataFrame):
            return units_df
        return pd.DataFrame()

    # Keep the NWB import local so session-aware helpers that only work with
    # in-memory session mappings do not require the full NWB stack.
    from pynwb import NWBHDF5IO

    session_path = Path(session_or_path)
    with NWBHDF5IO(str(session_path), mode="r", load_namespaces=True) as io:
        nwb = io.read()
        return nwb.units.to_dataframe() if nwb.units is not None else pd.DataFrame()


def get_available_areas_for_session(session_or_path: Union[Mapping[str, Any], PathLike]) -> List[str]:
    """Return sorted area labels available in a session or session mapping."""

    units_df = _units_df_from_session_or_path(session_or_path)
    if "location" not in units_df.columns:
        return []

    areas = [str(x).strip() for x in pd.unique(units_df["location"].astype(str).str.strip())]
    return sorted([area for area in areas if area and area.lower() != "nan"])