from __future__ import annotations

from typing import Any, Optional

import numpy as np


def _as_discrete_action(action: Any) -> int:
    if isinstance(action, np.ndarray):
        return int(action.item())
    if isinstance(action, np.generic):
        return int(action.item())
    try:
        return int(action)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"Cannot convert action to int: {action!r}") from exc


def select_action(
    model: Any,
    obs: Any,
    deterministic: bool,
    epsilon: float = 0.0,
    epsilon_random_mode: str = "uniform",
    env: Optional[Any] = None,
    rng: Optional[np.random.Generator] = None,
) -> int:
    """Select action with optional epsilon-greedy for deterministic inference.

    When deterministic=False (stochastic policy), epsilon is ignored by design.
    """

    if epsilon < 0.0 or epsilon > 1.0:
        raise ValueError(f"epsilon must be in [0, 1], got {epsilon}")

    mode = epsilon_random_mode.lower().strip()
    if mode not in {"uniform", "policy"}:
        raise ValueError(f"Unsupported epsilon_random_mode={epsilon_random_mode!r}. Use 'uniform' or 'policy'.")

    effective_epsilon = epsilon if deterministic else 0.0
    if effective_epsilon > 0.0:
        draw = float(rng.random()) if rng is not None else float(np.random.random())
        if draw < effective_epsilon:
            if mode == "uniform":
                if env is None or not hasattr(env, "action_space"):
                    raise ValueError("uniform epsilon_random_mode requires env.action_space.sample().")
                return _as_discrete_action(env.action_space.sample())
            action, _ = model.predict(obs, deterministic=False)
            return _as_discrete_action(action)

    action, _ = model.predict(obs, deterministic=deterministic)
    return _as_discrete_action(action)

