"""Cloud training command generators for GCP, AWS, RunPod, and Lambda."""

from __future__ import annotations

from dataclasses import dataclass


@dataclass
class CloudTrainingConfig:
    """Cloud-specific training configuration."""

    cloud: str  # "gcp", "aws", "runpod", "lambda"
    command: str  # Full CLI command or code snippet
    estimated_cost: str  # "~$5-15"
    gpu_type: str  # "a100", "h100", "l4"
    notes: list[str]


def generate_cloud_command(
    dataset_repo: str,
    policy_type: str,
    cloud: str,
    gpu_type: str = "a100",
    gpu_count: int = 1,
    steps: int = 100_000,
    batch_size: int = 8,
) -> CloudTrainingConfig:
    """Generate cloud training launch commands."""
    generators = {
        "gcp": _gcp_command,
        "aws": _aws_command,
        "runpod": _runpod_command,
        "lambda": _lambda_command,
    }
    gen = generators.get(cloud)
    if gen is None:
        raise ValueError(f"Unknown cloud: {cloud}. Choose from: {list(generators)}")
    return gen(
        dataset_repo=dataset_repo,
        policy_type=policy_type,
        gpu_type=gpu_type,
        gpu_count=gpu_count,
        steps=steps,
        batch_size=batch_size,
    )


def _gcp_command(
    dataset_repo: str,
    policy_type: str,
    gpu_type: str,
    gpu_count: int,
    **kwargs,
) -> CloudTrainingConfig:
    gcp_gpu_map = {"a100": "nvidia-a100-80gb", "h100": "nvidia-h100-80gb", "l4": "nvidia-l4"}
    gcp_gpu = gcp_gpu_map.get(gpu_type, "nvidia-l4")
    command = (
        f"gcloud run jobs create train-robot \\\n"
        f"  --image=orbit-robotics:latest \\\n"
        f"  --gpu={gpu_count} --gpu-type={gcp_gpu} \\\n"
        f"  --memory=32Gi \\\n"
        f'  --command="orbit train {dataset_repo} --policy {policy_type}" \\\n'
        f"  --region=us-central1"
    )
    return CloudTrainingConfig(
        cloud="gcp",
        command=command,
        estimated_cost="~$3-10/h per A100" if "a100" in gpu_type else "~$1-3/h per L4",
        gpu_type=gpu_type,
        notes=[
            "Requires gcloud CLI authenticated with a project",
            "Cloud Run Jobs supports GPU workloads",
            "Dataset must be accessible from the container (GCS bucket or HuggingFace Hub)",
        ],
    )


def _aws_command(
    dataset_repo: str,
    policy_type: str,
    gpu_type: str,
    gpu_count: int,
    steps: int,
    batch_size: int,
    **kwargs,
) -> CloudTrainingConfig:
    instance_map = {"a100": "ml.p4d.24xlarge", "h100": "ml.p5.48xlarge", "l4": "ml.g5.xlarge"}
    instance = instance_map.get(gpu_type, "ml.g5.xlarge")
    command = (
        f"# SageMaker training job\n"
        f"from sagemaker.pytorch import PyTorch\n\n"
        f"estimator = PyTorch(\n"
        f"    entry_point='train.py',\n"
        f"    instance_type='{instance}',\n"
        f"    instance_count={gpu_count},\n"
        f"    framework_version='2.1',\n"
        f"    hyperparameters={{\n"
        f"        'dataset': '{dataset_repo}',\n"
        f"        'policy': '{policy_type}',\n"
        f"        'steps': {steps},\n"
        f"        'batch_size': {batch_size},\n"
        f"    }},\n"
        f")\n"
        f"estimator.fit()"
    )
    return CloudTrainingConfig(
        cloud="aws",
        command=command,
        estimated_cost="~$4-12/h per A100 (p4d)" if "a100" in gpu_type else "~$1-3/h (g5)",
        gpu_type=gpu_type,
        notes=[
            "Requires AWS credentials and SageMaker setup",
            "Dataset should be in S3 or accessible from HuggingFace Hub",
            "Spot instances can reduce cost by 60-70%",
        ],
    )


def _runpod_command(
    dataset_repo: str,
    policy_type: str,
    gpu_type: str,
    **kwargs,
) -> CloudTrainingConfig:
    gpu_name_map = {
        "a100": "NVIDIA A100 80GB",
        "h100": "NVIDIA H100 80GB",
        "l4": "NVIDIA L4",
    }
    gpu_name = gpu_name_map.get(gpu_type, "NVIDIA A100 80GB")
    command = (
        f"runpodctl create pod \\\n"
        f'  --name train-robot \\\n'
        f'  --gpuType "{gpu_name}" \\\n'
        f"  --imageName orbit-robotics:latest \\\n"
        f'  --args "orbit train {dataset_repo} --policy {policy_type}"'
    )
    return CloudTrainingConfig(
        cloud="runpod",
        command=command,
        estimated_cost="~$1.5-3/h per A100 (community cloud)",
        gpu_type=gpu_type,
        notes=[
            "RunPod offers the cheapest GPU cloud for training",
            "Install: pip install runpodctl",
            "Community cloud GPUs are cheapest but less reliable",
        ],
    )


def _lambda_command(
    dataset_repo: str,
    policy_type: str,
    gpu_type: str,
    **kwargs,
) -> CloudTrainingConfig:
    instance_map = {
        "a100": "gpu_1x_a100_sxm4",
        "h100": "gpu_1x_h100_sxm5",
        "l4": "gpu_1x_a10",
    }
    instance = instance_map.get(gpu_type, "gpu_1x_a100_sxm4")
    command = (
        f"# Lambda Cloud instance\n"
        f"lambda-cloud launch \\\n"
        f"  --instance-type {instance} \\\n"
        f"  --ssh-key my-key \\\n"
        f"  --name train-robot\n\n"
        f"# Then SSH in and run:\n"
        f"pip install orbit-robotics lerobot\n"
        f"orbit train {dataset_repo} --policy {policy_type}"
    )
    return CloudTrainingConfig(
        cloud="lambda",
        command=command,
        estimated_cost="~$1-2/h per A100",
        gpu_type=gpu_type,
        notes=[
            "Lambda Cloud offers on-demand A100/H100 instances",
            "No container runtime — SSH in and run directly",
            "Often cheapest for long training runs",
        ],
    )
