from __future__ import annotations

from typing import Any

from .serving_spec import ResolvedServingSpec


def _ensure_hf_env_defaults(env: dict[str, str]) -> dict[str, str]:
    out = dict(env or {})
    # Prefer existing values if present, but always force offline mode in the platform.
    out.setdefault("HF_HOME", "/private/hf")
    out.setdefault("HUGGINGFACE_HUB_CACHE", "/private/hf/hub")
    out.setdefault("TRANSFORMERS_CACHE", "/private/hf/transformers")
    out["HF_HUB_OFFLINE"] = "1"
    return out


def build_llm_config_dict(
    resolved: ResolvedServingSpec,
    *,
    accelerator_type: str,
    runtime_env_env_vars: dict[str, str] | None,
    cpu_per_gpu: float = 1.0,
) -> dict[str, Any]:
    """
    Pure builder: maps a platform ResolvedServingSpec to a Ray Serve LLM-like config.

    We return a plain dict here to keep this layer unit-testable without depending on
    a specific Ray Serve LLM version. The reconciler (later milestone) can choose to
    instantiate `ray.serve.llm.LLMConfig` using this dict.
    """
    if not accelerator_type:
        raise ValueError("accelerator_type is required")
    if resolved.num_replicas < 1:
        raise ValueError("num_replicas must be >= 1")
    if resolved.gpus_per_replica < 1:
        raise ValueError("gpus_per_replica must be >= 1")
    if cpu_per_gpu <= 0:
        raise ValueError("cpu_per_gpu must be > 0")

    engine_kwargs: dict[str, Any] = dict(resolved.engine_kwargs or {})
    # Enforce tensor parallel mapping; user-provided value must not contradict requested GPUs.
    engine_kwargs["tensor_parallel_size"] = int(resolved.gpus_per_replica)

    # Ray Serve LLM (Ray 2.49.x) exposes `resources_per_bundle` instead of the older
    # `placement_group_config`. Use a single bundle that reserves the full GPU set
    # required by tensor-parallel execution.
    resources_per_bundle = {
        "GPU": float(resolved.gpus_per_replica),
        "CPU": float(cpu_per_gpu) * float(resolved.gpus_per_replica),
    }

    env_vars = _ensure_hf_env_defaults(dict(runtime_env_env_vars or {}))

    return {
        # Ray Serve LLM expects `model_loading_config` with model_id/model_source.
        "model_loading_config": {"model_id": resolved.model_id, "model_source": resolved.model_source},
        "accelerator_type": accelerator_type,
        "deployment_config": {"num_replicas": int(resolved.num_replicas)},
        "engine_kwargs": engine_kwargs,
        "resources_per_bundle": resources_per_bundle,
        "runtime_env": {"env_vars": env_vars},
    }
