"""Base configuration for DerivaML applications.

This module defines the base configuration and helper functions that simplify
creating hydra-zen configurations for both script execution and notebooks.

Simple Usage (notebooks using only BaseConfig fields):
    # In configs/my_notebook.py
    from deriva_ml.execution import notebook_config

    notebook_config(
        "my_notebook",
        defaults={"assets": "my_assets", "datasets": "my_dataset"},
    )

    # In notebook
    from deriva_ml.execution import run_notebook
    ml, execution, config = run_notebook("my_notebook")

Advanced Usage (notebooks with custom parameters):
    # In configs/my_analysis.py
    from dataclasses import dataclass
    from deriva_ml.execution import BaseConfig, notebook_config

    @dataclass
    class MyAnalysisConfig(BaseConfig):
        threshold: float = 0.5
        num_samples: int = 100

    notebook_config(
        "my_analysis",
        config_class=MyAnalysisConfig,
        defaults={"assets": "analysis_assets"},
    )

    # In notebook
    from deriva_ml.execution import run_notebook
    ml, execution, config = run_notebook("my_analysis")
    print(config.threshold)  # 0.5
"""

import importlib
import json
import os
import pkgutil
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, TypeVar, TYPE_CHECKING

from hydra_zen import builds, instantiate, launch, make_config, store

if TYPE_CHECKING:
    from deriva_ml import DerivaML
    from deriva_ml.execution import Execution

T = TypeVar("T")


# Standard hydra defaults for DerivaML applications.
# Projects can customize these or define their own defaults.
base_defaults = [
    "_self_",
    {"deriva_ml": "default_deriva"},
    {"datasets": "default_dataset"},
    {"assets": "default_asset"},
    {"workflow": "default_workflow"},
    {"model_config": "default_model"},
]


@dataclass
class BaseConfig:
    """Base configuration for DerivaML applications.

    This dataclass defines the common configuration structure shared by
    both script execution and notebook modes. Project-specific configs
    should inherit from this class to get the standard DerivaML fields.

    Note:
        Fields use ``Any`` type annotations because several DerivaML types
        (DerivaMLConfig, DatasetSpec) are Pydantic models which are not
        compatible with OmegaConf structured configs. The actual types at
        runtime are documented below.

    Attributes:
        deriva_ml: DerivaML connection configuration (DerivaMLConfig at runtime).
        datasets: List of dataset specifications (list[DatasetSpec] at runtime).
        assets: List of asset RIDs to load (list[str] at runtime).
        dry_run: If True, skip catalog writes (for testing/debugging).
        description: Human-readable description of this run.
        config_choices: Dictionary mapping config group names to selected config names.
            This is automatically populated by get_notebook_configuration() with the
            Hydra runtime choices (e.g., {"model_config": "cifar10_quick", "assets": "roc_quick"}).
            Useful for tracking which configurations were used in an execution.

    Example:
        >>> from dataclasses import dataclass
        >>> from deriva_ml.execution import BaseConfig
        >>>
        >>> @dataclass
        ... class MyConfig(BaseConfig):
        ...     learning_rate: float = 0.001
        ...     epochs: int = 10
    """
    deriva_ml: Any = None
    datasets: Any = field(default_factory=list)
    assets: Any = field(default_factory=list)
    dry_run: bool = False
    description: str = ""
    config_choices: dict[str, str] = field(default_factory=dict)


# Create and register the base config with hydra-zen store.
# This provides a ready-to-use base that experiments can inherit from.
DerivaBaseConfig = builds(
    BaseConfig,
    populate_full_signature=True,
    hydra_defaults=base_defaults,
)

store(DerivaBaseConfig, name="deriva_base")


def get_notebook_configuration(
    config_class: type[T],
    config_name: str,
    overrides: list[str] | None = None,
    job_name: str = "notebook",
    version_base: str = "1.3",
) -> T:
    """Load and return a hydra-zen configuration for use in notebooks.

    This function is the notebook equivalent of `run_model`. While `run_model`
    launches a full execution with model training, `get_notebook_configuration`
    simply resolves the configuration and returns it for interactive use.

    The function handles:
    - Adding configurations to the hydra store
    - Launching hydra-zen to resolve defaults and overrides
    - Returning the instantiated configuration object

    Args:
        config_class: The hydra-zen builds() class for the configuration.
            This should be a class created with `builds(YourConfig, ...)`.
        config_name: Name of the configuration in the hydra store.
            Must match the name used when calling `store(config_class, name=...)`.
        overrides: Optional list of Hydra override strings (e.g., ["param=value"]).
        job_name: Name for the Hydra job (default: "notebook").
        version_base: Hydra version base (default: "1.3").

    Returns:
        The instantiated configuration object with all defaults resolved.

    Example:
        In your notebook's configuration module (e.g., `configs/roc_analysis.py`):

        >>> from dataclasses import dataclass, field
        >>> from hydra_zen import builds, store
        >>> from deriva_ml.execution import BaseConfig
        >>>
        >>> @dataclass
        ... class ROCAnalysisConfig(BaseConfig):
        ...     execution_rids: list[str] = field(default_factory=list)
        >>>
        >>> ROCAnalysisConfigBuilds = builds(
        ...     ROCAnalysisConfig,
        ...     populate_full_signature=True,
        ...     hydra_defaults=["_self_", {"deriva_ml": "default_deriva"}],
        ... )
        >>> store(ROCAnalysisConfigBuilds, name="roc_analysis")

        In your notebook:

        >>> from configs import load_all_configs
        >>> from configs.roc_analysis import ROCAnalysisConfigBuilds
        >>> from deriva_ml.execution import get_notebook_configuration
        >>>
        >>> # Load all project configs into hydra store
        >>> load_all_configs()
        >>>
        >>> # Get resolved configuration
        >>> config = get_notebook_configuration(
        ...     ROCAnalysisConfigBuilds,
        ...     config_name="roc_analysis",
        ...     overrides=["execution_rids=[3JRC,3KT0]"],
        ... )
        >>>
        >>> # Use the configuration
        >>> print(config.execution_rids)  # ['3JRC', '3KT0']
        >>> print(config.deriva_ml.hostname)  # From default_deriva config

    Environment Variables:
        DERIVA_ML_HYDRA_OVERRIDES: JSON-encoded list of override strings.
            When running via `deriva-ml-run-notebook`, this is automatically
            set from command-line arguments. Overrides from this environment
            variable are applied first, then any overrides passed directly
            to this function are applied (taking precedence).
    """
    # Ensure configs are in the hydra store
    store.add_to_hydra_store(overwrite_ok=True)

    # Collect overrides from environment variable (set by run_notebook CLI)
    env_overrides_json = os.environ.get("DERIVA_ML_HYDRA_OVERRIDES")
    env_overrides = json.loads(env_overrides_json) if env_overrides_json else []

    # Merge overrides: env overrides first, then explicit overrides (higher precedence)
    all_overrides = env_overrides + (overrides or [])

    # Variable to capture choices from within the task function
    captured_choices: dict[str, str] = {}

    # Define a task function that instantiates and returns the config
    # The cfg from launch() is an OmegaConf DictConfig, so we need to
    # use hydra_zen.instantiate() to convert it to actual Python objects
    def return_instantiated_config(cfg: Any) -> T:
        nonlocal captured_choices
        # Capture the Hydra runtime choices (which config names were selected)
        # Filter out None values (some Hydra internal groups have None choices)
        try:
            from hydra.core.hydra_config import HydraConfig
            choices = HydraConfig.get().runtime.choices
            captured_choices = {k: v for k, v in choices.items() if v is not None}
        except Exception:
            # If HydraConfig is not available, leave choices empty
            pass
        return instantiate(cfg)

    # Launch hydra-zen to resolve the configuration
    result = launch(
        config_class,
        return_instantiated_config,
        version_base=version_base,
        config_name=config_name,
        job_name=job_name,
        overrides=all_overrides,
    )

    # Inject the captured choices into the config object
    config = result.return_value
    if hasattr(config, "config_choices"):
        config.config_choices = captured_choices

    return config


# ---------------------------------------------------------------------------
# Registry for notebook configurations
# ---------------------------------------------------------------------------
# Maps config_name -> (config_builds_class, config_name)
_notebook_configs: dict[str, tuple[Any, str]] = {}


def notebook_config(
    name: str,
    config_class: type[BaseConfig] | None = None,
    defaults: dict[str, str] | None = None,
    **field_defaults: Any,
) -> Any:
    """Register a notebook configuration with simplified syntax.

    This is the recommended way to create notebook configurations. It handles
    all the hydra-zen boilerplate (builds, store, defaults) automatically.

    For simple notebooks that only use BaseConfig fields (deriva_ml, datasets,
    assets, etc.), just specify which defaults to use. For notebooks with
    custom parameters, provide a config_class that inherits from BaseConfig.

    Args:
        name: Configuration name. Used both as the hydra config name and
            to look up the config in run_notebook().
        config_class: Optional dataclass inheriting from BaseConfig. If None,
            uses BaseConfig directly (suitable for notebooks that only need
            the standard fields).
        defaults: Dict mapping config group names to config names. These
            override the base defaults. Common groups:
            - "deriva_ml": Connection config (e.g., "default_deriva", "eye_ai")
            - "datasets": Dataset config (e.g., "cifar10_training")
            - "assets": Asset config (e.g., "model_weights")
            - "workflow": Workflow config (e.g., "default_workflow")
        **field_defaults: Default values for fields in config_class.

    Returns:
        The hydra-zen builds() class, in case you need to reference it directly.

    Examples:
        Simple notebook using only standard fields:

            # configs/roc_analysis.py
            from deriva_ml.execution import notebook_config

            notebook_config(
                "roc_analysis",
                defaults={"assets": "roc_comparison_probabilities"},
            )

        Notebook with custom parameters:

            # configs/training_analysis.py
            from dataclasses import dataclass
            from deriva_ml.execution import BaseConfig, notebook_config

            @dataclass
            class TrainingAnalysisConfig(BaseConfig):
                learning_rate: float = 0.001
                batch_size: int = 32

            notebook_config(
                "training_analysis",
                config_class=TrainingAnalysisConfig,
                defaults={"datasets": "cifar10_training"},
                learning_rate=0.01,  # Override default
            )
    """
    # Use BaseConfig if no custom class provided
    actual_class = config_class or BaseConfig

    # Build the hydra defaults list
    hydra_defaults = ["_self_"]

    # Start with base defaults, then apply overrides
    default_groups = {
        "deriva_ml": "default_deriva",
        "datasets": "default_dataset",
        "assets": "default_asset",
    }
    if defaults:
        default_groups.update(defaults)

    for group, config_name in default_groups.items():
        hydra_defaults.append({group: config_name})

    # Create the hydra-zen builds() class
    config_builds = builds(
        actual_class,
        populate_full_signature=True,
        hydra_defaults=hydra_defaults,
        **field_defaults,
    )

    # Register with hydra-zen store
    store(config_builds, name=name)

    # Also register in our internal registry for run_notebook()
    _notebook_configs[name] = (config_builds, name)

    return config_builds


def load_configs(package_name: str = "configs") -> list[str]:
    """Dynamically import all configuration modules from a package.

    This function discovers and imports all Python modules in the specified
    package. Each module is expected to register its configurations with
    the hydra-zen store as a side effect of being imported.

    Args:
        package_name: Name of the package containing config modules.
            Default is "configs" which works for the standard project layout.

    Returns:
        List of module names that were successfully loaded.

    Raises:
        ImportError: If a config module fails to import.

    Example:
        # In your main script or notebook
        from deriva_ml.execution import load_configs

        load_configs()  # Loads from "configs" package
        # or
        load_configs("my_project.configs")  # Custom package

    Note:
        The "experiments" module (if present) is loaded last because it
        typically depends on other configs being registered first.
    """
    loaded_modules = []

    try:
        package = importlib.import_module(package_name)
    except ImportError:
        # Package doesn't exist, return empty
        return []

    package_dir = Path(package.__file__).parent

    # Collect module names
    modules_to_load = []
    for module_info in pkgutil.iter_modules([str(package_dir)]):
        modules_to_load.append(module_info.name)

    # Sort modules but ensure 'experiments' is loaded last
    modules_to_load.sort()
    if "experiments" in modules_to_load:
        modules_to_load.remove("experiments")
        modules_to_load.append("experiments")

    for module_name in modules_to_load:
        importlib.import_module(f"{package_name}.{module_name}")
        loaded_modules.append(module_name)

    return sorted(loaded_modules)


def run_notebook(
    config_name: str,
    overrides: list[str] | None = None,
    workflow_name: str | None = None,
    workflow_type: str = "Analysis Notebook",
    ml_class: type["DerivaML"] | None = None,
    config_package: str = "configs",
) -> tuple["DerivaML", "Execution", BaseConfig]:
    """Initialize a notebook with DerivaML execution context.

    This is the main entry point for notebooks. It handles all the setup:
    1. Loads all config modules from the config package
    2. Resolves the hydra-zen configuration
    3. Creates the DerivaML connection
    4. Creates a workflow and execution context
    5. Downloads any specified datasets and assets

    Args:
        config_name: Name of the notebook configuration (registered via
            notebook_config() or store()).
        overrides: Optional list of Hydra override strings
            (e.g., ["assets=different_assets"]).
        workflow_name: Name for the workflow. Defaults to config_name.
        workflow_type: Type of workflow (default: "Analysis Notebook").
        ml_class: Optional DerivaML subclass to use. If None, uses DerivaML.
        config_package: Package containing config modules (default: "configs").

    Returns:
        Tuple of (ml_instance, execution, config):
        - ml_instance: Connected DerivaML (or subclass) instance
        - execution: Execution context with downloaded inputs
        - config: Resolved configuration object

    Example:
        # Simple usage
        from deriva_ml.execution import run_notebook

        ml, execution, config = run_notebook("roc_analysis")

        # Access config values
        print(config.assets)
        print(config.deriva_ml.hostname)

        # Use ml and execution
        for asset_table, paths in execution.asset_paths.items():
            for path in paths:
                print(f"Downloaded: {path.file_name}")

        # At the end of notebook
        execution.upload_execution_outputs()

    Example with overrides:
        ml, execution, config = run_notebook(
            "roc_analysis",
            overrides=["assets=roc_quick_probabilities"],
        )

    Example with custom ML class:
        from eye_ai import EyeAI

        ml, execution, config = run_notebook(
            "eye_analysis",
            ml_class=EyeAI,
        )
    """
    # Import here to avoid circular imports
    from deriva_ml import DerivaML
    from deriva_ml.execution import Execution, ExecutionConfiguration

    # Load all config modules
    load_configs(config_package)

    # Get the config builds class from our registry or try the store
    if config_name in _notebook_configs:
        config_builds, _ = _notebook_configs[config_name]
    else:
        # Fall back to looking up in hydra store by building a simple config
        # This handles configs registered the old way
        config_builds = DerivaBaseConfig

    # Resolve the configuration
    config = get_notebook_configuration(
        config_builds,
        config_name=config_name,
        overrides=overrides,
    )

    # Create DerivaML instance
    actual_ml_class = ml_class or DerivaML
    ml = actual_ml_class(
        hostname=config.deriva_ml.hostname,
        catalog_id=config.deriva_ml.catalog_id,
    )

    # Create workflow
    actual_workflow_name = workflow_name or config_name.replace("_", " ").title()
    workflow = ml.create_workflow(
        name=actual_workflow_name,
        workflow_type=workflow_type,
        description=config.description or f"Running {config_name}",
    )

    # Create execution configuration
    exec_config = ExecutionConfiguration(
        workflow=workflow,
        datasets=config.datasets if config.datasets else [],
        assets=config.assets if config.assets else [],
        description=config.description or f"Execution of {config_name}",
    )

    # Create execution context (downloads inputs)
    execution = Execution(configuration=exec_config, ml_object=ml)

    return ml, execution, config


class DescribedList(list):
    """A list with an attached description.

    This class extends list to add a `description` attribute while maintaining
    full list compatibility. This allows configuration values (like asset RIDs
    or dataset specs) to carry documentation without changing how they're used.

    When stored in hydra-zen and resolved via `instantiate()`, the result is a
    DescribedList that behaves like a regular list but has a `description` attribute.

    Attributes:
        description: Human-readable description of this configuration.

    Example:
        >>> from hydra_zen import store
        >>> from deriva_ml.execution import with_description
        >>>
        >>> asset_store = store(group="assets")
        >>> asset_store(
        ...     with_description(
        ...         ["3WMG", "3XPA"],
        ...         "Model weights from quick and extended training",
        ...     ),
        ...     name="comparison_weights",
        ... )
        >>>
        >>> # After instantiation, usage is identical to a regular list:
        >>> # config.assets[0]  # "3WMG"
        >>> # len(config.assets)  # 2
        >>> # for rid in config.assets: ...
        >>> # config.assets.description  # "Model weights from..."
    """

    def __init__(self, items: list | None = None, description: str = ""):
        """Initialize a DescribedList.

        Args:
            items: Initial list items. If None, creates empty list.
            description: Human-readable description of this list.
        """
        super().__init__(items or [])
        self.description = description

    def __repr__(self) -> str:
        """Return string representation including description."""
        if self.description:
            return f"DescribedList({list(self)!r}, description={self.description!r})"
        return f"DescribedList({list(self)!r})"


def _make_described_list(items: list, description: str = "") -> DescribedList:
    """Factory function for creating DescribedList instances.

    This is used internally by `with_description` to create a hydra-zen
    compatible config.
    """
    return DescribedList(items, description)


# Pre-built config for DescribedList
_DescribedListConfig = builds(_make_described_list, populate_full_signature=True)


def with_description(items: list, description: str) -> Any:
    """Create a hydra-zen config for a list with an attached description.

    Use this to add descriptions to configuration values like asset RIDs
    or dataset specifications. The result is a hydra-zen config that, when
    instantiated, produces a DescribedList.

    Args:
        items: List items (e.g., asset RIDs, dataset specs).
        description: Human-readable description of this configuration.

    Returns:
        A hydra-zen config that instantiates to a DescribedList.

    Example:
        >>> from hydra_zen import store
        >>> from deriva_ml.execution import with_description
        >>>
        >>> # Assets with description
        >>> asset_store = store(group="assets")
        >>> asset_store(
        ...     with_description(
        ...         ["3WMG", "3XPA"],
        ...         "Model weights from quick and extended training runs",
        ...     ),
        ...     name="comparison_weights",
        ... )
        >>>
        >>> # Datasets with description
        >>> from deriva_ml.dataset import DatasetSpecConfig
        >>> datasets_store = store(group="datasets")
        >>> datasets_store(
        ...     with_description(
        ...         [DatasetSpecConfig(rid="28CT", version="0.21.0")],
        ...         "Complete CIFAR-10 dataset with 10,000 images",
        ...     ),
        ...     name="cifar10_complete",
        ... )
        >>>
        >>> # After instantiation:
        >>> # config.assets is a DescribedList
        >>> # config.assets[0]  # "3WMG"
        >>> # config.assets.description  # "Model weights from..."

    Note:
        For model configs created with `builds()`, use the `zen_meta` parameter
        instead:

        >>> model_store(
        ...     Cifar10CNNConfig,
        ...     name="cifar10_quick",
        ...     epochs=3,
        ...     zen_meta={"description": "Quick training - 3 epochs"},
        ... )
    """
    return _DescribedListConfig(items=items, description=description)
