"""Generate ready-to-use project templates (Makefile, scripts, GH workflows)."""

from __future__ import annotations

from pathlib import Path


def generate_makefile(dataset: str, policy: str, min_grade: str = "C") -> str:
    """Generate a Makefile for the training pipeline."""
    return f"""\
# Generated by: orbit init
# Makefile for robot policy training pipeline

DATASET = {dataset}
POLICY = {policy}
MIN_GRADE = {min_grade}

.PHONY: check train verify all clean coach

check:
\torbit gate $(DATASET) -p $(POLICY) --min-grade $(MIN_GRADE)

train: check
\tlerobot-train \\
\t  --dataset.repo_id=$(DATASET) \\
\t  --policy.type=$(POLICY) \\
\t  --batch_size=32 \\
\t  --steps=100000

verify: train
\torbit verify ./outputs/train/latest/ -d $(DATASET)

clean:
\torbit clean $(DATASET) -p $(POLICY) --aggressive

coach:
\torbit coach $(DATASET) --target 50 --policy $(POLICY)

all: check train verify
"""


def generate_train_script(dataset: str, policy: str, min_grade: str = "C") -> str:
    """Generate a training shell script."""
    return f"""\
#!/bin/bash
# Generated by: orbit init
# Training pipeline script for {dataset}

set -euo pipefail

DATASET="{dataset}"
POLICY="{policy}"
MIN_GRADE="{min_grade}"

echo "=== Step 1: Quality Gate ==="
orbit gate "$DATASET" -p "$POLICY" --min-grade "$MIN_GRADE" || {{
    echo "Quality gate failed. Run: orbit analyze $DATASET"
    exit 1
}}

echo "=== Step 2: Training ==="
lerobot-train \\
  --dataset.repo_id="$DATASET" \\
  --policy.type="$POLICY" \\
  --batch_size=32 \\
  --steps=100000

echo "=== Step 3: Verification ==="
orbit verify ./outputs/train/latest/ -d "$DATASET"

echo "=== Done ==="
"""


def generate_github_workflow(dataset: str, policy: str, min_grade: str = "C") -> str:
    """Generate a GitHub Actions workflow for training."""
    return f"""\
# Generated by: orbit init
name: Train Robot Policy
on:
  push:
    paths: ['data/**']
  workflow_dispatch:

jobs:
  quality-check:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: '3.11'
      - run: pip install orbit-robotics[formats]
      - run: orbit gate {dataset} -p {policy} --min-grade {min_grade}

  train:
    needs: quality-check
    runs-on: [self-hosted, gpu]
    steps:
      - uses: actions/checkout@v4
      - run: pip install lerobot orbit-robotics
      - run: |
          lerobot-train \\
            --dataset.repo_id={dataset} \\
            --policy.type={policy} \\
            --batch_size=32 \\
            --steps=100000

  verify:
    needs: train
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - run: pip install orbit-robotics
      - run: orbit verify ./outputs/train/latest/ -d {dataset}
"""


def init_project(
    dataset: str,
    policy: str,
    output_dir: str = ".",
    min_grade: str = "C",
) -> list[str]:
    """Generate all project files.

    Returns:
        List of created file paths.
    """
    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)
    created = []

    # Makefile
    makefile = out / "Makefile"
    makefile.write_text(generate_makefile(dataset, policy, min_grade))
    created.append(str(makefile))

    # train.sh
    train_sh = out / "train.sh"
    train_sh.write_text(generate_train_script(dataset, policy, min_grade))
    train_sh.chmod(0o755)
    created.append(str(train_sh))

    # GitHub workflow
    gh_dir = out / ".github" / "workflows"
    gh_dir.mkdir(parents=True, exist_ok=True)
    workflow = gh_dir / "train.yml"
    workflow.write_text(generate_github_workflow(dataset, policy, min_grade))
    created.append(str(workflow))

    return created
