import os
import importlib
from pathlib import Path
from collections import defaultdict

def get_versioned_modules(version_mode: str, older_version: str = None) -> dict:
    """
    Load versioned Python modules based on the provided version mode.

    Parameters:
    -----------
    version_mode : str
        One of "newer", "stag", or "older". Determines which codebase to use.

    older_version : str, optional
        Required if version_mode is "older" (e.g., 'v2').
        Ignored otherwise.

    Returns:
    --------
    modules : dict
        Dictionary of module names mapped to their imported module objects.

    Example:
    --------
    >>> modules = get_versioned_modules("older", "v2")
    >>> modules["feature_transaction2"].run()
    >>> modules = get_versioned_modules("newer")
    >>> modules["feature_funding"].calculate()
    """

    if version_mode not in {"newer", "stag", "older"}:
        raise ValueError("version_mode must be one of: 'newer', 'stag', 'older'")

    project_root = Path(__file__).resolve().parent.parent
    lib_ff_path = project_root / "lib" / "FF"

    if version_mode == "older":
        if not older_version or not older_version.startswith("v"):
            raise ValueError("For 'older', provide a valid version string like 'v2'")
        version_int = int(older_version.lstrip("v"))
        base_path = "lib.FF.older"
        dir_path = lib_ff_path / "older"
    else:
        base_path = f"lib.FF.{version_mode}"
        dir_path = lib_ff_path / version_mode
        if not dir_path.exists():
            raise FileNotFoundError(f"Directory not found: {dir_path}")
        return _load_all_modules_from_dir(base_path, dir_path)

    # Load versioned modules from older with fallback
    versioned_files = defaultdict(dict)
    for f in os.listdir(dir_path):
        if f.endswith(".py") and f.startswith("v"):
            parts = f.split("_", 1)
            if len(parts) == 2:
                ver = parts[0]
                modname = parts[1][:-3]
                versioned_files[modname][ver] = f

    modules = {}
    for modname, versions in versioned_files.items():
        selected_ver = None
        for i in reversed(range(1, version_int + 1)):
            try_ver = f"v{i}"
            if try_ver in versions:
                selected_ver = try_ver
                break

        if selected_ver:
            py_module_name = versions[selected_ver][:-3]
            full_import_path = f"{base_path}.{py_module_name}"
            try:
                modules[modname] = importlib.import_module(full_import_path)
            except ModuleNotFoundError as e:
                raise ImportError(f"Failed to import module '{full_import_path}'") from e

    return modules


def _load_all_modules_from_dir(base_path: str, dir_path: Path) -> dict:
    """
    Helper to load all .py modules in a given directory.

    Parameters:
    -----------
    base_path : str
        Python import path (e.g., 'lib.FF.newer')

    dir_path : Path
        Filesystem path to the directory containing .py files

    Returns:
    --------
    modules : dict
        Dictionary of imported modules
    """
    modules = {}
    for file in os.listdir(dir_path):
        if file.endswith(".py") and not file.startswith("__"):
            modname = file[:-3]
            full_import_path = f"{base_path}.{modname}"
            try:
                modules[modname] = importlib.import_module(full_import_path)
            except ModuleNotFoundError as e:
                raise ImportError(f"Failed to import module '{full_import_path}'") from e
    return modules

  
## Load from newer
#modules = get_versioned_modules("newer")
#modules["feature_funding"].calculate()
#
## Load from stag
#modules = get_versioned_modules("stag")
#modules["feature_validation"].run()
#
## Load from older version (fallback logic)
#modules = get_versioned_modules("older", "v3")
#modules["feature_transaction2"].run()
