"""Base training configuration and framework registry.

Each framework generator produces the ACTUAL command that the framework
expects, validated against the current (April 2025) CLI of each project.
"""

from __future__ import annotations

from dataclasses import dataclass, field


@dataclass
class TrainingConfig:
    """Framework-agnostic training configuration."""

    framework: str  # "lerobot", "openvla", "openpi", "groot", "osmo", "custom"
    command: str  # Full CLI command or YAML content
    estimated_time: str  # "~3h on 1x A100"
    gpu_requirements: str  # "1x GPU with 24GB+ VRAM"
    dataset_format: str  # "lerobot-v3", "lerobot-v2", "rlds", "hdf5"
    conversion_needed: bool  # True if orbit convert is needed first
    conversion_command: str  # "orbit convert ./data --to lerobot-v2 -o ./converted/"
    pre_checks: list[str] = field(default_factory=list)
    tips: list[str] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)
    # Generated config files the user needs to create (path -> content)
    generated_files: dict[str, str] = field(default_factory=dict)


# Registry of available frameworks
FRAMEWORKS = ["lerobot", "openvla", "openpi", "groot", "osmo", "custom"]


def get_framework_config(
    framework: str,
    dataset_repo: str,
    policy_type: str,
    info,
    episode_count: int,
    action_dim: int,
    camera_count: int,
    resolution: str,
    gpu_memory_gb: int | None = None,
    batch_size: int = 8,
    steps: int = 100_000,
    learning_rate: float = 1e-4,
    chunk_size: int = 50,
    fps: float = 30.0,
) -> TrainingConfig:
    """Generate a TrainingConfig for the given framework.

    Args:
        framework: One of FRAMEWORKS.
        dataset_repo: Dataset path or HuggingFace repo ID.
        policy_type: Policy architecture name.
        info: DatasetInfo object.
        episode_count: Number of episodes.
        action_dim: Action dimensionality.
        camera_count: Number of cameras.
        resolution: Image resolution bucket ("480p", "720p", "1080p").
        gpu_memory_gb: Available GPU memory.
        batch_size: Batch size (from LeRobot suggester).
        steps: Training steps.
        learning_rate: Learning rate.
        chunk_size: Action chunk size.
        fps: Dataset FPS.

    Returns:
        TrainingConfig for the specified framework.
    """
    generators = {
        "lerobot": _lerobot_config,
        "openvla": _openvla_config,
        "openpi": _openpi_config,
        "groot": _groot_config,
        "osmo": _osmo_config,
        "custom": _custom_config,
    }
    gen = generators.get(framework)
    if gen is None:
        raise ValueError(f"Unknown framework: {framework}. Choose from: {FRAMEWORKS}")
    return gen(
        dataset_repo=dataset_repo,
        policy_type=policy_type,
        info=info,
        episode_count=episode_count,
        action_dim=action_dim,
        camera_count=camera_count,
        resolution=resolution,
        gpu_memory_gb=gpu_memory_gb,
        batch_size=batch_size,
        steps=steps,
        learning_rate=learning_rate,
        chunk_size=chunk_size,
        fps=fps,
    )


# ---------------------------------------------------------------------------
# LeRobot
# ---------------------------------------------------------------------------


def _lerobot_config(
    dataset_repo: str,
    policy_type: str,
    batch_size: int,
    steps: int,
    learning_rate: float,
    chunk_size: int,
    episode_count: int,
    **kwargs,
) -> TrainingConfig:
    output_dir = f"outputs/train/{dataset_repo.split('/')[-1]}-{policy_type}"
    parts = [
        "lerobot-train",
        f"  --dataset.repo_id={dataset_repo}",
        f"  --policy.type={policy_type}",
        f"  --batch_size={batch_size}",
        f"  --steps={steps}",
    ]
    if policy_type == "act":
        parts.append(f"  --policy.chunk_size={chunk_size}")
        parts.append(f"  --policy.n_action_steps={chunk_size}")
    elif policy_type == "diffusion_policy":
        parts.append(f"  --policy.horizon={chunk_size}")

    parts.extend([
        "  --log_freq=200",
        f"  --save_freq={max(1000, steps // 5)}",
        f"  --eval_freq={max(1000, steps // 5)}",
        "  --policy.device=cuda",
        f"  --output_dir={output_dir}",
    ])
    command = " \\\n".join(parts)

    time_est = "~2-4h on 1x A100" if episode_count < 200 else "~6-12h on 1x A100"

    return TrainingConfig(
        framework="lerobot",
        command=command,
        estimated_time=time_est,
        gpu_requirements="1x GPU with 16GB+ VRAM (24GB+ recommended)",
        dataset_format="lerobot-v3",
        conversion_needed=False,
        conversion_command="",
        pre_checks=[f"orbit gate {dataset_repo} -p {policy_type} --min-grade C"],
        tips=[
            f"Loss should drop below 0.1 by step {steps // 5}",
            "If loss plateaus above 0.2: demonstrations may be too inconsistent",
            "If robot doesn't move during eval: check action normalization stats",
        ],
        warnings=[],
    )


# ---------------------------------------------------------------------------
# OpenVLA (OFT recommended, LoRA as alternative)
# Source: https://github.com/openvla/openvla-oft
# ---------------------------------------------------------------------------


def _openvla_config(
    dataset_repo: str,
    gpu_memory_gb: int | None,
    episode_count: int,
    action_dim: int,
    camera_count: int,
    info,
    **kwargs,
) -> TrainingConfig:
    gpu_gb = gpu_memory_gb or 80
    dataset_name = dataset_repo.split("/")[-1]

    # OFT (Open Fine-Tuning) is the recommended path as of March 2025
    # Falls back to original LoRA for lower VRAM
    if gpu_gb >= 40:
        # OFT fine-tuning (recommended)
        command = (
            "torchrun --standalone --nnodes 1 --nproc-per-node 1 \\\n"
            "  vla-scripts/finetune.py \\\n"
            "  --vla_path openvla/openvla-7b-oft \\\n"
            f"  --data_root_dir ./data \\\n"
            f"  --dataset_name {dataset_name} \\\n"
            "  --run_root_dir outputs/openvla-oft \\\n"
            "  --adapter_tmp_dir ./adapter_tmp \\\n"
            "  --lora_rank 32 \\\n"
            "  --batch_size 16 \\\n"
            "  --learning_rate 2e-5 \\\n"
            "  --image_aug true \\\n"
            "  --save_steps 1000"
        )
        gpu_req = f"1x GPU with {gpu_gb}GB+ VRAM (OFT fine-tuning)"
        time_est = "~8-12h on 1x A100"
        tips = [
            "OFT (Open Fine-Tuning) produces faster, more successful policies than base OpenVLA",
            "OFT is the recommended fine-tuning recipe as of March 2025",
            "LoRA rank 32 is a good default — increase to 64 for more complex tasks",
            "Image augmentation (--image_aug true) significantly improves generalization",
        ]
    else:
        # Original LoRA for lower VRAM
        command = (
            "torchrun --standalone --nnodes 1 --nproc-per-node 1 \\\n"
            "  vla-scripts/finetune.py \\\n"
            "  --vla_path openvla/openvla-7b \\\n"
            f"  --data_root_dir ./data \\\n"
            f"  --dataset_name {dataset_name} \\\n"
            "  --run_root_dir outputs/openvla \\\n"
            "  --lora_rank 32 \\\n"
            "  --batch_size 8 \\\n"
            "  --learning_rate 2e-5 \\\n"
            "  --image_aug true"
        )
        gpu_req = "1x GPU with 24GB+ VRAM (LoRA fine-tuning)"
        time_est = "~10-15h on 1x A100"
        tips = [
            "Consider upgrading to OFT (openvla-7b-oft) for better results if you have 40GB+ VRAM",
            "LoRA fine-tuning preserves general capabilities while adapting to your task",
            "Image augmentation (--image_aug true) significantly improves generalization",
        ]

    # OpenVLA uses RLDS format
    needs_convert = True
    convert_cmd = f"orbit convert {dataset_repo} --to rlds -o ./data/"

    return TrainingConfig(
        framework="openvla",
        command=command,
        estimated_time=time_est,
        gpu_requirements=gpu_req,
        dataset_format="rlds",
        conversion_needed=needs_convert,
        conversion_command=convert_cmd,
        pre_checks=[f"orbit gate {dataset_repo} -p act --min-grade C"],
        tips=tips,
        warnings=[
            "Requires RLDS dataset format — run the conversion command first",
            "OpenVLA requires language task descriptions — ensure your data includes task prompts",
        ],
    )


# ---------------------------------------------------------------------------
# OpenPI / pi0
# Source: https://github.com/Physical-Intelligence/openpi
# Requires creating a Python TrainConfig class, not just a CLI flag.
# ---------------------------------------------------------------------------


def _openpi_config(
    dataset_repo: str,
    episode_count: int,
    action_dim: int,
    camera_count: int,
    info,
    **kwargs,
) -> TrainingConfig:
    dataset_name = dataset_repo.split("/")[-1].replace("-", "_")
    class_name = "".join(w.capitalize() for w in dataset_name.split("_")) + "Config"

    # Generate the TrainConfig Python file template
    camera_keys = []
    if hasattr(info, "camera_names") and info.camera_names:
        camera_keys = info.camera_names
    if not camera_keys:
        camera_keys = ["observation.image"] if camera_count > 0 else []

    state_dim = action_dim  # often same; user should verify

    config_template = (
        f'"""Training config for {dataset_name}.\n'
        f'\n'
        f'Generated by: orbit suggest {dataset_repo} --framework openpi\n'
        f'Place this file in: openpi/configs/{dataset_name}_config.py\n'
        f'"""\n'
        f'\n'
        f'import dataclasses\n'
        f'\n'
        f'from openpi.training.config import TrainConfig\n'
        f'from openpi.models.pi0 import Pi0Config\n'
        f'\n'
        f'\n'
        f'@dataclasses.dataclass\n'
        f'class {class_name}(TrainConfig):\n'
        f'    # Model\n'
        f'    model: Pi0Config = dataclasses.field(\n'
        f'        default_factory=lambda: Pi0Config(\n'
        f'            action_dim={action_dim},\n'
        f'            action_horizon=50,\n'
        f'        )\n'
        f'    )\n'
        f'\n'
        f'    # Dataset\n'
        f'    dataset_name: str = "{dataset_name}"\n'
        f'    data_path: str = "{dataset_repo}"\n'
        f'\n'
        f'    # Camera keys from your dataset\n'
        f'    image_keys: tuple[str, ...] = {tuple(camera_keys)}\n'
        f'\n'
        f'    # Action and state dimensions\n'
        f'    action_dim: int = {action_dim}\n'
        f'    state_dim: int = {state_dim}\n'
        f'\n'
        f'    # Training\n'
        f'    batch_size: int = 64\n'
        f'    num_train_steps: int = {max(10_000, episode_count * 200)}\n'
        f'    learning_rate: float = 1e-4\n'
        f'\n'
        f'    # Task prompt (customize for your task)\n'
        f'    default_prompt: str = "pick up the object"  # CHANGE THIS\n'
    )

    command = (
        f"# Step 1: Place config file (generated below) in openpi/configs/\n"
        f"# Step 2: Run training:\n"
        f"XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run scripts/train.py \\\n"
        f"  {dataset_name}_config \\\n"
        f"  --exp-name={dataset_name}_experiment \\\n"
        f"  --overwrite"
    )

    config_path = f"openpi/configs/{dataset_name}_config.py"

    return TrainingConfig(
        framework="openpi",
        command=command,
        estimated_time="~1-20h on 1x A100 (depends on data size)",
        gpu_requirements="1x A100 80GB+ recommended",
        dataset_format="lerobot-v3",
        conversion_needed=False,
        conversion_command="",
        pre_checks=[f"orbit gate {dataset_repo} -p act --min-grade C"],
        tips=[
            f"Save the config template to {config_path} and customize it",
            "Edit default_prompt to match your task description",
            "pi0 excels at dexterous manipulation with language conditioning",
            "pi0-FAST variant is faster for inference — use for real-time control",
        ],
        warnings=[
            f"You MUST create {config_path} before training — see generated template",
            "pi0 requires JAX — install with: pip install jax[cuda12]",
        ],
        generated_files={config_path: config_template},
    )


# ---------------------------------------------------------------------------
# GR00T N1.6
# Source: https://github.com/NVIDIA/Isaac-GR00T
# Actual command: python gr00t/experiment/launch_finetune.py
# Requires: LeRobot v2 format, modality config Python file, modality.json
# ---------------------------------------------------------------------------


def _groot_config(
    dataset_repo: str,
    episode_count: int,
    action_dim: int,
    info,
    camera_count: int,
    **kwargs,
) -> TrainingConfig:
    dataset_name = dataset_repo.split("/")[-1]
    robot_type = getattr(info, "robot_type", "so100") or "so100"
    num_gpus = min(8, max(1, episode_count // 25))

    # Actual GR00T N1.6 fine-tuning command
    command = (
        f"python gr00t/experiment/launch_finetune.py \\\n"
        f"  --base-model-path nvidia/GR00T-N1.6-3B \\\n"
        f"  --dataset-path ./{dataset_name} \\\n"
        f"  --embodiment-tag new_embodiment \\\n"
        f"  --modality-config-path ./{robot_type}_config.py \\\n"
        f"  --num-gpus {num_gpus} \\\n"
        f"  --output-dir outputs/groot"
    )

    # Generate modality config Python file
    camera_keys = []
    if hasattr(info, "camera_names") and info.camera_names:
        camera_keys = info.camera_names
    if not camera_keys:
        camera_keys = ["observation.image"] if camera_count > 0 else []

    video_keys_str = ", ".join(f'"{k}"' for k in camera_keys) if camera_keys else '"observation.image"'
    modality_config = (
        f'"""Modality config for {robot_type}.\n'
        f'\n'
        f'Generated by: orbit suggest {dataset_repo} --framework groot\n'
        f'Save as: ./{robot_type}_config.py\n'
        f'"""\n'
        f'\n'
        f'from gr00t.experiment.data_config import ModalityConfig\n'
        f'\n'
        f'\n'
        f'MODALITY_CONFIGS = {{\n'
        f'    "video": ModalityConfig(\n'
        f'        delta=False,\n'
        f'        keys=[{video_keys_str}],\n'
        f'    ),\n'
        f'    "state": ModalityConfig(\n'
        f'        delta=False,\n'
        f'        keys=["observation.state"],\n'
        f'    ),\n'
        f'    "action": ModalityConfig(\n'
        f'        delta=True,  # Use delta actions (relative)\n'
        f'        keys=["action"],\n'
        f'    ),\n'
        f'}}\n'
    )

    # Generate modality.json for dataset meta/ directory
    modality_json = (
        '{\n'
        '    "video": {\n'
        f'        "keys": [{video_keys_str}],\n'
        '        "delta": false\n'
        '    },\n'
        '    "state": {\n'
        '        "keys": ["observation.state"],\n'
        '        "delta": false\n'
        '    },\n'
        '    "action": {\n'
        '        "keys": ["action"],\n'
        '        "delta": true\n'
        '    }\n'
        '}'
    )

    # LeRobot v2 is required (not v3)
    needs_convert = True
    convert_cmd = f"orbit convert {dataset_repo} --to lerobot-v2 -o ./{dataset_name}/"

    return TrainingConfig(
        framework="groot",
        command=command,
        estimated_time=f"~4-8h on {num_gpus}x A100/H100",
        gpu_requirements=f"{num_gpus}x A100/H100 (NVIDIA GPUs required)",
        dataset_format="lerobot-v2",
        conversion_needed=needs_convert,
        conversion_command=convert_cmd,
        pre_checks=[f"orbit gate {dataset_repo} -p act --min-grade C"],
        tips=[
            "GR00T N1.6 requires LeRobot v2 format — NOT v3",
            f"Save the modality config to ./{robot_type}_config.py",
            f"Save modality.json to ./{dataset_name}/meta/modality.json",
            "Set --embodiment-tag to a unique name for your robot",
            "Works best with SO-100, ALOHA, and Franka robot types",
        ],
        warnings=[
            "Dataset must be in LeRobot v2 format — run the conversion command first",
            "Requires modality.json in your dataset's meta/ directory — see generated template",
            f"Requires modality config Python file — see generated {robot_type}_config.py",
            "Foundation model weights require NVIDIA NGC access",
        ],
        generated_files={
            f"{robot_type}_config.py": modality_config,
            f"{dataset_name}/meta/modality.json": modality_json,
        },
    )


# ---------------------------------------------------------------------------
# NVIDIA OSMO
# Source: https://github.com/NVIDIA/OSMO
# Updated YAML schema matching actual OSMO CLI format.
# ---------------------------------------------------------------------------


def _osmo_config(
    dataset_repo: str,
    policy_type: str,
    batch_size: int,
    steps: int,
    **kwargs,
) -> TrainingConfig:
    yaml_content = (
        f"# Generated by: orbit suggest {dataset_repo} --framework osmo\n"
        f"# Run with: osmo run workflow.yaml\n"
        f"workflow:\n"
        f"  tasks:\n"
        f"  - name: orbit-analyze\n"
        f"    image: orbit-robotics:latest\n"
        f"    command: orbit analyze {dataset_repo} --json > /output/analysis.json\n"
        f"    outputs:\n"
        f"    - dataset:\n"
        f"        name: analysis-results\n"
        f"\n"
        f"  - name: orbit-gate\n"
        f"    image: orbit-robotics:latest\n"
        f"    command: orbit gate {dataset_repo} -p {policy_type} --min-grade C\n"
        f"    inputs:\n"
        f"    - task: orbit-analyze\n"
        f"\n"
        f"  - name: train-policy\n"
        f"    image: nvcr.io/nvidia/pytorch:latest\n"
        f"    platform: gb200\n"
        f"    resources:\n"
        f"      gpu: 8\n"
        f"    inputs:\n"
        f"    - task: orbit-gate\n"
        f"    command: >-\n"
        f"      lerobot-train\n"
        f"      --dataset.repo_id={dataset_repo}\n"
        f"      --policy.type={policy_type}\n"
        f"      --batch_size={batch_size}\n"
        f"      --steps={steps}\n"
        f"    outputs:\n"
        f"    - dataset:\n"
        f"        name: trained-model\n"
        f"\n"
        f"  - name: orbit-verify\n"
        f"    image: orbit-robotics:latest\n"
        f"    command: orbit verify /input/trained-model/ -d {dataset_repo}\n"
        f"    inputs:\n"
        f"    - task: train-policy\n"
    )

    return TrainingConfig(
        framework="osmo",
        command=yaml_content,
        estimated_time="Depends on GPU allocation (typically ~2-6h)",
        gpu_requirements="Configured in OSMO workflow (default: 8x GB200)",
        dataset_format="lerobot-v3",
        conversion_needed=False,
        conversion_command="",
        pre_checks=[],
        tips=[
            "Save output to a file: orbit suggest ... --framework osmo > workflow.yaml",
            "Run with: osmo run workflow.yaml",
            "OSMO handles GPU scheduling, data staging, and checkpointing automatically",
            "Modify the platform field to match your cluster (gb200, h100, a100)",
        ],
        warnings=[
            "Requires NVIDIA OSMO CLI and cluster access",
        ],
    )


# ---------------------------------------------------------------------------
# Custom / generic
# ---------------------------------------------------------------------------


def _custom_config(
    dataset_repo: str,
    policy_type: str,
    action_dim: int,
    episode_count: int,
    batch_size: int,
    steps: int,
    learning_rate: float,
    fps: float,
    **kwargs,
) -> TrainingConfig:
    command = (
        f"python train.py \\\n"
        f"  --dataset {dataset_repo} \\\n"
        f"  --policy {policy_type} \\\n"
        f"  --action_dim {action_dim} \\\n"
        f"  --batch_size {batch_size} \\\n"
        f"  --steps {steps} \\\n"
        f"  --lr {learning_rate} \\\n"
        f"  --fps {fps}"
    )

    return TrainingConfig(
        framework="custom",
        command=command,
        estimated_time="Depends on model and hardware",
        gpu_requirements="1x GPU with 16GB+ VRAM recommended",
        dataset_format="varies",
        conversion_needed=False,
        conversion_command="",
        pre_checks=[f"orbit gate {dataset_repo} -p {policy_type} --min-grade C"],
        tips=[
            "Adapt the command above to your training script",
            "Key hyperparameters to tune: batch_size, learning_rate, steps",
            f"With {episode_count} episodes, expect training to take {steps // 1000}K steps",
        ],
        warnings=[],
    )
