from __future__ import annotations

from dataclasses import dataclass
from typing import Any


def _require(d: dict[str, Any], key: str) -> Any:
    if key not in d or d[key] in (None, ""):
        raise ValueError(f"missing required field: {key}")
    return d[key]


@dataclass(frozen=True)
class RayConfig:
    address: str
    shared_root: str
    entrypoint_num_cpus: float
    entrypoint_resources: dict[str, float]
    runtime_env_env_vars: dict[str, str]
    user_code_path: str
    verl_code_path: str

    @staticmethod
    def from_dict(d: dict[str, Any]) -> "RayConfig":
        # New format: root.ray
        # Backward compatible: root (flat)
        if isinstance(d.get("ray"), dict):
            d = d["ray"]  # type: ignore[assignment]

        runtime_env = d.get("runtime_env") or {}
        env_vars = (runtime_env.get("env_vars") or {}) if isinstance(runtime_env, dict) else {}
        if not isinstance(env_vars, dict):
            raise ValueError("runtime_env.env_vars must be a mapping")

        entrypoint_resources = d.get("entrypoint_resources") or {}
        if not isinstance(entrypoint_resources, dict):
            raise ValueError("entrypoint_resources must be a mapping")

        return RayConfig(
            address=str(_require(d, "address")),
            shared_root=str(_require(d, "shared_root")),
            entrypoint_num_cpus=float(d.get("entrypoint_num_cpus", 1)),
            entrypoint_resources={str(k): float(v) for k, v in entrypoint_resources.items()},
            runtime_env_env_vars={str(k): str(v) for k, v in env_vars.items()},
            user_code_path=str(d.get("user_code_path", f"{_require(d, 'shared_root')}/user/code")),
            verl_code_path=str(d.get("verl_code_path", f"{_require(d, 'shared_root')}/common/code/verl/verl_repo")),
        )

    def to_public_dict(self) -> dict[str, Any]:
        return {
            "address": self.address,
            "shared_root": self.shared_root,
            "entrypoint_num_cpus": self.entrypoint_num_cpus,
            "entrypoint_resources": self.entrypoint_resources,
            "runtime_env": {"env_vars": self.runtime_env_env_vars},
            "user_code_path": self.user_code_path,
            "verl_code_path": self.verl_code_path,
        }


@dataclass(frozen=True)
class JobSpec:
    workload: str  # ppo|grpo|sft
    submission_id: str | None
    code_path: str
    model_id: str

    train_file: str
    val_file: str

    nnodes: int
    n_gpus_per_node: int

    total_epochs: int
    total_training_steps: int | None

    save_freq: int
    test_freq: int | None

    trainer_device: str | None  # only for sft (driver-side device)

    @staticmethod
    def from_dict(d: dict[str, Any]) -> "JobSpec":
        workload = str(_require(d, "workload"))
        if workload not in ("ppo", "grpo", "sft"):
            raise ValueError(f"unsupported workload: {workload}")

        test_freq = d.get("test_freq", None)
        if test_freq in ("", "null"):
            test_freq = None

        total_training_steps = d.get("total_training_steps", None)
        if total_training_steps in ("", "null"):
            total_training_steps = None

        return JobSpec(
            workload=workload,
            submission_id=(str(d["submission_id"]) if d.get("submission_id") else None),
            code_path=str(_require(d, "code_path")),
            model_id=str(_require(d, "model_id")),
            train_file=str(_require(d, "train_file")),
            val_file=str(_require(d, "val_file")),
            nnodes=int(d.get("nnodes", 2)),
            n_gpus_per_node=int(d.get("n_gpus_per_node", 4)),
            total_epochs=int(d.get("total_epochs", 1)),
            # If omitted or explicitly null: do not override trainer.total_training_steps
            # (let VERL derive it from epochs and dataset length).
            total_training_steps=(int(total_training_steps) if total_training_steps is not None else None),
            save_freq=int(d.get("save_freq", 10)),
            test_freq=(int(test_freq) if test_freq is not None else None),
            trainer_device=(str(d.get("trainer_device")) if d.get("trainer_device") else None),
        )

    def to_public_dict(self) -> dict[str, Any]:
        out: dict[str, Any] = {
            "workload": self.workload,
            "submission_id": self.submission_id or "",
            "code_path": self.code_path,
            "model_id": self.model_id,
            "train_file": self.train_file,
            "val_file": self.val_file,
            "nnodes": self.nnodes,
            "n_gpus_per_node": self.n_gpus_per_node,
            "total_epochs": self.total_epochs,
            "total_training_steps": self.total_training_steps,
            "save_freq": self.save_freq,
            "test_freq": self.test_freq,
        }
        if self.workload == "sft":
            out["trainer_device"] = self.trainer_device or "cpu"
        return out


@dataclass(frozen=True)
class AdvancedTaskSpec:
    kind: str  # advanced
    workload: str  # always "advanced" for classification/ID prefix
    submission_id: str | None
    nnodes: int
    n_gpus_per_node: int
    command: str

    @staticmethod
    def from_dict(d: dict[str, Any]) -> "AdvancedTaskSpec":
        kind = str(_require(d, "kind"))
        if kind != "advanced":
            raise ValueError(f"unsupported taskspec kind: {kind}")

        command = str(_require(d, "command"))

        return AdvancedTaskSpec(
            kind=kind,
            workload="advanced",
            submission_id=(str(d["submission_id"]) if d.get("submission_id") else None),
            nnodes=int(_require(d, "nnodes")),
            n_gpus_per_node=int(_require(d, "n_gpus_per_node")),
            command=command,
        )

    def to_public_dict(self) -> dict[str, Any]:
        return {
            "kind": self.kind,
            "workload": self.workload,
            "submission_id": self.submission_id or "",
            "nnodes": self.nnodes,
            "n_gpus_per_node": self.n_gpus_per_node,
            "command": self.command,
        }


def parse_taskspec(d: dict[str, Any]) -> JobSpec | AdvancedTaskSpec:
    """
    v3.x compatibility rule:
    - If no "kind": treat as legacy/basic JobSpec.
    - If kind == "advanced": parse AdvancedTaskSpec.
    """

    kind = d.get("kind")
    if kind in (None, ""):
        return JobSpec.from_dict(d)
    if kind == "advanced":
        return AdvancedTaskSpec.from_dict(d)
    raise ValueError(f"unsupported taskspec kind: {kind}")
