"""Environment registry — resolve environment names to factory functions.

Usage:
    from wafer_core.environments.registry import get_environment_factory, list_environments

    factory = get_environment_factory("coding")
    env = factory(working_dir=Path.cwd(), enabled_tools=["bash", "read"])

Environments are registered by name. Each factory takes **kwargs and returns
an Environment instance. The kwargs come from CLI --args (key=value pairs).

To add a new environment:
    1. Add entry to ENVIRONMENTS dict below
    2. Factory is imported lazily to avoid heavy deps at CLI startup
"""

from __future__ import annotations

from collections.abc import Callable
from pathlib import Path
from typing import Any

from wafer_core.rollouts.dtypes import Environment

# Factory type: takes string kwargs from CLI --args, returns an Environment
EnvironmentFactory = Callable[..., Environment]


def _coding_factory(**kwargs: Any) -> Environment:  # noqa: ANN401
    from wafer_core.environments.coding import CodingEnvironment
    from wafer_core.sandbox import SandboxMode

    return CodingEnvironment(
        working_dir=Path(kwargs.get("dir", ".")).resolve(),
        enabled_tools=kwargs.get("tools", "").split(",") if kwargs.get("tools") else None,
        sandbox_mode=SandboxMode.DISABLED if kwargs.get("no_sandbox") else SandboxMode.ENABLED,
        allow_network=bool(kwargs.get("allow_network", False)),
    )


def _remote_target_factory(**kwargs: Any) -> Environment:  # noqa: ANN401
    from wafer_core.environments.remote_target import RemoteTargetEnvironment

    target = kwargs.get("target")
    if not target:
        raise ValueError("remote-target environment requires --args target=<name>")

    return RemoteTargetEnvironment(
        target_name=target,
        working_dir=Path(kwargs.get("dir", ".")).resolve(),
        remote_working_dir=kwargs.get("remote_dir", "/tmp/wafer-workspace"),
        bash_timeout=int(kwargs.get("bash_timeout", "300")),
    )


def _runpod_sandbox_factory(**kwargs: Any) -> Environment:  # noqa: ANN401
    try:
        from optimize_kernelbench_eval.environments.pooled_runpod import (
            PooledRunPodEnvironment,
        )
    except ImportError as e:
        raise ImportError(
            "runpod-sandbox environment requires optimize_kernelbench_eval package. "
            "Install with: uv sync --extra evals"
        ) from e

    return PooledRunPodEnvironment(
        provider=kwargs.get("provider", "runpod"),
        working_dir=Path(kwargs.get("dir", ".")).resolve(),
        gpu=kwargs.get("gpu"),
        keep=kwargs.get("keep", "true").lower() != "false",
    )


def _modal_sandbox_factory(**kwargs: Any) -> Environment:  # noqa: ANN401
    from wafer_core.environments.modal_sandbox import (
        ModalSandboxConfig,
        ModalSandboxEnvironment,
    )

    config = ModalSandboxConfig(
        gpu=kwargs.get("gpu", "A100"),
        timeout_seconds=int(kwargs.get("timeout", "600")),
        app_name=kwargs.get("app_name", "kernelbench-eval"),
    )

    return ModalSandboxEnvironment(
        config=config,
        working_dir=Path(kwargs.get("dir", ".")).resolve(),
    )


# name → (factory_fn, description, needs_asyncio_bridge)
ENVIRONMENTS: dict[str, tuple[EnvironmentFactory, str, bool]] = {
    "coding": (_coding_factory, "Local filesystem with bash/read/write/edit/glob/grep", False),
    "remote-target": (_remote_target_factory, "Proxy tools to remote GPU via SSH", False),
    "runpod-sandbox": (
        _runpod_sandbox_factory,
        "RunPod GPU instance via SSH (provisions on start)",
        False,
    ),
    "modal-sandbox": (_modal_sandbox_factory, "Modal Sandbox with GPU (CUDA)", True),
}


def get_environment_factory(name: str) -> EnvironmentFactory:
    """Get factory function for named environment."""
    entry = ENVIRONMENTS.get(name)
    if entry is None:
        available = ", ".join(sorted(ENVIRONMENTS))
        raise ValueError(f"Unknown environment: {name!r}. Available: {available}")
    return entry[0]


def needs_asyncio_bridge(name: str) -> bool:
    """Check if environment requires trio_asyncio bridge for asyncio SDK calls."""
    entry = ENVIRONMENTS.get(name)
    if entry is None:
        return False
    return entry[2]


def list_environments() -> list[tuple[str, str]]:
    """Return list of (name, description) for all registered environments."""
    return [(name, desc) for name, (_, desc, _) in sorted(ENVIRONMENTS.items())]
