"""Training framework knowledge base.

Verified CLI syntax, requirements, and best practices for every
supported robotics training framework.
"""

from __future__ import annotations

FRAMEWORK_KNOWLEDGE: dict[str, dict] = {
    "lerobot": {
        "name": "LeRobot",
        "policies": ["act", "diffusion", "vqbet", "smolvla"],
        "command_template": (
            "python lerobot/scripts/train.py "
            "--dataset.repo_id={dataset} "
            "--policy.type={policy} "
            "--training.num_epochs={epochs}"
        ),
        "format_required": "lerobot-v3",
        "gpu_minimum": "1x GPU, 8GB+ VRAM (24GB+ for SmolVLA)",
        "typical_data": "50-200 demonstrations for single-task ACT",
        "typical_time": "Few hours on 1 GPU for ACT/Diffusion",
        "tips": [
            "50 demos sufficient for single-task ACT, ~100K steps",
            "Diffusion Policy handles multimodal demonstrations better than ACT",
            "VQ-BeT good for discrete action spaces",
            "SmolVLA is a vision-language-action model, needs more data (100+)",
            "Use --training.eval_freq=5000 to monitor progress",
        ],
        "hyperparameters": {
            "act": {"chunk_size": 100, "batch_size": 8, "lr": 1e-5, "steps": 100000},
            "diffusion": {"chunk_size": 16, "batch_size": 64, "lr": 1e-4, "steps": 200000},
        },
    },
    "openvla": {
        "name": "OpenVLA (OFT recommended)",
        "policies": ["openvla-7b", "openvla-7b-oft"],
        "command_template": (
            "torchrun --standalone --nnodes 1 --nproc-per-node 1 "
            "vla-scripts/finetune.py "
            "--vla_path openvla/openvla-7b-oft "
            "--data_root_dir {data_root} "
            "--dataset_name {dataset_name} "
            "--run_root_dir {output_dir} "
            "--batch_size 16 "
            "--learning_rate 2e-5 "
            "--max_steps {steps}"
        ),
        "lora_command_template": (
            "torchrun --standalone --nnodes 1 --nproc-per-node 1 "
            "vla-scripts/finetune.py "
            "--vla_path openvla/openvla-7b-oft "
            "--data_root_dir {data_root} "
            "--dataset_name {dataset_name} "
            "--run_root_dir {output_dir} "
            "--lora_rank 32 "
            "--batch_size 16 "
            "--learning_rate 5e-4 "
            "--max_steps {steps}"
        ),
        "format_required": "RLDS or custom",
        "gpu_minimum": "1x A100 80GB (OFT full) or 1x 24GB+ (LoRA)",
        "typical_data": "100+ demonstrations",
        "typical_time": "10-15 hours on 1x A100",
        "tips": [
            "OFT recipe is recommended since March 2025 — faster and better than original",
            "Use LoRA (--lora_rank 32) if VRAM < 40GB",
            "Needs data_root_dir with specific directory structure",
            "RLDS format works natively, or use custom data loader",
            "Fine-tuning from OFT base significantly outperforms from-scratch VLA",
        ],
    },
    "openpi": {
        "name": "OpenPI / pi0 (Physical Intelligence)",
        "policies": ["pi0", "pi0-fast"],
        "command_template": (
            "XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 "
            "uv run scripts/train.py {config_name} "
            "--exp-name={exp_name}"
        ),
        "format_required": "Custom per-config (usually LeRobot compatible)",
        "gpu_minimum": "1x A100+ (80GB preferred)",
        "typical_data": "50-500 demonstrations depending on task complexity",
        "typical_time": "1-20 hours depending on data size",
        "tips": [
            "REQUIRES a custom TrainConfig Python class defining data source, cameras, actions",
            "Config must be saved to openpi/configs/ and registered in __init__.py",
            "Uses JAX/XLA backend — set XLA_PYTHON_CLIENT_MEM_FRACTION=0.9",
            "pi0 is the base model; pi0-fast uses flow matching for faster inference",
            "Can work with LeRobot format datasets with proper config",
        ],
        "config_template": (
            'from openpi.training.config import TrainConfig\n'
            'from openpi.models.pi0 import Pi0Config\n\n'
            'def {config_fn}() -> TrainConfig:\n'
            '    return TrainConfig(\n'
            '        model=Pi0Config(),\n'
            '        dataset="{dataset}",\n'
            '        num_train_steps={steps},\n'
            '        batch_size={batch_size},\n'
            '    )\n'
        ),
    },
    "groot": {
        "name": "GR00T N1.6 (NVIDIA)",
        "policies": ["groot-n1.6-3b"],
        "command_template": (
            "python gr00t/experiment/launch_finetune.py "
            "--dataset-path {dataset} "
            "--base-model-path nvidia/GR00T-N1.6-3B "
            "--embodiment-tag {embodiment_tag} "
            "--modality-config-path {modality_config} "
            "--num-epochs {epochs} "
            "--batch-size {batch_size}"
        ),
        "format_required": "lerobot-v2",
        "gpu_minimum": "1x A4090/A6000/H100, CUDA 12.4+",
        "typical_data": "50-200 demonstrations",
        "typical_time": "2-8 hours on A6000/H100",
        "tips": [
            "REQUIRES LeRobot v2 format — NOT v3. Convert with: orbit convert <ds> --to lerobot-v2",
            "REQUIRES --embodiment-tag: use NEW_EMBODIMENT for custom robots",
            "REQUIRES modality.json in dataset's meta/ directory",
            "REQUIRES a modality config Python file (ModalityConfig)",
            "Use orbit suggest --framework groot to generate modality.json and config",
            "CUDA 12.4+ required",
        ],
    },
    "osmo": {
        "name": "NVIDIA OSMO (Workflow Orchestration)",
        "policies": [],
        "command_template": "osmo workflow submit {workflow_yaml} --pool {pool_name}",
        "format_required": "Depends on training task",
        "gpu_minimum": "Kubernetes cluster with OSMO deployed",
        "typical_data": "N/A (orchestration layer)",
        "typical_time": "Depends on pipeline tasks",
        "tips": [
            "OSMO orchestrates pipelines — not a training framework itself",
            "YAML-based: tasks with name, image, platform, resources, inputs",
            "Platforms: gb200, rtx-pro-6000, jetson-agx-thor",
            "Can chain orbit gate → training → orbit verify in one workflow",
            "Now has agent context files for coding agents",
        ],
        "workflow_template": (
            "workflow:\n"
            "  tasks:\n"
            "  - name: quality-gate\n"
            "    image: orbit-robotics:latest\n"
            "    command: orbit gate {dataset} -p {policy} --min-grade {min_grade}\n"
            "\n"
            "  - name: train-policy\n"
            "    image: nvcr.io/nvidia/pytorch:latest\n"
            "    platform: {platform}\n"
            "    resources:\n"
            "      gpu: {gpu_count}\n"
            "    inputs:\n"
            "    - task: quality-gate\n"
            "    command: {train_cmd}\n"
            "\n"
            "  - name: verify-result\n"
            "    image: orbit-robotics:latest\n"
            "    command: orbit verify /output/ -d {dataset}\n"
            "    inputs:\n"
            "    - task: train-policy\n"
        ),
    },
}


def get_framework_summary() -> str:
    """Generate a concise text summary of all frameworks for the system prompt."""
    parts: list[str] = []
    for key, fw in FRAMEWORK_KNOWLEDGE.items():
        parts.append(f"### {fw['name']}")
        parts.append(f"- Command: `{fw['command_template']}`")
        parts.append(f"- Format required: {fw['format_required']}")
        parts.append(f"- GPU: {fw['gpu_minimum']}")
        parts.append(f"- Data: {fw['typical_data']}, Time: {fw['typical_time']}")
        for tip in fw["tips"]:
            parts.append(f"- {tip}")
        parts.append("")
    return "\n".join(parts)
