"""Context search space utilities for hyperparameter optimization.

This module provides utilities for creating hyperopt search spaces
for context parameters in LeCrapaud.

Example:
    >>> from lecrapaud import create_context_search_space
    >>> from hyperopt import hp
    >>> space = create_context_search_space(
    ...     optimization_metric=["LOGLOSS", "ROC_AUC"],
    ...     use_class_weights=[True, False],
    ...     percentile=[10, 20, 30],
    ... )
"""

from typing import Any, Callable, Optional


def create_context_search_space(**kwargs) -> dict:
    """
    Create a hyperopt search space for ANY context parameters.

    Pass lists of values to create hp.choice, or hp.* objects directly.

    Args:
        **kwargs: Parameter names and their possible values.
            - Lists: Converted to hp.choice
            - hp.* objects: Used as-is
            - Single values: Used as-is (no choice)
            - None: Ignored

    Returns:
        dict: Hyperopt search space dictionary

    Examples:
        # Simple choices
        create_context_search_space(
            optimization_metric=["LOGLOSS", "ROC_AUC"],
            use_class_weights=[True, False],
            percentile=[10, 20, 30],
        )

        # With continuous ranges
        from hyperopt import hp
        create_context_search_space(
            val_size=hp.uniform("val_size", 0.1, 0.3),
            max_p_value=hp.loguniform("max_p_value", -5, -1),
        )
    """
    from hyperopt import hp

    space = {}

    for param_name, values in kwargs.items():
        if values is None:
            continue

        # Already a hyperopt expression (has pos_args attribute from hp.* functions)
        if hasattr(values, "pos_args"):
            space[param_name] = values
        # List of choices
        elif isinstance(values, list):
            if len(values) == 0:
                continue  # Skip empty lists
            elif len(values) == 1:
                space[param_name] = values[0]  # Single value, no choice
            else:
                space[param_name] = hp.choice(param_name, values)
        # Single value
        else:
            space[param_name] = values

    return space


def build_context_from_params(
    base_context: dict,
    sampled_params: dict,
    context_builder: Optional[Callable[[dict], dict]] = None,
) -> dict:
    """
    Build a full context dict from base context and sampled hyperparams.

    Args:
        base_context: Base context dictionary with default values
        sampled_params: Parameters sampled from the hyperopt search space
        context_builder: Optional function to post-process the context
            (e.g., compute derived parameters)

    Returns:
        dict: Complete context dictionary ready for LeCrapaud
    """
    context = base_context.copy()
    context.update(sampled_params)

    if context_builder:
        context = context_builder(context)

    return context


def get_metric_direction(metric: str) -> str:
    """
    Get optimization direction for a metric.

    Uses METRICS_CONFIG from utils.py, with additional support for business metrics.

    Args:
        metric: Metric name (case-insensitive)

    Returns:
        str: "minimize" or "maximize"

    Raises:
        ValueError: If metric is unknown
    """
    from lecrapaud.utils import get_metric_direction as utils_get_metric_direction
    from lecrapaud.utils import METRICS_CONFIG

    # Business metrics not in METRICS_CONFIG
    if metric.lower() == "business_cost":
        return "minimize"

    # Try utils version (handles uppercase)
    try:
        return utils_get_metric_direction(metric.upper())
    except ValueError:
        pass

    # Unknown metric
    valid_metrics = list(METRICS_CONFIG.keys()) + ["business_cost"]
    raise ValueError(f"Unknown metric: {metric}. Valid metrics: {valid_metrics}")


def get_default_search_space() -> dict:
    """
    Get a default search space for common context parameters.

    Returns:
        dict: Default search space with common optimization parameters
    """
    from hyperopt import hp

    return create_context_search_space(
        # Feature selection
        percentile=[10, 20, 30, 40],
        corr_threshold=[70, 80, 90],
        cumulative_importance=[0.7, 0.8, 0.9],
        # Model selection
        use_class_weights=[True, False],
        # Data split
        val_size=[0.15, 0.20, 0.25],
    )


def _extract_choices(node) -> list | None:
    """Extract the list of choices from an hp.choice hyperopt node."""
    from hyperopt import pyll

    if not isinstance(node, pyll.base.Apply):
        return None

    # hp.choice nodes have name 'switch' at the top level
    # The structure is: switch(randint, option0, option1, ...)
    if node.name == "switch":
        choices = []
        # Skip first arg (the randint index sampler), rest are the choices
        for arg in node.pos_args[1:]:
            choices.append(_extract_literal(arg))
        return choices

    # Walk children to find a switch node (hp.choice wraps in pos_args)
    for arg in node.pos_args:
        if isinstance(arg, pyll.base.Apply):
            result = _extract_choices(arg)
            if result is not None:
                return result

    return None


def _extract_literal(node):
    """Extract a literal Python value from a hyperopt pyll node."""
    from hyperopt import pyll

    if isinstance(node, pyll.base.Literal):
        return node.obj

    if isinstance(node, pyll.base.Apply):
        # Try to recursively extract from pos_args/literal nodes
        if node.name == "pos_args":
            return [_extract_literal(a) for a in node.pos_args]
        if node.name == "dict":
            # Reconstruct dict from named_args
            return {k: _extract_literal(v) for k, v in node.named_args}

    return repr(node)


def describe_search_space(space: dict) -> str:
    """
    Return a human-readable JSON description of a hyperopt search space.

    Extracts choices directly from hp.choice nodes for clear display.

    Args:
        space: Hyperopt search space dictionary

    Returns:
        str: JSON-formatted search space description
    """
    import json
    from hyperopt import pyll

    result = {}
    for param_name, param_space in space.items():
        if isinstance(param_space, pyll.base.Apply):
            choices = _extract_choices(param_space)
            if choices is not None:
                result[param_name] = choices
            else:
                result[param_name] = repr(param_space)
        else:
            result[param_name] = param_space

    return json.dumps(result, indent=2, default=str)


def print_search_space(space: dict) -> None:
    """Print a human-readable description of a hyperopt search space."""
    print(describe_search_space(space))


def merge_search_spaces(*spaces: dict) -> dict:
    """
    Merge multiple search spaces into one.

    Later spaces override earlier ones for conflicting keys.

    Args:
        *spaces: Search space dictionaries to merge

    Returns:
        dict: Merged search space
    """
    result = {}
    for space in spaces:
        result.update(space)
    return result
