from __future__ import annotations

from dataclasses import dataclass

from .models import JobSpec


@dataclass(frozen=True)
class BuiltCommand:
    argv: list[str]


def build_training_argv(spec: JobSpec, submission_id: str, job_dir: str) -> BuiltCommand:
    """
    Returns argv for the actual training process (Hydra overrides preserved).
    This argv is executed by a lightweight Python driver entrypoint.
    """
    if spec.workload in ("ppo", "grpo"):
        algo_overrides: list[str] = []
        if spec.workload == "grpo":
            algo_overrides.append("algorithm.adv_estimator=grpo")

        test_freq = spec.test_freq if spec.test_freq is not None else -1
        maybe_total_steps: list[str] = []
        if spec.total_training_steps is not None:
            maybe_total_steps.append(f"trainer.total_training_steps={spec.total_training_steps}")

        argv = [
            "python3",
            "-m",
            "verl.trainer.main_ppo",
            f"data.train_files={spec.train_file}",
            f"data.val_files={spec.val_file}",
            "data.train_batch_size=256",
            "data.max_prompt_length=512",
            "data.max_response_length=512",
            f"actor_rollout_ref.model.path={spec.model_id}",
            "actor_rollout_ref.actor.optim.lr=1e-6",
            "actor_rollout_ref.actor.ppo_mini_batch_size=64",
            "actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4",
            "actor_rollout_ref.rollout.name=vllm",
            "actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8",
            "actor_rollout_ref.rollout.tensor_model_parallel_size=1",
            "actor_rollout_ref.rollout.gpu_memory_utilization=0.4",
            "actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4",
            "critic.optim.lr=1e-5",
            f"critic.model.path={spec.model_id}",
            "critic.ppo_micro_batch_size_per_gpu=4",
            "algorithm.kl_ctrl.kl_coef=0.001",
            *algo_overrides,
            "trainer.logger=console",
            "trainer.val_before_train=False",
            f"trainer.n_gpus_per_node={spec.n_gpus_per_node}",
            f"trainer.nnodes={spec.nnodes}",
            f"trainer.save_freq={spec.save_freq}",
            f"trainer.test_freq={test_freq}",
            f"trainer.total_epochs={spec.total_epochs}",
            *maybe_total_steps,
            "trainer.resume_mode=disable",
            f"trainer.default_local_dir={job_dir}/checkpoints",
            "+ray_kwargs.ray_init.address=auto",
            f"hydra.run.dir={job_dir}/logs/hydra",
        ]
        return BuiltCommand(argv=argv)

    if spec.workload == "sft":
        trainer_device = spec.trainer_device or "cpu"
        maybe_total_steps: list[str] = []
        if spec.total_training_steps is not None:
            maybe_total_steps.append(f"trainer.total_training_steps={spec.total_training_steps}")

        argv = [
            "python3",
            "-m",
            "verl.trainer.sft_trainer_ray",
            f"model.path={spec.model_id}",
            f"data.train_files={spec.train_file}",
            f"data.val_files={spec.val_file}",
            "data.train_batch_size=64",
            "data.micro_batch_size_per_gpu=1",
            "data.max_token_len_per_gpu=2048",
            "data.max_length=1024",
            "trainer.logger=console",
            "trainer.project_name=mvp11-sft",
            f"trainer.experiment_name={submission_id}",
            f"trainer.total_epochs={spec.total_epochs}",
            *maybe_total_steps,
            f"trainer.save_freq={spec.save_freq}",
            "trainer.test_freq=-1",
            "trainer.resume_mode=disable",
            f"trainer.device={trainer_device}",
            f"trainer.default_local_dir={job_dir}/checkpoints",
            f"trainer.nnodes={spec.nnodes}",
            f"trainer.n_gpus_per_node={spec.n_gpus_per_node}",
            f"hydra.run.dir={job_dir}/logs/hydra",
        ]
        return BuiltCommand(argv=argv)

    raise ValueError(f"unsupported workload: {spec.workload}")


def build_advanced_argv(command: str) -> BuiltCommand:
    """
    Execute an arbitrary user-provided command under bash for maximum shell compatibility
    (multiline, env vars, line continuations, etc.).
    """

    return BuiltCommand(argv=["bash", "-lc", command])
