"""Lambda Labs cloud GPU integration.

Lambda Labs provides affordable GPU instances for training:
- A100 40GB: ~$1.10/hour
- H100: ~$2.00/hour
- A10: ~$0.60/hour

API docs: https://cloud.lambdalabs.com/api/v1/docs

Usage:
    # Set API key
    export LAMBDA_API_KEY=your_key_here

    # List available instances
    python -m openadapt_ml.cloud.lambda_labs list

    # Launch instance for training
    python -m openadapt_ml.cloud.lambda_labs launch --type gpu_1x_a100

    # Check running instances
    python -m openadapt_ml.cloud.lambda_labs status

    # Terminate instance
    python -m openadapt_ml.cloud.lambda_labs terminate <instance_id>
"""

from __future__ import annotations

import json
import os
import shlex
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import requests


API_BASE = "https://cloud.lambdalabs.com/api/v1"

# Default port for HTTP server
DEFAULT_SERVER_PORT = 8765


def start_dashboard_server(
    output_dir: Path, port: int = DEFAULT_SERVER_PORT
) -> tuple[subprocess.Popen, str]:
    """Start a background HTTP server for the dashboard.

    Args:
        output_dir: Directory containing dashboard files
        port: Port to serve on

    Returns:
        (process, url): The server process and the dashboard URL
    """

    # Start simple HTTP server in background thread
    server_proc = subprocess.Popen(
        [sys.executable, "-m", "http.server", str(port)],
        cwd=str(output_dir),
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )

    url = f"http://localhost:{port}/dashboard.html"

    # Give server time to start
    time.sleep(0.5)

    return server_proc, url


def open_dashboard_in_browser(output_dir: Path, port: int = DEFAULT_SERVER_PORT):
    """Start HTTP server and open dashboard in browser.

    Args:
        output_dir: Directory containing dashboard files
        port: Port to serve on

    Returns:
        Server process (caller should call .terminate() when done), or None if failed
    """
    import webbrowser

    try:
        server_proc, url = start_dashboard_server(output_dir, port)
        webbrowser.open(url)
        print(f"Dashboard: {url}")
        print("  Stop Training button enabled in dashboard")
        return server_proc
    except Exception as e:
        print(f"Warning: Could not start dashboard server: {e}")
        return None


def setup_capture_screenshots_symlink(
    output_dir: Path, capture_path: str | Path
) -> bool:
    """Create symlink from output_dir/screenshots to capture's screenshots folder.

    This allows the dashboard to serve screenshots via relative paths.

    Args:
        output_dir: Training output directory (e.g., training_output/job_id/)
        capture_path: Path to capture directory (local)

    Returns:
        True if symlink created successfully
    """
    capture_path = Path(capture_path)
    screenshots_src = capture_path / "screenshots"
    screenshots_dst = output_dir / "screenshots"

    if not screenshots_src.exists():
        return False

    # Remove existing symlink or directory
    if screenshots_dst.is_symlink():
        screenshots_dst.unlink()
    elif screenshots_dst.exists():
        return False  # Don't overwrite real directory

    try:
        screenshots_dst.symlink_to(screenshots_src.resolve())
        return True
    except Exception:
        return False


def rewrite_evaluation_paths(
    evaluations: list[dict], remote_prefix: str = "/home/ubuntu/capture/"
) -> list[dict]:
    """Rewrite Lambda paths in evaluations to relative paths.

    Converts: /home/ubuntu/capture/screenshots/foo.png -> screenshots/foo.png

    Args:
        evaluations: List of evaluation dicts with image_path
        remote_prefix: The Lambda path prefix to replace

    Returns:
        Evaluations with rewritten paths
    """
    for ev in evaluations:
        if "image_path" in ev and ev["image_path"].startswith(remote_prefix):
            ev["image_path"] = ev["image_path"].replace(remote_prefix, "")
    return evaluations


def download_checkpoints_from_instance(
    instance_ip: str, output_dir: Path, ssh_key: str | None = None
) -> bool:
    """Download checkpoints from Lambda instance.

    Args:
        instance_ip: IP address of Lambda instance
        output_dir: Local directory to save checkpoints
        ssh_key: Path to SSH key (uses default if not provided)

    Returns:
        True if download succeeded
    """
    checkpoints_dir = output_dir / "checkpoints"
    checkpoints_dir.mkdir(parents=True, exist_ok=True)

    ssh_key = ssh_key or str(Path.home() / ".ssh" / "lambda_id_ed25519")
    ssh_opts = (
        f"-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -i {ssh_key}"
    )

    # Download checkpoints from remote
    remote_path = f"ubuntu@{instance_ip}:~/openadapt-ml/checkpoints/"
    local_path = str(checkpoints_dir) + "/"

    cmd = f"rsync -avz --progress -e 'ssh {ssh_opts}' {remote_path} {local_path}"
    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)

    if result.returncode == 0:
        return True
    return False


def check_stop_signal(output_dir: Path) -> bool:
    """Check if stop signal file exists.

    The dashboard can create this file to signal training should stop.
    """
    stop_file = output_dir / "STOP_TRAINING"
    return stop_file.exists()


@dataclass
class InstanceType:
    """Lambda Labs instance type."""

    name: str
    price_cents_per_hour: int
    description: str
    gpu_count: int
    gpu_type: str
    vcpus: int
    memory_gb: int
    storage_gb: int
    available_regions: list[str]

    @property
    def price_per_hour(self) -> float:
        return self.price_cents_per_hour / 100

    def __str__(self) -> str:
        regions = ", ".join(self.available_regions[:3])
        if len(self.available_regions) > 3:
            regions += f" (+{len(self.available_regions) - 3} more)"
        return (
            f"{self.name}: ${self.price_per_hour:.2f}/hr | "
            f"{self.gpu_count}x {self.gpu_type} | {self.vcpus} vCPUs | "
            f"{self.memory_gb}GB RAM | {self.storage_gb}GB SSD | "
            f"Regions: {regions}"
        )


@dataclass
class Instance:
    """Running Lambda Labs instance."""

    id: str
    name: str
    instance_type: str
    status: str
    ip: str | None
    region: str
    ssh_key_names: list[str]

    def __str__(self) -> str:
        ip_str = self.ip or "pending"
        return f"{self.id[:8]}... | {self.instance_type} | {self.status} | IP: {ip_str} | {self.region}"


class LambdaLabsClient:
    """Client for Lambda Labs API."""

    def __init__(self, api_key: str | None = None):
        # Try provided key, then settings, then env var
        if not api_key:
            from openadapt_ml.config import settings

            api_key = settings.lambda_api_key or os.environ.get("LAMBDA_API_KEY")

        self.api_key = api_key
        if not self.api_key:
            raise ValueError(
                "Lambda Labs API key required. Set LAMBDA_API_KEY in .env file "
                "or get one at https://cloud.lambdalabs.com/api-keys"
            )
        self.session = requests.Session()
        self.session.headers["Authorization"] = f"Bearer {self.api_key}"

    def _get(self, endpoint: str) -> dict[str, Any]:
        """Make GET request to API."""
        resp = self.session.get(f"{API_BASE}{endpoint}")
        resp.raise_for_status()
        return resp.json()

    def _post(self, endpoint: str, data: dict[str, Any]) -> dict[str, Any]:
        """Make POST request to API."""
        resp = self.session.post(f"{API_BASE}{endpoint}", json=data)
        if not resp.ok:
            error = resp.json().get("error", {})
            raise RuntimeError(f"API error: {error.get('message', resp.text)}")
        return resp.json()

    def list_instance_types(self) -> list[InstanceType]:
        """List available GPU instance types."""
        data = self._get("/instance-types")
        types = []

        for name, info in data.get("data", {}).items():
            specs = info.get("instance_type", {}).get("specs", {})
            regions = [
                r["name"] for r in info.get("regions_with_capacity_available", [])
            ]

            types.append(
                InstanceType(
                    name=name,
                    price_cents_per_hour=info.get("instance_type", {}).get(
                        "price_cents_per_hour", 0
                    ),
                    description=info.get("instance_type", {}).get("description", ""),
                    gpu_count=specs.get("gpus", 0),
                    gpu_type=info.get("instance_type", {}).get("gpu_description", ""),
                    vcpus=specs.get("vcpus", 0),
                    memory_gb=specs.get("memory_gib", 0),
                    storage_gb=specs.get("storage_gib", 0),
                    available_regions=regions,
                )
            )

        # Sort by price
        types.sort(key=lambda t: t.price_cents_per_hour)
        return types

    def list_ssh_keys(self) -> list[dict[str, str]]:
        """List registered SSH keys."""
        data = self._get("/ssh-keys")
        return data.get("data", [])

    def add_ssh_key(self, name: str, public_key: str) -> dict[str, str]:
        """Add an SSH key."""
        data = self._post("/ssh-keys", {"name": name, "public_key": public_key})
        return data.get("data", {})

    def list_instances(self) -> list[Instance]:
        """List running instances."""
        data = self._get("/instances")
        instances = []

        for inst in data.get("data", []):
            # ssh_key_names can be list of strings or list of dicts
            ssh_keys = inst.get("ssh_key_names", [])
            if ssh_keys and isinstance(ssh_keys[0], dict):
                ssh_key_names = [k["name"] for k in ssh_keys]
            else:
                ssh_key_names = ssh_keys  # Already list of strings

            instances.append(
                Instance(
                    id=inst["id"],
                    name=inst.get("name", ""),
                    instance_type=inst.get("instance_type", {}).get("name", "unknown"),
                    status=inst.get("status", "unknown"),
                    ip=inst.get("ip"),
                    region=inst.get("region", {}).get("name", "unknown"),
                    ssh_key_names=ssh_key_names,
                )
            )

        return instances

    def launch_instance(
        self,
        instance_type: str,
        region: str | None = None,
        ssh_key_names: list[str] | None = None,
        name: str | None = None,
    ) -> Instance:
        """Launch a new GPU instance.

        Args:
            instance_type: Instance type name (e.g., 'gpu_1x_a100')
            region: Region name (auto-selects if None)
            ssh_key_names: SSH key names to use
            name: Optional instance name

        Returns:
            Launched instance
        """
        # If no region specified, find one with capacity
        if not region:
            types = self.list_instance_types()
            for t in types:
                if t.name == instance_type and t.available_regions:
                    region = t.available_regions[0]
                    break
            if not region:
                raise RuntimeError(f"No regions available for {instance_type}")

        # If no SSH key specified, use first available
        if not ssh_key_names:
            keys = self.list_ssh_keys()
            if not keys:
                raise RuntimeError(
                    "No SSH keys found. Add one at https://cloud.lambdalabs.com/ssh-keys"
                )
            ssh_key_names = [keys[0]["name"]]

        payload = {
            "region_name": region,
            "instance_type_name": instance_type,
            "ssh_key_names": ssh_key_names,
        }
        if name:
            payload["name"] = name

        data = self._post("/instance-operations/launch", payload)
        instance_ids = data.get("data", {}).get("instance_ids", [])

        if not instance_ids:
            raise RuntimeError("Failed to launch instance")

        # Wait for instance to be ready
        print(f"Instance {instance_ids[0]} launched, waiting for IP...")
        instance = None
        for _ in range(60):  # Wait up to 5 minutes for IP
            instances = self.list_instances()
            for inst in instances:
                if inst.id == instance_ids[0] and inst.ip:
                    instance = inst
                    break
            if instance:
                break
            time.sleep(5)

        if not instance:
            raise RuntimeError("Timed out waiting for instance IP")

        # Wait for SSH to be ready - be patient, instances can take a while to boot
        print(f"Instance IP: {instance.ip}, waiting for SSH...")
        for attempt in range(60):  # Wait up to 5 minutes for SSH
            try:
                result = subprocess.run(
                    [
                        "ssh",
                        "-o",
                        "StrictHostKeyChecking=no",
                        "-o",
                        "ConnectTimeout=10",
                        f"ubuntu@{instance.ip}",
                        "echo ready",
                    ],
                    capture_output=True,
                    text=True,
                    timeout=20,
                )
                if result.returncode == 0:
                    print("SSH ready!")
                    return instance
            except subprocess.TimeoutExpired:
                pass
            if attempt % 6 == 5:  # Log progress every 30 seconds
                print(f"  Still waiting for SSH ({(attempt + 1) * 5}s elapsed)...")
            time.sleep(5)

        print("Warning: SSH may not be ready yet, continuing anyway...")
        return instance

    def terminate_instance(self, instance_id: str) -> bool:
        """Terminate an instance."""
        data = self._post(
            "/instance-operations/terminate", {"instance_ids": [instance_id]}
        )
        terminated = data.get("data", {}).get("terminated_instances", [])
        return any(t.get("id") == instance_id for t in terminated)

    def get_ssh_command(self, instance: Instance, user: str = "ubuntu") -> str:
        """Get SSH command for an instance."""
        if not instance.ip:
            return "# Instance IP not yet available"
        return f"ssh {user}@{instance.ip}"

    def ssh_run(
        self,
        instance: Instance,
        command: str,
        timeout: int | None = None,
        retries: int = 3,
    ) -> subprocess.CompletedProcess:
        """Run a command on an instance via SSH.

        Args:
            instance: Instance to run on
            command: Shell command to run
            timeout: Optional timeout in seconds
            retries: Number of retries on connection failure

        Returns:
            CompletedProcess with stdout/stderr
        """
        if not instance.ip:
            raise RuntimeError("Instance has no IP address")

        ssh_cmd = [
            "ssh",
            "-o",
            "StrictHostKeyChecking=no",
            "-o",
            "ConnectTimeout=30",  # Increased from 10
            "-o",
            "ServerAliveInterval=60",  # Keep connection alive
            "-o",
            "ServerAliveCountMax=3",
            f"ubuntu@{instance.ip}",
            command,
        ]

        last_error = None
        for attempt in range(retries):
            try:
                return subprocess.run(
                    ssh_cmd,
                    capture_output=True,
                    text=True,
                    timeout=timeout,
                )
            except subprocess.TimeoutExpired as e:
                last_error = e
                if attempt < retries - 1:
                    print(f"  SSH timeout, retrying ({attempt + 1}/{retries})...")
                    time.sleep(5)

        raise last_error if last_error else RuntimeError("SSH failed")

    def setup_instance(
        self,
        instance: Instance,
        repo_url: str = "https://github.com/OpenAdaptAI/openadapt-ml.git",
        clean_gpu: bool = True,
    ) -> bool:
        """Set up training environment on instance.

        Clones repo, installs uv, syncs dependencies.
        Optionally clears GPU memory from previous runs.
        Returns True if successful.
        """
        print(f"Setting up instance {instance.ip}...")

        # Clean GPU memory if requested (don't fail if this doesn't work)
        if clean_gpu:
            print("  Clearing GPU memory...")
            try:
                self.ssh_run(
                    instance,
                    """
python3 -c "
import torch
if torch.cuda.is_available():
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    print('GPU memory cleared')
" 2>/dev/null || true
# Kill any stale python processes using GPU
pkill -f "python.*train" 2>/dev/null || true
""",
                    timeout=60,
                )
            except Exception as e:
                print(f"  GPU cleanup skipped: {e}")

        setup_script = f"""
set -e
cd ~

# Install uv via official installer (most robust)
if ! command -v uv &> /dev/null; then
    curl -LsSf https://astral.sh/uv/install.sh | sh
fi
export PATH="$HOME/.local/bin:$HOME/.cargo/bin:$PATH"

# Clone or update repo
if [ ! -d "openadapt-ml" ]; then
    git clone {repo_url}
else
    cd openadapt-ml && git pull origin main && cd ~
fi

cd openadapt-ml

# Remove uv.sources (local path deps that don't exist on remote)
sed -i '/\\[tool.uv.sources\\]/,/^$/d' pyproject.toml 2>/dev/null || true

uv sync
echo "SETUP_COMPLETE"
"""

        try:
            result = self.ssh_run(
                instance, setup_script, timeout=900
            )  # 15 min timeout for setup

            if "SETUP_COMPLETE" in result.stdout:
                print("  Environment ready")
                return True
            else:
                stderr_preview = result.stderr[:500] if result.stderr else "(no stderr)"
                print(f"  Setup failed: {stderr_preview}")
                return False
        except subprocess.TimeoutExpired:
            print("  Setup timed out after 15 minutes")
            return False
        except Exception as e:
            print(f"  Setup failed: {e}")
            return False

    def sync_local_code(
        self, instance: Instance, local_repo_path: str = ".", retries: int = 3
    ) -> bool:
        """Sync local code to remote instance via git archive.

        Uses `git archive HEAD` to send only committed, tracked files (~10MB)
        instead of rsync'ing the full working directory (~1.8GB with binary
        artifacts, checkpoints, benchmark results, etc.).

        Args:
            instance: Instance to sync to
            local_repo_path: Local repository path (must be a git repo)
            retries: Number of retry attempts

        Returns:
            True if successful
        """
        if not instance.ip:
            raise RuntimeError("Instance has no IP address")

        print(f"Syncing local code to {instance.ip} via git archive...")

        ssh_opts = (
            "ssh -o StrictHostKeyChecking=no "
            "-o ConnectTimeout=30 -o ServerAliveInterval=60"
        )

        # git archive HEAD produces a tarball of only committed tracked files,
        # piped over SSH and extracted on the remote — clean and minimal.
        cmd = (
            f"cd {shlex.quote(local_repo_path)} && "
            f"git archive HEAD | "
            f"{ssh_opts} ubuntu@{instance.ip} "
            f"'mkdir -p ~/openadapt-ml && tar -xf - -C ~/openadapt-ml/'"
        )

        for attempt in range(retries):
            result = subprocess.run(cmd, shell=True)
            if result.returncode == 0:
                print("  Code synced (git archive)")
                return True
            if attempt < retries - 1:
                print(f"  Sync failed, retrying ({attempt + 1}/{retries})...")
                time.sleep(5)

        return False

    def upload_capture(
        self,
        instance: Instance,
        local_path: str,
        remote_path: str = "~/capture",
        retries: int = 3,
    ) -> bool:
        """Upload a capture directory to instance via rsync.

        Args:
            instance: Instance to upload to
            local_path: Local path to capture directory
            remote_path: Remote path (default: ~/capture)
            retries: Number of retry attempts

        Returns:
            True if successful
        """
        if not instance.ip:
            raise RuntimeError("Instance has no IP address")

        print(f"Uploading capture to {instance.ip}:{remote_path}...")

        # SSH options for more robust connection
        ssh_opts = "ssh -o StrictHostKeyChecking=no -o ConnectTimeout=30 -o ServerAliveInterval=60"

        rsync_cmd = [
            "rsync",
            "-avz",
            "--progress",
            "--timeout=120",  # 2 minute timeout per file
            "-e",
            ssh_opts,
            f"{local_path}/",
            f"ubuntu@{instance.ip}:{remote_path}/",
        ]

        for attempt in range(retries):
            result = subprocess.run(rsync_cmd)
            if result.returncode == 0:
                return True
            if attempt < retries - 1:
                print(f"  Upload failed, retrying ({attempt + 1}/{retries})...")
                time.sleep(5)

        return False

    def run_training(
        self,
        instance: Instance,
        config: str = "configs/qwen3vl_capture.yaml",
        capture: str | None = None,
        goal: str | None = None,
        jsonl: str | None = None,
        background: bool = True,
    ) -> subprocess.Popen | subprocess.CompletedProcess:
        """Run training on instance.

        Args:
            instance: Instance to train on
            config: Config file path (relative to repo)
            capture: Remote capture path (if uploaded)
            goal: Task goal description
            jsonl: Remote path to JSONL training data
            background: Run in background (returns Popen) or foreground

        Returns:
            Popen if background=True, CompletedProcess if background=False
        """
        if not instance.ip:
            raise RuntimeError("Instance has no IP address")

        # Build training command (quote user-supplied values for shell safety)
        train_cmd = f"uv run python -m openadapt_ml.scripts.train --config {shlex.quote(config)}"
        if jsonl:
            train_cmd += f" --jsonl {shlex.quote(jsonl)}"
        elif capture:
            train_cmd += f" --capture {shlex.quote(capture)}"
        if goal:
            train_cmd += f" --goal {shlex.quote(goal)}"

        # Full script with environment setup
        script = f"""
cd ~/openadapt-ml
export PATH="$HOME/.local/bin:$PATH"
{train_cmd}
"""

        ssh_cmd = [
            "ssh",
            "-o",
            "StrictHostKeyChecking=no",
            f"ubuntu@{instance.ip}",
            script,
        ]

        print(f"Running training on {instance.ip}...")
        print(f"  Config: {config}")
        if capture:
            print(f"  Capture: {capture}")

        if background:
            # Run in background, return Popen for monitoring
            return subprocess.Popen(
                ssh_cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
            )
        else:
            # Run in foreground, stream output
            return subprocess.run(ssh_cmd)

    def download_results(
        self,
        instance: Instance,
        remote_path: str = "~/openadapt-ml",
        local_path: str = ".",
        include_checkpoint: bool = True,
        include_logs: bool = True,
    ) -> bool:
        """Download training results from instance.

        Args:
            instance: Instance to download from
            remote_path: Remote openadapt-ml directory
            local_path: Local directory to download to
            include_checkpoint: Download checkpoint weights
            include_logs: Download training logs and dashboard

        Returns:
            True if successful
        """
        if not instance.ip:
            raise RuntimeError("Instance has no IP address")

        print(f"Downloading results from {instance.ip}...")
        success = True

        # Download training output (logs, dashboard)
        if include_logs:
            print("  Downloading training logs...")
            rsync_cmd = [
                "rsync",
                "-avz",
                "-e",
                "ssh -o StrictHostKeyChecking=no",
                f"ubuntu@{instance.ip}:{remote_path}/training_output/",
                f"{local_path}/training_output_lambda/",
            ]
            result = subprocess.run(rsync_cmd, capture_output=True)
            if result.returncode == 0:
                print("    Training logs downloaded to training_output_lambda/")
            else:
                print("    Warning: Failed to download logs")
                success = False

        # Download checkpoint
        if include_checkpoint:
            print("  Downloading checkpoint...")
            rsync_cmd = [
                "rsync",
                "-avz",
                "-e",
                "ssh -o StrictHostKeyChecking=no",
                f"ubuntu@{instance.ip}:{remote_path}/checkpoints/",
                f"{local_path}/checkpoints_lambda/",
            ]
            result = subprocess.run(rsync_cmd, capture_output=True)
            if result.returncode == 0:
                print("    Checkpoint downloaded to checkpoints_lambda/")
            else:
                print("    Warning: Failed to download checkpoint (may not exist yet)")

        # Regenerate all dashboards with static navigation and correct status
        if include_logs:
            try:
                from openadapt_ml.training.trainer import regenerate_all_dashboards

                output_dir = Path(local_path) / "training_output_lambda"
                if output_dir.exists():
                    print("  Regenerating dashboards with static navigation...")
                    regenerate_all_dashboards(output_dir)
            except Exception as e:
                print(f"    Warning: Failed to regenerate dashboards: {e}")

        return success

    def get_training_status(self, instance: Instance) -> dict:
        """Check training status by reading training_log.json on instance."""
        result = self.ssh_run(
            instance,
            "cat ~/openadapt-ml/training_output/training_log.json 2>/dev/null || echo '{}'",
            timeout=10,
        )
        try:
            import json

            return json.loads(result.stdout.strip())
        except Exception:
            return {}


def setup_lambda_ssh_key(client: LambdaLabsClient) -> str:
    """Set up SSH key for Lambda Labs if not already done.

    Returns the SSH key name that was added/found.
    """
    # Check if we already have keys
    keys = client.list_ssh_keys()
    if keys:
        print(f"Using existing SSH key: {keys[0]['name']}")
        return keys[0]["name"]

    # Look for local SSH key
    ssh_key_path = Path.home() / ".ssh" / "id_rsa.pub"
    if not ssh_key_path.exists():
        ssh_key_path = Path.home() / ".ssh" / "id_ed25519.pub"

    if not ssh_key_path.exists():
        raise RuntimeError(
            "No SSH key found at ~/.ssh/id_rsa.pub or ~/.ssh/id_ed25519.pub\n"
            "Generate one with: ssh-keygen -t ed25519"
        )

    public_key = ssh_key_path.read_text().strip()
    key_name = f"openadapt-{os.environ.get('USER', 'user')}"

    print(f"Adding SSH key '{key_name}' to Lambda Labs...")
    client.add_ssh_key(key_name, public_key)
    return key_name


def main():
    """CLI for Lambda Labs."""
    import argparse

    parser = argparse.ArgumentParser(description="Lambda Labs GPU management")
    subparsers = parser.add_subparsers(dest="command", help="Command")

    # List instances command
    subparsers.add_parser("list", help="List available instance types")

    # Status command
    subparsers.add_parser("status", help="Show running instances")

    # Launch command
    launch_parser = subparsers.add_parser("launch", help="Launch a GPU instance")
    launch_parser.add_argument(
        "--type",
        "-t",
        default="gpu_1x_a100",
        help="Instance type (default: gpu_1x_a100)",
    )
    launch_parser.add_argument(
        "--region", "-r", help="Region (auto-selects if not specified)"
    )
    launch_parser.add_argument("--name", "-n", help="Instance name")

    # Terminate command
    term_parser = subparsers.add_parser("terminate", help="Terminate an instance")
    term_parser.add_argument("instance_id", help="Instance ID to terminate")

    # SSH command - run commands or get interactive shell
    ssh_parser = subparsers.add_parser(
        "ssh", help="SSH into Lambda instance or run command"
    )
    ssh_parser.add_argument(
        "instance_id", nargs="?", help="Instance ID (uses first if not specified)"
    )
    ssh_parser.add_argument(
        "--cmd", "-c", help="Command to run (opens shell if not specified)"
    )
    ssh_parser.add_argument(
        "--timeout", "-t", type=int, default=60, help="Command timeout in seconds"
    )

    # Serve command - start dashboard server with stop button support
    serve_parser = subparsers.add_parser(
        "serve", help="Start dashboard server with stop button support"
    )
    serve_parser.add_argument(
        "--output",
        "-o",
        default="training_output",
        help="Output directory (default: training_output)",
    )
    serve_parser.add_argument(
        "--port", "-p", type=int, default=8765, help="Port (default: 8765)"
    )
    serve_parser.add_argument(
        "--open", action="store_true", help="Open dashboard in browser"
    )

    # Rsync command - copy files to/from Lambda instance
    rsync_parser = subparsers.add_parser(
        "rsync", help="Rsync files to/from Lambda instance"
    )
    rsync_parser.add_argument(
        "source", help="Source path (prefix with 'remote:' for remote paths)"
    )
    rsync_parser.add_argument(
        "dest", help="Destination path (prefix with 'remote:' for remote paths)"
    )
    rsync_parser.add_argument(
        "instance_id", nargs="?", help="Instance ID (uses first if not specified)"
    )
    rsync_parser.add_argument(
        "--delete", action="store_true", help="Delete extraneous files from dest"
    )

    # Setup command
    subparsers.add_parser("setup", help="Set up SSH key for Lambda Labs")

    # Train command - full automated training pipeline
    train_parser = subparsers.add_parser("train", help="Run training on Lambda GPU")
    train_parser.add_argument("--capture", "-c", help="Local path to capture directory")
    train_parser.add_argument("--goal", "-g", help="Task goal description")
    train_parser.add_argument(
        "--config",
        default="configs/qwen3vl_capture_4bit.yaml",
        help="Config file (default: 4bit for memory efficiency)",
    )
    train_parser.add_argument(
        "--type", "-t", default="gpu_1x_a10", help="Instance type"
    )
    train_parser.add_argument(
        "--instance", "-i", help="Use existing instance ID instead of launching new"
    )
    train_parser.add_argument(
        "--bundle",
        "-b",
        help="Local bundle directory (with training_data.jsonl + images/) to upload",
    )
    train_parser.add_argument(
        "--demo-dir",
        help="Directory with annotated demo JSON files (auto-converts to bundle)",
    )
    train_parser.add_argument(
        "--captures-dir",
        help="Parent directory containing capture directories (for screenshot resolution)",
    )
    train_parser.add_argument(
        "--mapping",
        help="Pre-computed screenshot_mapping.json (optional, auto-generated if omitted)",
    )
    train_parser.add_argument(
        "--gpu-wait",
        type=int,
        default=0,
        metavar="MINUTES",
        help="Wait up to N minutes for GPU availability (0=fail immediately, default: 0)",
    )
    train_parser.add_argument(
        "--no-terminate",
        action="store_true",
        help="Don't terminate instance after training",
    )
    train_parser.add_argument(
        "--max-runtime",
        type=int,
        default=60,
        help="Max runtime in minutes before auto-terminate (default: 60)",
    )
    train_parser.add_argument(
        "--open",
        action="store_true",
        help="Open dashboard in browser when training starts",
    )

    # Training status command
    train_status_parser = subparsers.add_parser(
        "train-status", help="Check training status on instance"
    )
    train_status_parser.add_argument("instance_id", nargs="?", help="Instance ID")

    # Monitor command - live dashboard for Lambda training
    monitor_parser = subparsers.add_parser(
        "monitor", help="Monitor Lambda training with live dashboard"
    )
    monitor_parser.add_argument("instance_id", nargs="?", help="Instance ID")
    monitor_parser.add_argument(
        "--open", action="store_true", help="Open dashboard in browser"
    )
    monitor_parser.add_argument(
        "--interval", type=int, default=5, help="Poll interval in seconds (default: 5)"
    )
    monitor_parser.add_argument(
        "--capture", type=str, help="Local capture path for screenshot symlink"
    )
    monitor_parser.add_argument(
        "--auto-stop-loss",
        type=float,
        default=0.5,
        help="Auto-terminate when loss drops below this (default: 0.5)",
    )
    monitor_parser.add_argument(
        "--download-checkpoints",
        action="store_true",
        default=True,
        help="Auto-download checkpoints each epoch",
    )
    monitor_parser.add_argument(
        "--no-download-checkpoints",
        action="store_false",
        dest="download_checkpoints",
        help="Disable checkpoint download",
    )
    monitor_parser.add_argument(
        "--stub",
        action="store_true",
        help="Use stub training provider (no GPU, instant simulation)",
    )

    # Refresh command - one-shot dashboard update
    refresh_parser = subparsers.add_parser(
        "refresh", help="One-shot refresh of training dashboard"
    )
    refresh_parser.add_argument("instance_id", nargs="?", help="Instance ID")
    refresh_parser.add_argument(
        "--open", action="store_true", help="Open dashboard in browser"
    )
    refresh_parser.add_argument(
        "--capture", type=str, help="Local capture path for screenshot preview"
    )

    # Checkpoints command - list remote checkpoints
    checkpoints_parser = subparsers.add_parser(
        "checkpoints", help="List checkpoints on remote instance"
    )
    checkpoints_parser.add_argument("instance_id", nargs="?", help="Instance ID")

    # Download results command
    download_parser = subparsers.add_parser(
        "download", help="Download training results from instance"
    )
    download_parser.add_argument("instance_id", nargs="?", help="Instance ID")
    download_parser.add_argument(
        "--output", "-o", default=".", help="Local output directory"
    )

    # Check files on instance
    files_parser = subparsers.add_parser(
        "files", help="List training files on instance"
    )
    files_parser.add_argument("instance_id", nargs="?", help="Instance ID")
    files_parser.add_argument(
        "--path", "-p", default="~/openadapt-ml", help="Path to check"
    )

    # Kill command - terminate training processes
    kill_parser = subparsers.add_parser(
        "kill", help="Kill training/inference processes on instance"
    )
    kill_parser.add_argument("instance_id", nargs="?", help="Instance ID")
    kill_parser.add_argument(
        "--local", action="store_true", help="Also kill local Lambda-related processes"
    )
    kill_parser.add_argument(
        "--all",
        action="store_true",
        help="Kill all Python processes on instance (careful!)",
    )

    # Check command - analyze training status and early stopping
    check_parser = subparsers.add_parser(
        "check", help="Check training health and early stopping status"
    )
    check_parser.add_argument("instance_id", nargs="?", help="Instance ID")
    check_parser.add_argument(
        "--threshold",
        "-t",
        type=float,
        default=0.01,
        help="Early stopping threshold (loss improvement over last N steps)",
    )
    check_parser.add_argument(
        "--window",
        "-w",
        type=int,
        default=10,
        help="Number of recent steps to check for improvement",
    )

    # Compare command - run comparison on Lambda and sync back
    compare_parser = subparsers.add_parser(
        "compare", help="Run human vs AI comparison on Lambda"
    )
    compare_parser.add_argument("instance_id", nargs="?", help="Instance ID")
    compare_parser.add_argument(
        "--checkpoint", "-c", help="Checkpoint to use (default: latest)"
    )
    compare_parser.add_argument(
        "--epoch", "-e", type=int, help="Use checkpoint from specific epoch"
    )
    compare_parser.add_argument(
        "--open", action="store_true", help="Open viewer after generation"
    )

    # Results viewer command - downloads and generates comparison viewer
    results_parser = subparsers.add_parser(
        "results", help="Download results and generate comparison viewer"
    )
    results_parser.add_argument(
        "--capture",
        "-c",
        required=True,
        help="Local capture directory (for comparison)",
    )
    results_parser.add_argument("--goal", "-g", help="Task goal description")
    results_parser.add_argument(
        "--open", action="store_true", help="Open viewer in browser"
    )
    results_parser.add_argument("instance_id", nargs="?", help="Instance ID")

    # Sync command - sync training output and regenerate navigation for file:// protocol
    sync_parser = subparsers.add_parser(
        "sync", help="Sync training output from Lambda and regenerate navigation"
    )
    sync_parser.add_argument("instance_id", nargs="?", help="Instance ID")
    sync_parser.add_argument(
        "--output",
        "-o",
        default="training_output",
        help="Local output directory (default: training_output)",
    )
    sync_parser.add_argument(
        "--open", action="store_true", help="Open dashboard in browser after sync"
    )

    # Viewer command - regenerate local viewer (no Lambda required)
    viewer_parser = subparsers.add_parser(
        "viewer", help="Regenerate local viewer (no Lambda required)"
    )
    viewer_parser.add_argument(
        "--output",
        "-o",
        default="training_output",
        help="Training output directory (default: training_output)",
    )
    viewer_parser.add_argument(
        "--dashboard",
        "-d",
        action="store_true",
        help="Regenerate dashboard instead of viewer",
    )
    viewer_parser.add_argument(
        "--open",
        action="store_true",
        help="Open in browser (use 'serve' instead for better experience)",
    )

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        return

    try:
        client = LambdaLabsClient()
    except ValueError as e:
        print(f"Error: {e}")
        print("\nGet your API key at https://cloud.lambdalabs.com/api-keys")
        print("Then set it: export LAMBDA_API_KEY=your_key_here")
        return

    if args.command == "list":
        print("Available GPU instances:\n")
        types = client.list_instance_types()
        for t in types:
            print(f"  {t}")
        print(f"\nTotal: {len(types)} instance types")
        print(
            "\nLaunch with: python -m openadapt_ml.cloud.lambda_labs launch --type <name>"
        )

    elif args.command == "status":
        instances = client.list_instances()
        if not instances:
            print("No running instances.")
        else:
            print("Running instances:\n")
            for inst in instances:
                print(f"  {inst}")
            print(f"\nTotal: {len(instances)} instances")

    elif args.command == "launch":
        # Ensure SSH key is set up
        ssh_key = setup_lambda_ssh_key(client)

        print(f"Launching {args.type}...")
        instance = client.launch_instance(
            instance_type=args.type,
            region=args.region,
            ssh_key_names=[ssh_key],
            name=args.name,
        )
        print("\nInstance launched!")
        print(f"  ID: {instance.id}")
        print(f"  IP: {instance.ip}")
        print(f"  Type: {instance.instance_type}")
        print(f"  Region: {instance.region}")
        print(f"\nConnect with: ssh ubuntu@{instance.ip}")
        print(
            f"\nTerminate with: python -m openadapt_ml.cloud.lambda_labs terminate {instance.id}"
        )

    elif args.command == "terminate":
        if client.terminate_instance(args.instance_id):
            print(f"Instance {args.instance_id} terminated.")
        else:
            print(f"Failed to terminate instance {args.instance_id}")

    elif args.command == "ssh":
        instances = client.list_instances()
        if not instances:
            print("No running instances.")
            return

        if args.instance_id:
            instance = next(
                (i for i in instances if i.id.startswith(args.instance_id)), None
            )
            if not instance:
                print(f"Instance {args.instance_id} not found.")
                return
        else:
            instance = instances[0]

        if hasattr(args, "cmd") and args.cmd:
            # Run single command
            print(f"Running on {instance.ip}: {args.cmd}")
            result = client.ssh_run(instance, args.cmd, timeout=args.timeout)
            if result.stdout:
                print(result.stdout)
            if result.stderr:
                print(f"[stderr] {result.stderr}", file=sys.stderr)
            if result.returncode != 0:
                sys.exit(result.returncode)
        else:
            # Print SSH command for interactive use
            print(client.get_ssh_command(instance))

    elif args.command == "rsync":
        # Rsync files to/from Lambda instance
        instances = client.list_instances()
        if not instances:
            print("No running instances.")
            return

        if args.instance_id:
            instance = next(
                (i for i in instances if i.id.startswith(args.instance_id)), None
            )
            if not instance:
                print(f"Instance {args.instance_id} not found.")
                return
        else:
            instance = instances[0]

        # Parse source and dest - 'remote:' prefix indicates remote path
        source = args.source
        dest = args.dest

        if source.startswith("remote:"):
            source = f"ubuntu@{instance.ip}:{source[7:]}"
        if dest.startswith("remote:"):
            dest = f"ubuntu@{instance.ip}:{dest[7:]}"

        rsync_cmd = [
            "rsync",
            "-avz",
            "--progress",
            "-e",
            "ssh -o StrictHostKeyChecking=no",
        ]
        if args.delete:
            rsync_cmd.append("--delete")
        rsync_cmd.extend([source, dest])

        print(f"Running: {' '.join(rsync_cmd)}")
        result = subprocess.run(rsync_cmd)
        sys.exit(result.returncode)

    elif args.command == "setup":
        ssh_key = setup_lambda_ssh_key(client)
        print(f"SSH key '{ssh_key}' is configured.")

    elif args.command == "train":
        # Full automated training pipeline
        import time as time_module

        instance = None
        start_time = time_module.time()
        training_completed = False  # Track if training actually finished

        # --- Step 0: Auto-convert demos to bundle if --demo-dir provided ---
        if args.demo_dir and not args.bundle:
            from openadapt_ml.training.convert_demos import prepare_bundle

            if not args.captures_dir:
                print("Error: --captures-dir is required with --demo-dir")
                return

            print("=" * 50)
            print("Step 0: Converting demos to training bundle")
            print("=" * 50)
            try:
                bundle_dir = prepare_bundle(
                    demo_dir=args.demo_dir,
                    captures_dir=args.captures_dir,
                    mapping_path=args.mapping,
                )
                args.bundle = str(bundle_dir)
                print(f"Bundle ready at: {bundle_dir}\n")
            except Exception as e:
                print(f"Error converting demos: {e}")
                return

        # Instance pricing (approximate $/hr)
        INSTANCE_PRICES = {
            "gpu_1x_a10": 0.75,
            "gpu_1x_a100": 1.29,
            "gpu_1x_a100_sxm4": 1.29,
            "gpu_1x_h100_pcie": 2.49,
            "gpu_1x_h100_sxm5": 3.29,
        }

        # Get or launch instance
        if args.instance:
            instances = client.list_instances()
            instance = next(
                (i for i in instances if i.id.startswith(args.instance)), None
            )
            if not instance:
                print(f"Error: Instance {args.instance} not found")
                return
        else:
            # Check for existing instances
            instances = client.list_instances()
            if instances:
                print(f"Using existing instance: {instances[0].id[:8]}...")
                instance = instances[0]
            else:
                # Launch new instance (with optional GPU wait/retry)
                ssh_key = setup_lambda_ssh_key(client)
                gpu_wait = getattr(args, "gpu_wait", 0)
                deadline = time_module.time() + gpu_wait * 60
                retry_interval = 60  # check every minute

                while True:
                    try:
                        print(f"Launching {args.type}...")
                        instance = client.launch_instance(
                            instance_type=args.type,
                            ssh_key_names=[ssh_key],
                            name="openadapt-training",
                        )
                        print(
                            f"Instance launched: {instance.id[:8]}... at {instance.ip}"
                        )
                        break
                    except RuntimeError as e:
                        if "No regions available" not in str(e) or gpu_wait <= 0:
                            raise
                        remaining = deadline - time_module.time()
                        if remaining <= 0:
                            print(f"Error: GPU not available after {gpu_wait} min wait")
                            return
                        mins_left = int(remaining / 60)
                        print(
                            f"  GPU unavailable, retrying in {retry_interval}s "
                            f"({mins_left} min remaining)..."
                        )
                        time_module.sleep(retry_interval)

        price_per_hour = INSTANCE_PRICES.get(instance.instance_type, 1.00)
        print(f"  Instance type: {instance.instance_type} (~${price_per_hour:.2f}/hr)")
        print(f"  Max runtime: {args.max_runtime} minutes")

        # Generate initial dashboard with setup status
        from pathlib import Path
        from openadapt_ml.training.trainer import (
            TrainingState,
            TrainingConfig,
            generate_training_dashboard,
            setup_job_directory,
        )
        import time as time_module

        job_id = time_module.strftime("%Y%m%d_%H%M%S")
        output_dir = setup_job_directory("training_output", job_id)
        dashboard_path = output_dir / "dashboard.html"
        log_path = output_dir / "training_log.json"

        def update_dashboard(
            status: str, logs: list, step: int = 0, loss: float = 0.0, epoch: int = 0
        ):
            """Update dashboard with current setup/training status."""
            state = TrainingState(job_id=job_id)
            state.cloud_provider = "lambda"
            state.cloud_dashboard_url = "https://cloud.lambda.ai/instances"
            state.cloud_instance_id = instance.id
            state.instance_ip = instance.ip or ""
            state.instance_type = instance.instance_type
            state.setup_status = status
            state.setup_logs = logs
            state.epoch = epoch
            state.step = step
            state.loss = loss
            state.start_time = start_time
            config = TrainingConfig(num_train_epochs=5, learning_rate=5e-5)
            dashboard_path.write_text(generate_training_dashboard(state, config))
            # Also write log for polling
            log_path.write_text(json.dumps(state.to_dict(), indent=2))

        # Initial dashboard
        setup_logs = [
            f"Lambda Cloud instance: {instance.id[:8]}...",
            f"Instance type: {instance.instance_type} (~${price_per_hour:.2f}/hr)",
            f"IP address: {instance.ip or 'pending...'}",
        ]
        update_dashboard("booting", setup_logs)

        # Open dashboard in browser via HTTP server
        server_proc = None
        if args.open:
            server_proc = open_dashboard_in_browser(output_dir)

        try:
            # Set up environment with retries at the command level
            setup_logs.append("Connecting to instance...")
            update_dashboard("booting", setup_logs)

            setup_success = False
            for setup_attempt in range(3):
                setup_logs.append(f"Setup attempt {setup_attempt + 1}/3...")
                update_dashboard("installing", setup_logs)
                if client.setup_instance(instance):
                    setup_success = True
                    setup_logs.append("Instance setup complete!")
                    update_dashboard("installing", setup_logs)
                    break
                if setup_attempt < 2:
                    setup_logs.append(
                        f"Setup attempt {setup_attempt + 1} failed, retrying in 30s..."
                    )
                    update_dashboard("booting", setup_logs)
                    print(
                        f"  Setup attempt {setup_attempt + 1} failed, retrying in 30s..."
                    )
                    time_module.sleep(30)

            if not setup_success:
                setup_logs.append("ERROR: Failed to set up instance after 3 attempts")
                update_dashboard("booting", setup_logs)
                print("\nError: Failed to set up instance after 3 attempts")
                print(f"Instance still running: {instance.ip}")
                print("Debug via: ssh ubuntu@" + instance.ip)
                print(
                    f"Terminate with: python -m openadapt_ml.cloud.lambda_labs terminate {instance.id}"
                )
                return  # Don't terminate - let user debug

            # Sync local code to ensure remote has latest changes
            setup_logs.append("Syncing local code to instance...")
            update_dashboard("installing", setup_logs)
            if not client.sync_local_code(instance):
                setup_logs.append(
                    "Warning: Failed to sync local code, using remote repo version"
                )
                update_dashboard("installing", setup_logs)
                print("Warning: Failed to sync local code, using remote repo version")
            else:
                setup_logs.append("Code synced successfully")
                update_dashboard("installing", setup_logs)

            # Upload capture if provided
            remote_capture = None
            if args.capture:
                setup_logs.append("Uploading capture data...")
                update_dashboard("installing", setup_logs)
                if client.upload_capture(instance, args.capture, "~/capture"):
                    remote_capture = "~/capture"
                    setup_logs.append(f"Capture uploaded to {instance.ip}:~/capture")
                    update_dashboard("installing", setup_logs)
                    print(f"Capture uploaded to {instance.ip}:~/capture")
                else:
                    setup_logs.append("ERROR: Failed to upload capture after retries")
                    update_dashboard("installing", setup_logs)
                    print("\nError: Failed to upload capture after retries")
                    print(f"Instance still running: {instance.ip}")
                    print("Debug via: ssh ubuntu@" + instance.ip)
                    print(
                        f"Terminate with: python -m openadapt_ml.cloud.lambda_labs terminate {instance.id}"
                    )
                    return  # Don't terminate - let user debug

            # Upload bundle if provided
            remote_jsonl = None
            if args.bundle:
                setup_logs.append("Uploading training bundle...")
                update_dashboard("installing", setup_logs)
                if client.upload_capture(instance, args.bundle, "~/training_data"):
                    remote_jsonl = "~/training_data/training_data.jsonl"
                    setup_logs.append(
                        f"Bundle uploaded to {instance.ip}:~/training_data"
                    )
                    update_dashboard("installing", setup_logs)
                    print(f"Bundle uploaded to {instance.ip}:~/training_data")
                else:
                    setup_logs.append("ERROR: Failed to upload bundle after retries")
                    update_dashboard("installing", setup_logs)
                    print("\nError: Failed to upload bundle after retries")
                    print(f"Instance still running: {instance.ip}")
                    print("Debug via: ssh ubuntu@" + instance.ip)
                    print(
                        f"Terminate with: python -m openadapt_ml.cloud.lambda_labs terminate {instance.id}"
                    )
                    return  # Don't terminate - let user debug

            # Run training in background and poll for status
            setup_logs.append("Installing dependencies and starting training...")
            update_dashboard("training", setup_logs)
            print("\n" + "=" * 50)
            print("Starting training...")
            print("=" * 50 + "\n")

            client.run_training(
                instance,
                config=args.config,
                capture=remote_capture,
                goal=args.goal,
                jsonl=remote_jsonl,
                background=True,  # Run in background so we can poll
            )

            # Poll for training status and update dashboard
            poll_interval = 10  # seconds
            sync_interval = 300  # sync training_log.json every 5 minutes
            last_step = 0
            last_epoch = 0
            last_sync_time = time_module.time()
            print(
                f"Polling training status every {poll_interval}s (Ctrl+C to stop)...\n"
            )

            while True:
                try:
                    status = client.get_training_status(instance)

                    if status and status.get("step", 0) > 0:
                        step = status.get("step", 0)
                        epoch = status.get("epoch", 0)
                        loss = status.get("loss", 0)
                        elapsed_training = status.get("elapsed_time", 0)
                        total_epochs = status.get("total_epochs", 5)

                        # Print progress when step changes
                        if step > last_step or epoch > last_epoch:
                            print(
                                f"  Epoch {epoch + 1}/{total_epochs} | Step {step} | Loss: {loss:.4f} | Elapsed: {elapsed_training:.0f}s"
                            )
                            last_step = step
                            last_epoch = epoch

                        # Update local training_log.json (dashboard polls this)
                        status["total_epochs"] = total_epochs
                        if not status.get("instance_ip"):
                            status["instance_ip"] = instance.ip
                        if not status.get("instance_type"):
                            status["instance_type"] = instance.instance_type
                        # Add cloud provider info
                        status["cloud_provider"] = "lambda"
                        status["cloud_dashboard_url"] = (
                            "https://cloud.lambda.ai/instances"
                        )
                        status["cloud_instance_id"] = instance.id
                        status["setup_status"] = "training"
                        status["setup_logs"] = setup_logs
                        log_path.write_text(json.dumps(status, indent=2))

                        # Regenerate dashboard with updated data
                        state = TrainingState()
                        state.job_id = status.get("job_id", "")
                        state.hostname = status.get("hostname", "lambda")
                        state.instance_ip = instance.ip or ""
                        state.instance_type = instance.instance_type
                        state.epoch = epoch
                        state.step = step
                        state.total_epochs = total_epochs
                        state.loss = loss
                        state.learning_rate = status.get("learning_rate", 5e-5)
                        state.losses = status.get("losses", [])
                        state.evaluations = status.get("evaluations", [])
                        state.start_time = time_module.time() - elapsed_training
                        state.cloud_provider = "lambda"
                        state.cloud_dashboard_url = "https://cloud.lambda.ai/instances"
                        state.cloud_instance_id = instance.id
                        state.setup_status = "training"
                        state.setup_logs = setup_logs

                        config = TrainingConfig(
                            num_train_epochs=total_epochs,
                            learning_rate=status.get("learning_rate", 5e-5),
                        )
                        dashboard_path.write_text(
                            generate_training_dashboard(state, config)
                        )

                        # Check if training is complete (all epochs done)
                        if epoch >= total_epochs - 1:
                            # Check if step count stopped increasing
                            time_module.sleep(poll_interval)
                            new_status = client.get_training_status(instance)
                            if new_status and new_status.get("step", 0) == step:
                                print("\n" + "=" * 50)
                                print("Training complete!")
                                print("=" * 50)
                                training_completed = True
                                break
                    else:
                        # Training not started yet, show setup status
                        print("  Waiting for training to start...")

                except Exception as e:
                    print(f"  Poll error: {e}")

                # Periodic sync of training artifacts (every 5 min)
                now = time_module.time()
                if now - last_sync_time >= sync_interval:
                    try:
                        subprocess.run(
                            [
                                "rsync",
                                "-az",
                                "-e",
                                "ssh -o StrictHostKeyChecking=no -o ConnectTimeout=10",
                                f"ubuntu@{instance.ip}:~/openadapt-ml/training_output/training_log.json",
                                str(log_path),
                            ],
                            capture_output=True,
                            timeout=30,
                        )
                    except Exception:
                        pass  # Non-critical — best-effort sync
                    last_sync_time = now

                time_module.sleep(poll_interval)

        except KeyboardInterrupt:
            print("\n\nTraining interrupted by user")
        finally:
            # Clean up HTTP server if running
            if server_proc:
                server_proc.terminate()
                print("Dashboard server stopped.")

        # Only auto-terminate if training completed successfully or user requested it
        elapsed = time_module.time() - start_time
        cost = (elapsed / 3600) * price_per_hour

        if training_completed and not args.no_terminate:
            # Run comparison on Lambda before downloading and terminating (if capture was provided)
            if args.capture:
                print("\n" + "=" * 50)
                print("Running comparison on Lambda instance...")
                print("=" * 50)

                # Determine the final checkpoint path (main checkpoint after training)
                checkpoint_path = (
                    "/home/ubuntu/openadapt-ml/checkpoints/qwen3vl2b_capture_lora"
                )

                # Check if checkpoint exists
                result = client.ssh_run(
                    instance,
                    f"ls {checkpoint_path}/adapter_config.json 2>/dev/null && echo 'exists'",
                    timeout=30,
                )

                if "exists" in result.stdout:
                    # Run comparison on Lambda
                    output_name = f"comparison_{time_module.strftime('%H%M%S')}.html"
                    cmd = f"""cd ~/openadapt-ml && source .venv/bin/activate && \
                        python -m openadapt_ml.scripts.compare \
                        --capture ~/capture \
                        --checkpoint {checkpoint_path} \
                        --output training_output/{output_name} 2>&1"""

                    print(
                        "  Generating comparison viewer (this may take a few minutes)..."
                    )
                    result = client.ssh_run(instance, cmd, timeout=600)

                    if result.returncode == 0:
                        print(f"  Comparison generated: {output_name}")
                    else:
                        print("  Warning: Comparison generation failed")
                        if result.stderr:
                            print(f"  Error: {result.stderr}")
                else:
                    print("  Warning: Final checkpoint not found, skipping comparison")

            # Download results (including comparison if generated)
            print("\n" + "=" * 50)
            print("Downloading results...")
            print("=" * 50)
            client.download_results(instance)

            print(f"\nTerminating instance {instance.id[:8]}...")
            client.terminate_instance(instance.id)
            print("Instance terminated.")
            print(f"\nFinal cost: ~${cost:.2f} ({elapsed / 60:.1f} minutes)")
        else:
            print(f"\nInstance still running: {instance.ip}")
            print(f"  Current cost: ~${cost:.2f}")
            if not training_completed:
                print("  (Not terminating - training did not complete successfully)")
            print(
                f"Terminate with: python -m openadapt_ml.cloud.lambda_labs terminate {instance.id}"
            )

    elif args.command == "train-status":
        instances = client.list_instances()
        if not instances:
            print("No running instances.")
            return

        if args.instance_id:
            instance = next(
                (i for i in instances if i.id.startswith(args.instance_id)), None
            )
            if not instance:
                print(f"Instance {args.instance_id} not found.")
                return
        else:
            instance = instances[0]

        print(f"Checking training status on {instance.ip}...")
        status = client.get_training_status(instance)

        if status:
            print(f"  Epoch: {status.get('epoch', 'N/A')}")
            print(f"  Step: {status.get('step', 'N/A')}")
            print(f"  Loss: {status.get('loss', 'N/A')}")
            print(f"  Elapsed: {status.get('elapsed_time', 0):.1f}s")
        else:
            print("  No training log found (training may not have started yet)")

    elif args.command == "checkpoints":
        # List checkpoints on remote instance
        instances = client.list_instances()
        if not instances:
            print("No running instances.")
            return

        if args.instance_id:
            instance = next(
                (i for i in instances if i.id.startswith(args.instance_id)), None
            )
            if not instance:
                print(f"Instance {args.instance_id} not found.")
                return
        else:
            instance = instances[0]

        print(f"Checking checkpoints on {instance.ip}...")

        ssh_cmd = [
            "ssh",
            "-o",
            "StrictHostKeyChecking=no",
            "-o",
            "ConnectTimeout=10",
            f"ubuntu@{instance.ip}",
            "ls -la ~/openadapt-ml/checkpoints/ 2>/dev/null && "
            "du -sh ~/openadapt-ml/checkpoints/ 2>/dev/null || echo 'No checkpoints directory found'",
        ]

        result = subprocess.run(ssh_cmd, capture_output=True, text=True)
        if result.returncode == 0:
            print(result.stdout)
        else:
            print("No checkpoints found yet")
            if result.stderr:
                print(f"  Error: {result.stderr}")

    elif args.command == "refresh":
        # One-shot dashboard refresh
        import time as time_module
        from pathlib import Path
        from openadapt_ml.training.trainer import (
            TrainingState,
            TrainingConfig,
            generate_training_dashboard,
        )

        instances = client.list_instances()
        if not instances:
            print("No running instances.")
            return

        if args.instance_id:
            instance = next(
                (i for i in instances if i.id.startswith(args.instance_id)), None
            )
            if not instance:
                print(f"Instance {args.instance_id} not found.")
                return
        else:
            instance = instances[0]

        # Use current job directory via symlink
        from openadapt_ml.training.trainer import (
            get_current_job_directory,
            setup_job_directory,
        )

        base_dir = Path("training_output")
        base_dir.mkdir(exist_ok=True)

        status = client.get_training_status(instance)

        if status and status.get("step", 0) > 0:
            # Get or create job directory based on remote job_id
            remote_job_id = status.get("job_id", "")
            if remote_job_id:
                output_dir = setup_job_directory(base_dir, remote_job_id)
            else:
                output_dir = get_current_job_directory(base_dir) or base_dir
            dashboard_path = output_dir / "dashboard.html"
            log_path = output_dir / "training_log.json"

            # Setup screenshots symlink if local capture path provided
            local_capture = (
                args.capture if hasattr(args, "capture") and args.capture else None
            )
            if local_capture:
                setup_capture_screenshots_symlink(output_dir, local_capture)

            # Rewrite evaluation paths from Lambda to relative
            if "evaluations" in status:
                status["evaluations"] = rewrite_evaluation_paths(status["evaluations"])

            # Ensure instance metadata is present
            status["instance_ip"] = instance.ip
            status["instance_type"] = instance.instance_type
            status["total_epochs"] = status.get("total_epochs", 5)

            # Save log
            log_path.write_text(json.dumps(status, indent=2))

            # Generate dashboard
            state = TrainingState(job_id=remote_job_id)
            state.job_id = remote_job_id
            state.hostname = status.get("hostname", "lambda")
            state.instance_ip = instance.ip or ""
            state.instance_type = instance.instance_type
            state.config_path = status.get("config_path", "")
            # Use local capture path for screenshots if provided, else remote path
            state.capture_path = (
                args.capture if args.capture else status.get("capture_path", "")
            )
            state.epoch = status.get("epoch", 0)
            state.step = status.get("step", 0)
            state.loss = status.get("loss", 0)
            state.learning_rate = status.get("learning_rate", 5e-5)
            state.losses = status.get("losses", [])
            state.evaluations = status.get("evaluations", [])
            state.total_epochs = status.get("total_epochs", 5)
            state.start_time = time_module.time() - status.get("elapsed_time", 0)
            # Cloud provider info
            state.cloud_provider = "lambda"
            state.cloud_dashboard_url = "https://cloud.lambda.ai/instances"
            state.cloud_instance_id = instance.id
            state.setup_status = status.get("setup_status", "training")
            state.setup_logs = status.get("setup_logs", [])

            config = TrainingConfig(
                num_train_epochs=status.get("total_epochs", 5),
                learning_rate=status.get("learning_rate", 5e-5),
            )

            dashboard_path.write_text(generate_training_dashboard(state, config))

            # Regenerate navigation for file:// protocol
            try:
                from openadapt_ml.training.trainer import regenerate_all_dashboards

                regenerate_all_dashboards(output_dir)
            except Exception:
                pass  # Silent fail for navigation

            epoch = status.get("epoch", 0)
            step = status.get("step", 0)
            loss = status.get("loss", 0)
            elapsed = status.get("elapsed_time", 0)
            print(
                f"Epoch {epoch + 1}/{state.total_epochs} | Step {step} | Loss: {loss:.4f} | Elapsed: {elapsed:.0f}s"
            )
            print(f"Dashboard: {dashboard_path.absolute()}")

            if args.open:
                import subprocess as sp

                sp.run(["open", str(dashboard_path)], capture_output=True)
        else:
            print("No training data yet")

    elif args.command == "monitor":
        # Live dashboard monitoring for Lambda training
        # Updates training_output/training_log.json so the existing dashboard auto-refreshes
        import time as time_module
        from pathlib import Path

        # Stub mode - simulate training without actual GPU
        if getattr(args, "stub", False):
            from openadapt_ml.training.stub_provider import StubTrainingProvider
            from openadapt_ml.training.trainer import (
                TrainingState,
                TrainingConfig,
                generate_training_dashboard,
            )

            print("\n[Stub Mode] Simulating training without GPU...")
            output_dir = Path("training_output")
            output_dir.mkdir(exist_ok=True)

            # Start dashboard server if requested
            server_proc = None
            if args.open:
                server_proc = open_dashboard_in_browser(output_dir)

            # Run stub training
            stub = StubTrainingProvider(
                output_dir=output_dir,
                epochs=5,
                steps_per_epoch=10,
                step_delay=0.3,  # Fast simulation
            )

            def update_dashboard(status):
                """Regenerate dashboard after each step."""
                state = TrainingState()
                state.job_id = status.get("job_id", "")
                state.hostname = status.get("hostname", "stub")
                state.instance_ip = "127.0.0.1"
                state.instance_type = "stub"
                state.epoch = status.get("epoch", 0)
                state.step = status.get("step", 0)
                state.loss = status.get("loss", 0)
                state.learning_rate = status.get("learning_rate", 5e-5)
                state.losses = status.get("losses", [])
                state.evaluations = status.get("evaluations", [])
                state.cloud_provider = "stub"
                state.setup_status = "training"

                config = TrainingConfig(
                    num_train_epochs=status.get("total_epochs", 5),
                    learning_rate=state.learning_rate,
                )

                dashboard_path = output_dir / "dashboard.html"
                dashboard_path.write_text(generate_training_dashboard(state, config))

            try:
                stub.run(callback=update_dashboard)
            except KeyboardInterrupt:
                print("\n[Stub] Interrupted by user.")
            finally:
                if server_proc:
                    server_proc.terminate()
                    print("[Stub] Dashboard server stopped.")

            print(f"\n[Stub] Results in: {output_dir}")
            return

        instances = client.list_instances()
        if not instances:
            print("No running instances.")
            return

        if args.instance_id:
            instance = next(
                (i for i in instances if i.id.startswith(args.instance_id)), None
            )
            if not instance:
                print(f"Instance {args.instance_id} not found.")
                return
        else:
            instance = instances[0]

        if instance.status == "booting" or not instance.ip:
            print(f"Instance {instance.id[:8]} is still booting, waiting for IP...")
            while True:
                time_module.sleep(5)
                instances = client.list_instances()
                instance = next((i for i in instances if i.id == instance.id), None)
                if not instance:
                    print("Instance terminated or not found.")
                    return
                if instance.ip and instance.status == "active":
                    print(f"Instance ready at {instance.ip}")
                    break
                print(f"  Status: {instance.status}...")

        # Use job-scoped directory structure
        from openadapt_ml.training.trainer import (
            TrainingState,
            TrainingConfig,
            generate_training_dashboard,
            setup_job_directory,
            get_current_job_directory,
        )

        base_dir = Path("training_output")
        base_dir.mkdir(exist_ok=True)

        # Get current job directory or wait for first status to determine job_id
        output_dir = get_current_job_directory(base_dir) or base_dir
        dashboard_path = output_dir / "dashboard.html"
        log_path = output_dir / "training_log.json"

        # Check for existing log with job_id
        current_job_id = None
        if log_path.exists():
            try:
                existing_log = json.loads(log_path.read_text())
                current_job_id = existing_log.get("job_id")
            except (json.JSONDecodeError, IOError):
                pass

        print(f"\nMonitoring Lambda training on {instance.ip}")
        print(f"Dashboard: {dashboard_path.absolute()}")
        print(f"Polling every {args.interval}s (Ctrl+C to stop)\n")

        # Generate initial dashboard if it doesn't exist
        if not dashboard_path.exists():
            state = TrainingState(job_id=current_job_id or "")
            state.cloud_provider = "lambda"
            state.cloud_dashboard_url = "https://cloud.lambda.ai/instances"
            state.cloud_instance_id = instance.id
            state.instance_ip = instance.ip or ""
            state.instance_type = instance.instance_type
            state.setup_status = "booting"
            state.setup_logs = [
                "Starting Lambda Cloud instance...",
                f"Instance ID: {instance.id[:8]}...",
                f"Instance type: {instance.instance_type}",
            ]
            config = TrainingConfig(num_train_epochs=5, learning_rate=5e-5)
            dashboard_path.write_text(generate_training_dashboard(state, config))

        # Open dashboard if requested via HTTP server
        server_proc = None
        if args.open:
            server_proc = open_dashboard_in_browser(output_dir)

        last_step = 0
        last_epoch = -1
        auto_stop_loss = getattr(args, "auto_stop_loss", 0.5)
        download_checkpoints = getattr(args, "download_checkpoints", True)
        step_stall_count = 0  # Track how many times step hasn't increased

        print(f"  Auto-stop loss threshold: {auto_stop_loss}")
        print(
            f"  Checkpoint download: {'enabled' if download_checkpoints else 'disabled'}"
        )

        try:
            while True:
                # Check for stop signal from dashboard
                if check_stop_signal(output_dir):
                    print("\n  Stop signal received from dashboard!")
                    print("  Downloading final checkpoints...")
                    if download_checkpoints:
                        download_checkpoints_from_instance(instance.ip, output_dir)

                    # Update status with termination info before terminating
                    termination_status = {
                        "termination_status": "user_stop",
                        "termination_message": "Training stopped by user via dashboard",
                    }
                    current_log = log_path.read_text() if log_path.exists() else "{}"
                    import json as json_module

                    current_data = json_module.loads(current_log)
                    current_data.update(termination_status)
                    log_path.write_text(json_module.dumps(current_data, indent=2))

                    print(f"  Terminating instance {instance.id}...")
                    client.terminate_instance(instance.id)
                    # Remove stop signal
                    (output_dir / "STOP_TRAINING").unlink(missing_ok=True)
                    print("  Training stopped by user.")
                    break

                try:
                    # Fetch training log from remote
                    status = client.get_training_status(instance)

                    if status and status.get("step", 0) > 0:
                        step = status.get("step", 0)
                        epoch = status.get("epoch", 0)
                        loss = status.get("loss", 0)
                        elapsed = status.get("elapsed_time", 0)
                        remote_job_id = status.get("job_id")

                        # Detect job_id change - clear old data if new job started
                        if (
                            remote_job_id
                            and current_job_id
                            and remote_job_id != current_job_id
                        ):
                            print(
                                f"\n  New job detected: {remote_job_id} (was: {current_job_id})"
                            )
                            print("  Clearing old job data...")
                            last_step = 0  # Reset step tracking
                            current_job_id = remote_job_id

                        # Update local training log (dashboard polls this file)
                        # Add total_epochs to status for dashboard
                        status["total_epochs"] = status.get("total_epochs", 5)
                        # Ensure instance metadata is present
                        if not status.get("instance_ip"):
                            status["instance_ip"] = instance.ip
                        if not status.get("instance_type"):
                            status["instance_type"] = instance.instance_type
                        # Add cloud provider info
                        status["cloud_provider"] = "lambda"
                        status["cloud_dashboard_url"] = (
                            "https://cloud.lambda.ai/instances"
                        )
                        status["cloud_instance_id"] = instance.id
                        status["setup_status"] = status.get("setup_status", "training")

                        # Setup screenshots symlink if local capture path provided
                        local_capture = (
                            args.capture
                            if hasattr(args, "capture") and args.capture
                            else None
                        )
                        if local_capture:
                            setup_capture_screenshots_symlink(output_dir, local_capture)

                        # Rewrite evaluation paths from Lambda to relative
                        if "evaluations" in status:
                            status["evaluations"] = rewrite_evaluation_paths(
                                status["evaluations"]
                            )

                        log_path.write_text(json.dumps(status, indent=2))

                        if step > last_step:
                            print(
                                f"  Epoch {epoch + 1} | Step {step} | Loss: {loss:.4f} | Elapsed: {elapsed:.0f}s"
                            )
                            last_step = step
                            step_stall_count = (
                                0  # Reset stall counter when step increases
                            )
                            if not current_job_id:
                                current_job_id = remote_job_id

                            # Regenerate dashboard with updated data
                            state = TrainingState()
                            state.job_id = status.get("job_id", "")
                            state.hostname = status.get("hostname", "lambda")
                            state.instance_ip = instance.ip or ""
                            state.instance_type = instance.instance_type
                            state.epoch = epoch
                            state.step = step
                            state.loss = loss
                            state.learning_rate = status.get("learning_rate", 5e-5)
                            state.losses = status.get("losses", [])
                            state.evaluations = status.get("evaluations", [])
                            state.start_time = time_module.time() - elapsed
                            # Cloud provider info
                            state.cloud_provider = "lambda"
                            state.cloud_dashboard_url = (
                                "https://cloud.lambda.ai/instances"
                            )
                            state.cloud_instance_id = instance.id
                            state.setup_status = status.get("setup_status", "training")
                            state.setup_logs = status.get("setup_logs", [])
                            state.termination_status = status.get(
                                "termination_status", ""
                            )
                            state.termination_message = status.get(
                                "termination_message", ""
                            )

                            config = TrainingConfig(
                                num_train_epochs=status.get("total_epochs", 5),
                                learning_rate=status.get("learning_rate", 5e-5),
                            )

                            dashboard_path.write_text(
                                generate_training_dashboard(state, config)
                            )

                            # Download checkpoints on epoch change
                            if download_checkpoints and epoch > last_epoch:
                                print(
                                    f"  Epoch {epoch + 1} completed - downloading checkpoints..."
                                )
                                if download_checkpoints_from_instance(
                                    instance.ip, output_dir
                                ):
                                    print(
                                        f"  Checkpoints saved to {output_dir}/checkpoints/"
                                    )
                                else:
                                    print("  Warning: checkpoint download failed")
                                last_epoch = epoch

                            # Auto-terminate when loss is low enough
                            if loss < auto_stop_loss and loss > 0:
                                print(
                                    f"\n  Loss {loss:.4f} < threshold {auto_stop_loss}"
                                )
                                print("  Downloading final checkpoints...")
                                if download_checkpoints:
                                    download_checkpoints_from_instance(
                                        instance.ip, output_dir
                                    )

                                # Update status with termination info
                                status["termination_status"] = "auto_low_loss"
                                status["termination_message"] = (
                                    f"Training auto-stopped: loss {loss:.4f} < threshold {auto_stop_loss}"
                                )
                                log_path.write_text(json.dumps(status, indent=2))

                                print(f"  Auto-terminating instance {instance.id}...")
                                client.terminate_instance(instance.id)
                                print("  Training completed (auto-stopped)!")
                                break
                        else:
                            # Step didn't increase - check if training is complete
                            step_stall_count += 1
                            total_epochs = status.get("total_epochs", 5)

                            # If on last epoch and step hasn't increased for 3 polls, training is complete
                            if epoch >= total_epochs - 1 and step_stall_count >= 3:
                                print(
                                    f"\n  Training complete (epoch {epoch + 1}/{total_epochs}, step stopped increasing)"
                                )
                                print("  Downloading final checkpoints...")
                                if download_checkpoints:
                                    download_checkpoints_from_instance(
                                        instance.ip, output_dir
                                    )

                                # Update status with termination info
                                status["termination_status"] = "auto_complete"
                                status["termination_message"] = (
                                    f"Training completed successfully ({epoch + 1}/{total_epochs} epochs)"
                                )
                                log_path.write_text(json.dumps(status, indent=2))

                                print(f"  Terminating instance {instance.id}...")
                                client.terminate_instance(instance.id)
                                print("  Instance terminated.")
                                break

                    else:
                        print("  Waiting for training to start...")

                except Exception as e:
                    print(f"  Poll error: {e}")

                time_module.sleep(args.interval)

        except KeyboardInterrupt:
            print("\n\nMonitoring stopped.")
            print(f"Dashboard: {dashboard_path.absolute()}")
        finally:
            # Clean up HTTP server if running
            if server_proc:
                server_proc.terminate()
                print("Dashboard server stopped.")

    elif args.command == "files":
        instances = client.list_instances()
        if not instances:
            print("No running instances.")
            return

        if args.instance_id:
            instance = next(
                (i for i in instances if i.id.startswith(args.instance_id)), None
            )
            if not instance:
                print(f"Instance {args.instance_id} not found.")
                return
        else:
            instance = instances[0]

        print(f"Files on {instance.ip} at {args.path}:")
        result = client.ssh_run(
            instance,
            f"find {args.path} -type f -name '*.pt' -o -name '*.json' -o -name '*.bin' 2>/dev/null | head -20",
            timeout=30,
        )
        if result.stdout:
            for line in result.stdout.strip().split("\n"):
                print(f"  {line}")
        else:
            print("  (no checkpoint files found)")

    elif args.command == "kill":
        # Kill training/inference processes
        instances = client.list_instances()
        if not instances:
            print("No running instances.")
            if args.local:
                print("\nKilling local Lambda-related processes...")
                subprocess.run(
                    ["pkill", "-f", "ssh.*ubuntu@.*openadapt"], capture_output=True
                )
                subprocess.run(["pkill", "-f", "lambda_labs"], capture_output=True)
                print("Done.")
            return

        if args.instance_id:
            instance = next(
                (i for i in instances if i.id.startswith(args.instance_id)), None
            )
            if not instance:
                print(f"Instance {args.instance_id} not found.")
                return
        else:
            instance = instances[0]

        print(f"Checking processes on {instance.ip}...")

        # List Python processes first
        result = client.ssh_run(
            instance,
            "ps aux | grep python | grep -v grep | grep -v jupyter",
            timeout=30,
        )
        if result.stdout.strip():
            print("Found Python processes:")
            for line in result.stdout.strip().split("\n"):
                print(f"  {line[:100]}...")
        else:
            print("No training/inference Python processes found.")
            return

        if args.all:
            print("\nKilling ALL Python processes (except jupyter)...")
            cmd = (
                "pkill -f 'python.*train\\|python.*compare\\|python.*openadapt' || true"
            )
        else:
            print("\nKilling training and inference processes...")
            cmd = "pkill -f 'python.*train' ; pkill -f 'python.*compare' || true"

        result = client.ssh_run(instance, cmd, timeout=30)
        print("Remote processes killed.")

        if args.local:
            print("\nKilling local Lambda-related processes...")
            subprocess.run(
                ["pkill", "-f", "ssh.*ubuntu@.*openadapt"], capture_output=True
            )
            subprocess.run(["pkill", "-f", "lambda_labs.*train"], capture_output=True)
            print("Local processes killed.")

        print("\nDone. Current status:")
        result = client.ssh_run(
            instance,
            "ps aux | grep python | grep -v grep | grep -v jupyter | wc -l",
            timeout=30,
        )
        count = result.stdout.strip()
        print(f"  {count} Python processes remaining on instance")

    elif args.command == "check":
        # Analyze training status and early stopping
        instances = client.list_instances()
        if not instances:
            print("No running instances.")
            return

        if args.instance_id:
            instance = next(
                (i for i in instances if i.id.startswith(args.instance_id)), None
            )
            if not instance:
                print(f"Instance {args.instance_id} not found.")
                return
        else:
            instance = instances[0]

        print(f"Checking training on {instance.ip}...")

        # Get training log
        result = client.ssh_run(
            instance,
            "cat ~/openadapt-ml/training_output/training_log.json 2>/dev/null",
            timeout=30,
        )

        if not result.stdout.strip():
            print("No training log found.")
            return

        try:
            data = json.loads(result.stdout)
            losses = data.get("losses", [])
        except json.JSONDecodeError:
            print("Could not parse training log.")
            return

        if not losses:
            print("No training data yet.")
            return

        total_steps = len(losses)
        epochs = sorted(set(loss["epoch"] for loss in losses))
        total_epochs = data.get("total_epochs", 5)
        min_loss = min(loss["loss"] for loss in losses)
        current_loss = losses[-1]["loss"]

        print(f"\n{'=' * 50}")
        print("TRAINING STATUS")
        print(f"{'=' * 50}")
        print(f"Steps: {total_steps}")
        print(f"Epochs: {max(epochs) + 1}/{total_epochs}")
        print(f"Current loss: {current_loss:.4f}")
        print(f"Min loss: {min_loss:.4f}")

        # Check if training is running
        proc_result = client.ssh_run(
            instance, "ps aux | grep 'python.*train' | grep -v grep | wc -l", timeout=30
        )
        is_running = int(proc_result.stdout.strip()) > 0

        if is_running:
            print("Status: RUNNING")
        else:
            print("Status: STOPPED")

        # Early stopping analysis
        window = min(args.window, len(losses))
        if window < 2:
            print("\nNot enough data for early stopping analysis.")
        else:
            recent_losses = [loss["loss"] for loss in losses[-window:]]
            older_losses = (
                [loss["loss"] for loss in losses[-window * 2 : -window]]
                if len(losses) >= window * 2
                else [loss["loss"] for loss in losses[:window]]
            )

            recent_avg = sum(recent_losses) / len(recent_losses)
            older_avg = (
                sum(older_losses) / len(older_losses) if older_losses else recent_avg
            )

            improvement = (older_avg - recent_avg) / older_avg if older_avg > 0 else 0
            loss_variance = max(recent_losses) - min(recent_losses)

            print(f"\n{'=' * 50}")
            print(f"EARLY STOPPING ANALYSIS (window={window})")
            print(f"{'=' * 50}")
            print(f"Recent avg loss: {recent_avg:.4f}")
            print(f"Prior avg loss: {older_avg:.4f}")
            print(f"Improvement: {improvement * 100:.2f}%")
            print(f"Loss variance: {loss_variance:.4f}")

            should_stop = improvement < args.threshold and loss_variance < 0.1
            if should_stop:
                print("\n⚠️  EARLY STOPPING RECOMMENDED")
                print(f"   Loss has plateaued (improvement < {args.threshold * 100}%)")
                if not is_running:
                    print("   (Training already stopped)")
                else:
                    print(
                        "\n   To stop: uv run python -m openadapt_ml.cloud.lambda_labs kill"
                    )
            else:
                print("\n✓ Training still improving, continue.")

        # Time estimate
        if is_running and len(losses) >= 2:
            avg_time_per_step = (
                losses[-1].get("time", 0) / len(losses)
                if losses[-1].get("time")
                else 50
            )
            steps_per_epoch = len(losses) / (max(epochs) + 1)
            remaining_epochs = total_epochs - max(epochs) - 1
            remaining_steps = remaining_epochs * steps_per_epoch
            eta_seconds = remaining_steps * avg_time_per_step
            eta_mins = eta_seconds / 60

            print(f"\n{'=' * 50}")
            print("TIME ESTIMATE")
            print(f"{'=' * 50}")
            print(f"Remaining epochs: {remaining_epochs}")
            print(f"Est. remaining steps: {remaining_steps:.0f}")
            print(f"ETA: {eta_mins:.1f} minutes")

    elif args.command == "compare":
        # Run comparison on Lambda and sync back
        instances = client.list_instances()
        if not instances:
            print("No running instances.")
            return

        if args.instance_id:
            instance = next(
                (i for i in instances if i.id.startswith(args.instance_id)), None
            )
            if not instance:
                print(f"Instance {args.instance_id} not found.")
                return
        else:
            instance = instances[0]

        # Determine checkpoint to use
        if args.checkpoint:
            checkpoint_path = args.checkpoint
        elif args.epoch is not None:
            checkpoint_path = (
                f"/home/ubuntu/openadapt-ml/checkpoints/epoch_{args.epoch}"
            )
        else:
            # Use latest (main checkpoint)
            checkpoint_path = (
                "/home/ubuntu/openadapt-ml/checkpoints/qwen3vl2b_capture_lora"
            )

        # Check if checkpoint exists
        result = client.ssh_run(
            instance,
            f"ls {checkpoint_path}/adapter_config.json 2>/dev/null && echo 'exists'",
            timeout=30,
        )
        if "exists" not in result.stdout:
            print(f"Checkpoint not found at {checkpoint_path}")
            # List available checkpoints
            result = client.ssh_run(
                instance, "ls -la ~/openadapt-ml/checkpoints/", timeout=30
            )
            print(f"Available checkpoints:\n{result.stdout}")
            return

        print(f"Running comparison on {instance.ip}...")
        print(f"Using checkpoint: {checkpoint_path}")

        # Run comparison on Lambda
        output_name = f"comparison_{time.strftime('%H%M%S')}.html"
        cmd = f"""cd ~/openadapt-ml && source .venv/bin/activate && \
            python -m openadapt_ml.scripts.compare \
            --capture ~/capture \
            --checkpoint {checkpoint_path} \
            --output training_output/{output_name} 2>&1"""

        print("Generating predictions (this may take a few minutes)...")
        result = client.ssh_run(instance, cmd, timeout=600)

        if result.returncode != 0:
            print(f"Comparison failed:\n{result.stderr}")
            return

        # Check if file was created
        result = client.ssh_run(
            instance, f"ls -la ~/openadapt-ml/training_output/{output_name}", timeout=30
        )
        if result.returncode != 0:
            print("Comparison file not created.")
            return

        print(f"Comparison generated: {output_name}")

        # Sync back to local
        local_output = Path("training_output") / output_name
        local_output.parent.mkdir(parents=True, exist_ok=True)

        print(f"Syncing to {local_output}...")
        subprocess.run(
            [
                "rsync",
                "-avz",
                f"ubuntu@{instance.ip}:~/openadapt-ml/training_output/{output_name}",
                str(local_output),
            ],
            capture_output=True,
        )

        print(f"Done! Comparison saved to: {local_output}")

        if args.open:
            subprocess.run(["open", str(local_output)], capture_output=True)
            print("Opened in browser.")

    elif args.command == "download":
        instances = client.list_instances()
        if not instances:
            print("No running instances.")
            return

        if args.instance_id:
            instance = next(
                (i for i in instances if i.id.startswith(args.instance_id)), None
            )
            if not instance:
                print(f"Instance {args.instance_id} not found.")
                return
        else:
            instance = instances[0]

        client.download_results(instance, local_path=args.output)

    elif args.command == "results":
        # Download results and generate comparison viewer
        instances = client.list_instances()
        if not instances:
            print("No running instances.")
            return

        if args.instance_id:
            instance = next(
                (i for i in instances if i.id.startswith(args.instance_id)), None
            )
            if not instance:
                print(f"Instance {args.instance_id} not found.")
                return
        else:
            instance = instances[0]

        # Download results
        print("Step 1: Downloading training results...")
        client.download_results(instance)

        # Generate comparison viewer
        print("\nStep 2: Generating comparison viewer...")
        checkpoint_path = "checkpoints_lambda/qwen3vl2b_capture_lora"

        import subprocess as sp

        cmd = [
            "uv",
            "run",
            "python",
            "-m",
            "openadapt_ml.scripts.compare",
            "--capture",
            args.capture,
            "--checkpoint",
            checkpoint_path,
        ]
        if args.goal:
            cmd.extend(["--goal", args.goal])
        if args.open:
            cmd.append("--open")

        result = sp.run(cmd)
        if result.returncode == 0:
            print("\nComparison viewer generated!")
            if not args.open:
                print(f"Open with: open {args.capture}/comparison.html")
        else:
            print("Warning: Failed to generate comparison viewer")

    elif args.command == "serve":
        # Start web server for live dashboard with stop button support
        import http.server
        import socketserver
        import time as time_module
        from pathlib import Path

        output_dir = (
            Path(args.output) if hasattr(args, "output") else Path("training_output")
        )
        port = args.port

        if not output_dir.exists():
            print(f"No {output_dir} directory. Run 'refresh' first.")
            return

        # Define handler with /api/stop support
        class Handler(http.server.SimpleHTTPRequestHandler):
            def __init__(self, *args, **kwargs):
                super().__init__(*args, directory=str(output_dir), **kwargs)

            def do_POST(self):
                if self.path == "/api/stop":
                    # Create stop signal file
                    stop_file = output_dir / "STOP_TRAINING"
                    stop_file.touch()
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Access-Control-Allow-Origin", "*")
                    self.end_headers()
                    self.wfile.write(b'{"status": "stop signal created"}')
                    print(f"  Stop signal created: {stop_file}")
                else:
                    self.send_error(404)

            def do_OPTIONS(self):
                # Handle CORS preflight
                self.send_response(200)
                self.send_header("Access-Control-Allow-Origin", "*")
                self.send_header("Access-Control-Allow-Methods", "POST, OPTIONS")
                self.send_header("Access-Control-Allow-Headers", "Content-Type")
                self.end_headers()

            def log_message(self, format, *args):
                pass  # Suppress log messages

        # Start web server
        with socketserver.TCPServer(("", port), Handler) as httpd:
            url = f"http://localhost:{port}/dashboard.html"
            print(f"\nDashboard server started at {url}")
            print("Press Ctrl+C to stop\n")

            if args.open:
                subprocess.run(["open", url], capture_output=True)

            try:
                httpd.serve_forever()
            except KeyboardInterrupt:
                print("\nServer stopped.")

    elif args.command == "sync":
        # Sync training output from Lambda and regenerate navigation for file:// protocol
        from pathlib import Path
        from openadapt_ml.training.trainer import (
            TrainingState,
            TrainingConfig,
            generate_training_dashboard,
            regenerate_all_dashboards,
        )

        instances = client.list_instances()
        if not instances:
            print("No running instances.")
            return

        if args.instance_id:
            instance = next(
                (i for i in instances if i.id.startswith(args.instance_id)), None
            )
            if not instance:
                print(f"Instance {args.instance_id} not found.")
                return
        else:
            instance = instances[0]

        output_dir = Path(args.output)
        output_dir.mkdir(exist_ok=True)

        print(f"Syncing training output from {instance.ip}...")

        # Sync all training output files
        rsync_cmd = [
            "rsync",
            "-avz",
            "--progress",
            "-e",
            "ssh -o StrictHostKeyChecking=no",
            f"ubuntu@{instance.ip}:~/openadapt-ml/training_output/",
            str(output_dir) + "/",
        ]
        result = subprocess.run(rsync_cmd, capture_output=False)

        if result.returncode != 0:
            print("Warning: rsync may have had issues")

        # Update dashboard with instance metadata
        log_path = output_dir / "training_log.json"
        dashboard_path = output_dir / "dashboard.html"

        if log_path.exists():
            try:
                import time as time_module

                status = json.loads(log_path.read_text())

                # Update with instance info
                status["instance_ip"] = instance.ip
                status["instance_type"] = instance.instance_type
                status["cloud_provider"] = "lambda"
                status["cloud_dashboard_url"] = "https://cloud.lambda.ai/instances"
                status["cloud_instance_id"] = instance.id

                log_path.write_text(json.dumps(status, indent=2))

                # Generate updated dashboard
                state = TrainingState()
                state.job_id = status.get("job_id", "")
                state.hostname = status.get("hostname", "lambda")
                state.instance_ip = instance.ip or ""
                state.instance_type = instance.instance_type
                state.config_path = status.get("config_path", "")
                state.capture_path = status.get("capture_path", "")
                state.epoch = status.get("epoch", 0)
                state.step = status.get("step", 0)
                state.loss = status.get("loss", 0)
                state.learning_rate = status.get("learning_rate", 5e-5)
                state.losses = status.get("losses", [])
                state.evaluations = status.get("evaluations", [])
                state.total_epochs = status.get("total_epochs", 5)
                state.start_time = time_module.time() - status.get("elapsed_time", 0)
                state.cloud_provider = "lambda"
                state.cloud_dashboard_url = "https://cloud.lambda.ai/instances"
                state.cloud_instance_id = instance.id

                config = TrainingConfig(
                    num_train_epochs=status.get("total_epochs", 5),
                    learning_rate=status.get("learning_rate", 5e-5),
                )

                dashboard_path.write_text(generate_training_dashboard(state, config))
            except Exception as e:
                print(f"Warning: Could not update dashboard: {e}")

        # Regenerate ALL dashboards with static navigation (for file:// protocol)
        print("Regenerating navigation links...")
        try:
            regenerated = regenerate_all_dashboards(output_dir)
            print(f"  Updated {len(regenerated)} files with static navigation")
        except Exception as e:
            print(f"Warning: Navigation regeneration failed: {e}")

        # Summary
        files = list(output_dir.glob("*.html"))
        print(f"\nSynced {len(files)} HTML files to {output_dir}/")
        for f in sorted(files):
            print(f"  - {f.name}")

        print(f"\nDashboard: {dashboard_path.absolute()}")

        if args.open:
            subprocess.run(["open", str(dashboard_path)], capture_output=True)

    elif args.command == "viewer":
        # Regenerate and open local viewer (no Lambda required)
        from pathlib import Path
        from openadapt_ml.training.trainer import regenerate_all_dashboards
        import re

        output_dir = Path(args.output)

        if not output_dir.exists():
            print(f"Error: {output_dir} does not exist")
            print("Run training or sync first to populate the directory.")
            return

        if not (output_dir / "training_log.json").exists():
            print(f"Error: No training_log.json found in {output_dir}")
            print("This directory doesn't contain training results.")
            return

        # Auto-link local screenshots if available
        screenshots_link = output_dir / "screenshots"
        if not screenshots_link.exists():
            # Try to find capture ID from training log or predictions
            try:
                capture_id = None

                # First try training log
                log_data = json.loads((output_dir / "training_log.json").read_text())
                capture_path = log_data.get("capture_path", "")
                capture_match = re.search(r"capture_(\d+)", capture_path)
                if capture_match:
                    capture_id = capture_match.group(1)

                # If not found, try predictions JSON files
                if not capture_id:
                    for pred_file in output_dir.glob("predictions_*.json"):
                        pred_data = json.loads(pred_file.read_text())
                        base_data = pred_data.get("base_data", [])
                        if base_data:
                            image_path = base_data[0].get("image_path", "")
                            capture_match = re.search(r"capture_(\d+)", image_path)
                            if capture_match:
                                capture_id = capture_match.group(1)
                                break

                if capture_id:
                    # Search for local screenshots in openadapt-capture
                    openadapt_capture_dir = (
                        Path.home() / "oa" / "src" / "openadapt-capture"
                    )
                    if openadapt_capture_dir.exists():
                        for capture_dir in openadapt_capture_dir.iterdir():
                            if capture_dir.is_dir():
                                screenshots_dir = capture_dir / "screenshots"
                                if screenshots_dir.exists():
                                    # Check if this capture has our screenshots
                                    sample_file = list(
                                        screenshots_dir.glob(
                                            f"capture_{capture_id}_step_*.png"
                                        )
                                    )
                                    if sample_file:
                                        print(
                                            f"Found local screenshots in {screenshots_dir}"
                                        )
                                        screenshots_link.symlink_to(screenshots_dir)
                                        print(
                                            f"  Linked: {screenshots_link} -> {screenshots_dir}"
                                        )
                                        break
            except Exception:
                pass  # Silently continue if auto-link fails

        print(f"Regenerating viewer from {output_dir}...")
        regenerated = regenerate_all_dashboards(output_dir)
        print(f"  Updated {len(regenerated)} files")

        # Show path info
        if args.dashboard:
            target = output_dir / "dashboard.html"
        else:
            target = output_dir / "viewer.html"

        print(f"\nGenerated: {target.absolute()}")
        print("View with: uv run python -m openadapt_ml.cloud.lambda_labs serve --open")

        if args.open:
            subprocess.run(["open", str(target)], capture_output=True)


if __name__ == "__main__":
    main()
