adam/tools/oasis_adapter.py

"""Execution bridge for the connected Oasis action world model trainer."""

from __future__ import annotations

import importlib.util
import json
import os
import queue
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Any

from adam.config import ConfigManager
from adam.executor import ToolCancelled, ToolContext, ToolExecutionError
from adam.oasis_dataset import dataset_directories, validate_oasis_dataset
from adam.process_control import set_process_tree_paused, terminate_process_tree


FORCE_STOP_TIMEOUT_SECONDS = 30
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}


def _external_oasis_root(context: ToolContext) -> Path | None:
    path = context.root / "config" / "external_tools.json"
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None
    for entry in payload.get("tools", []):
        if not isinstance(entry, dict):
            continue
        if str(entry.get("id", "")) != "external_oasis_game_trainer":
            continue
        root = Path(str(entry.get("backend", {}).get("root", ""))).expanduser()
        return root if root.is_dir() else None
    return None


def _oasis_root(context: ToolContext) -> Path:
    folders = ConfigManager(context.root).get("tool_folders", {})
    raw = str(folders.get("oasis_trainer", "")) if isinstance(folders, dict) else ""
    root = Path(raw).expanduser() if raw else (_external_oasis_root(context) or Path())
    script = root / "roblox_action_flow_app.py"
    if not script.is_file():
        raise ToolExecutionError(
            "Oasis roblox_action_flow_app.py was not found. Connect the Oasis-Game-Trainer folder in Settings."
        )
    return root.resolve()


def _safe_model_name(value: str) -> str:
    name = value.strip()
    if not name or len(name) > 96 or any(char in name for char in "<>:\"/\\|?*\x00"):
        raise ToolExecutionError("Choose a short Oasis model name without filesystem-reserved characters.")
    return name


def _valid_oasis_resolution(value: str) -> str:
    """Validate the frame-size invariant required by the external trainer."""
    try:
        width_text, height_text = str(value).lower().split("x")
        width, height = int(width_text), int(height_text)
    except (TypeError, ValueError):
        raise ToolExecutionError(
            "Oasis resolution must be WIDTHxHEIGHT, for example 256x144 or 512x288."
        ) from None
    if width <= 0 or height <= 0 or width * 9 != height * 16 or width % 16 or height % 16:
        raise ToolExecutionError(
            "Oasis resolution must be 16:9 and both dimensions must be divisible by 16 "
            "(use 256x144 or 512x288)."
        )
    return f"{width}x{height}"


def _valid_action_model(path: Path) -> bool:
    try:
        info = json.loads((path / "action_flow_model_info.json").read_text(encoding="utf-8"))
        model_type = info.get("model_type")
        has_unet = (path / "unet" / "config.json").is_file()
        has_vae = (path / "vae" / "config.json").is_file() and (path / "vae" / "pytorch_model.bin").is_file()
        return model_type in {
            "action_conditioned_rectified_flow_video",
            "action_conditioned_latent_vae_flow_video",
            "action_conditioned_temporal_latent_flow",
            "action_conditioned_temporal_pixel_flow",
        } and has_unet and (model_type in {"action_conditioned_rectified_flow_video", "action_conditioned_temporal_pixel_flow"} or has_vae)
    except (OSError, ValueError, TypeError, json.JSONDecodeError):
        return False


def _latest_preview(folder: Path) -> Path | None:
    try:
        images = [
            path for path in folder.rglob("*")
            if path.is_file()
            and path.suffix.casefold() in IMAGE_EXTENSIONS
            and any(token in path.name.casefold() for token in ("preview", "sample", "epoch"))
        ]
        return max(images, key=lambda path: path.stat().st_mtime) if images else None
    except OSError:
        return None


def _parse_event(context: ToolContext, event: dict[str, Any], epochs: int, output: Path, preview_enabled: bool, preview_every: int, preview_steps: int) -> None:
    kind = str(event.get("type", ""))
    if kind == "start":
        context.log(
            f"Oasis dataset ready: {event.get('transitions', 0):,} transitions, "
            f"{event.get('training', 0):,} training, {event.get('validation', 0):,} validation on {event.get('device', 'device')}."
        )
        context.progress(1, "Starting Oasis trainer")
    elif kind == "progress":
        update = int(event.get("update", 0) or 0)
        total = int(event.get("total_updates", 0) or 0)
        epoch = int(event.get("epoch", 0) or 0)
        loss = event.get("loss")
        eta = event.get("eta")
        message = f"Oasis epoch {epoch} of {epochs}"
        if total:
            message += f", step {update:,} of {total:,}"
        if isinstance(loss, (int, float)):
            message += f", loss {loss:.4f}"
        context.progress(
            max(1, min(99, round(update * 100 / total))) if total else max(1, min(99, round(epoch * 100 / max(epochs, 1)))),
            message,
            current_step=update,
            total_steps=total,
            epoch=epoch,
            total_epochs=epochs,
            eta_seconds=eta if isinstance(eta, (int, float)) else None,
            unit="step" if total else "epoch",
        )
    elif kind == "validation":
        validation = event.get("validation", {})
        loss = validation.get("total_loss") if isinstance(validation, dict) else None
        context.log(f"Oasis validation epoch {event.get('epoch')}: loss {loss:.4f}" if isinstance(loss, (int, float)) else f"Oasis validation epoch {event.get('epoch')}.")
    elif kind == "temporal_validation":
        metrics = event.get("metrics", {})
        rollout = metrics.get("rollout_rgb_l1") if isinstance(metrics, dict) else None
        context.log(
            f"Temporal Oasis validation epoch {event.get('epoch')}: "
            + (f"generated rollout RGB L1 {rollout}." if rollout else "completed.")
        )
    elif kind == "preview" and preview_enabled:
        path = Path(str(event.get("path", ""))) if event.get("path") else _latest_preview(output)
        if path:
            epoch = int(event.get("epoch", 0) or 0)
            context.preview(path, epoch=epoch, next_epoch=epoch + preview_every, steps=preview_steps, kind="training")
    elif kind == "saved":
        context.log(f"Oasis checkpoint saved at epoch {event.get('epoch')}.")
    elif kind == "best_checkpoint":
        context.log(
            f"Oasis selected epoch {event.get('epoch')} as the current best held-out checkpoint."
        )
    elif kind == "recovery_saved":
        context.log(f"Oasis recovery checkpoint saved at batch {event.get('batch')}.")
    elif kind == "warning":
        context.log(str(event.get("message", "Oasis trainer warning.")))
    elif kind == "complete":
        context.progress(100, "Oasis training completed")


def train_oasis(
    context: ToolContext,
    dataset_dir: str | list[str],
    model_name: str,
    epochs: int,
    output_dir: str,
    resume_from: str = "",
    resolution: str = "256x144",
    batch_size: int = 2,
    learning_rate: float = 0.00002,
    workers: int = 2,
    mixed_precision: str = "fp16",
    gradient_accumulation: int = 1,
    frame_gap: int = 1,
    sequence_context: int = 1,
    model_engine: str = "pixel_flow",
    context_frames: int = 4,
    rollout_frames: int = 3,
    vae_epochs: int = 5,
    action_aggregation: str = "window",
    validation_split: float = 0.1,
    validation_batches: int = 8,
    save_every: int = 5,
    best_checkpoint_min_improvement: float = 0.1,
    preview_enabled: bool = True,
    preview_every: int = 5,
    preview_steps: int = 1,
    seed: int = 1234,
    base_model: str = "",
    condition_noise: float = 0.03,
    temporal_loss_weight: float = 0.1,
    motion_loss_weight: float = 2.0,
    action_input_scale: float = 8.0,
    neutral_action_dropout: float = 0.15,
    action_contrast_weight: float = 0.35,
    action_contrast_margin: float = 0.02,
    contrast_every: int = 4,
    contrast_samples: int = 2,
    chunk_size: int = 0,
    chunk_mode: str = "balanced",
    chunk_offset: int = 0,
    replay_older_percent: float = 50.0,
    include_older_data: bool = True,
    recovery_minutes: int = 30,
    benchmark_batches: int = 8,
    tf32: bool = True,
    gradient_checkpointing: bool = False,
    balance_actions: bool = True,
) -> dict[str, object]:
    root = _oasis_root(context)
    script = root / "roblox_action_flow_app.py"
    model_name = _safe_model_name(model_name)
    resolution = _valid_oasis_resolution(resolution)
    dataset_value = (
        ";".join(str(item) for item in dataset_dir)
        if isinstance(dataset_dir, list)
        else str(dataset_dir)
    )
    if not include_older_data:
        directories = dataset_directories(dataset_value)
        if directories:
            dataset_value = str(directories[0])
    output = Path(output_dir).expanduser().resolve()
    output_root = (root / "output_action_flow_models").resolve()
    try:
        output.relative_to(output_root)
    except ValueError as exc:
        raise ToolExecutionError("Oasis outputs must stay inside output_action_flow_models.") from exc
    resume = Path(resume_from).expanduser().resolve() if resume_from else None
    if resume:
        try:
            output.relative_to(resume)
            overlaps_resume = True
        except ValueError:
            try:
                resume.relative_to(output)
                overlaps_resume = True
            except ValueError:
                overlaps_resume = False
        if overlaps_resume:
            raise ToolExecutionError("Oasis continuation output must be a new folder, separate from the source model.")
    if output.exists():
        raise ToolExecutionError("The chosen Oasis output folder already exists; ADAM will not overwrite it.")
    if model_engine not in {"pixel_flow", "vae_cpu_lite", "temporal_latent", "temporal_pixel_flow"}:
        raise ToolExecutionError("Choose pixel_flow, vae_cpu_lite, temporal_latent, or temporal_pixel_flow for Oasis.")
    if mixed_precision not in {"fp16", "no", "fp32"}:
        raise ToolExecutionError("Oasis precision must be fp16 or no.")
    if mixed_precision == "fp32":
        mixed_precision = "no"
    if chunk_mode not in {"balanced", "random", "sequential"}:
        raise ToolExecutionError("Oasis chunk selection must be balanced, random, or sequential.")
    if model_engine in {"temporal_latent", "temporal_pixel_flow"}:
        if int(context_frames) not in {1, 4, 8} or int(rollout_frames) not in {1, 3}:
            raise ToolExecutionError("Temporal Oasis uses 1, 4, or 8 history frames and 1 or 3 future frames.")
        # The external trainer owns the timing contract; use its actual sequence
        # inspection instead of accepting old endpoint-only validation here.
        if str(root) not in sys.path:
            sys.path.insert(0, str(root))
        try:
            from temporal_oasis import read_sequences
            from roblox_action_flow_app import BINARY_ACTION_NAMES
            train_windows, val_windows, temporal_report = read_sequences(
                dataset_directories(dataset_value), BINARY_ACTION_NAMES,
                int(context_frames), int(rollout_frames), float(validation_split),
            )
        except (OSError, ValueError, ImportError) as exc:
            raise ToolExecutionError(f"Temporal Oasis dataset validation failed: {exc}") from exc
        context.log(
            f"Temporal Oasis dataset validated: {len(train_windows):,} training windows, "
            f"{len(val_windows):,} validation windows, {temporal_report['context_frames']} frames of history."
        )
    else:
        report = validate_oasis_dataset(dataset_value, frame_gap=int(frame_gap))
        if not report.ok:
            raise ToolExecutionError("Oasis dataset validation failed: " + " ".join(report.errors[:12]))
        context.log(
            f"Oasis dataset validated: {report.valid_rows:,} labelled frames, "
            f"{report.valid_transitions:,} transitions, resolution {report.resolution}."
        )
    missing_packages = [
        package for package in ("torch", "torchvision", "diffusers", "PIL")
        if importlib.util.find_spec(package) is None
    ]
    if missing_packages:
        raise ToolExecutionError(
            "ADAM's Python environment is missing Oasis packages: "
            + ", ".join(missing_packages)
            + ". Install the Oasis requirements, then restart ADAM."
        )
    output.parent.mkdir(parents=True, exist_ok=True)
    if resume and not _valid_action_model(resume):
        raise ToolExecutionError("Choose a valid Oasis action model folder to resume.")
    base = Path(base_model).expanduser().resolve() if base_model and model_engine == "pixel_flow" else None
    command = [
        sys.executable, str(script), "--train-worker",
        "--dataset-dir", dataset_value,
        "--output-dir", str(output),
        "--model-name", model_name,
        "--model-engine", model_engine,
        "--context-frames", str(int(context_frames)),
        "--rollout-frames", str(int(rollout_frames)),
        "--vae-epochs", str(int(vae_epochs)),
        "--epochs", str(int(epochs)),
        "--resolution", str(resolution),
        "--batch-size", str(int(batch_size)),
        "--learning-rate", str(float(learning_rate)),
        "--workers", str(int(workers)),
        "--gradient-accumulation", str(int(gradient_accumulation)),
        "--condition-noise", str(float(condition_noise)),
        "--temporal-loss-weight", str(float(temporal_loss_weight)),
        "--motion-loss-weight", str(float(motion_loss_weight)),
        "--validation-split", str(float(validation_split)),
        "--validation-batches", str(int(validation_batches)),
        "--metrics-every", "10",
        "--frame-gap", str(int(frame_gap)),
        "--action-aggregation", str(action_aggregation),
        "--action-input-scale", str(float(action_input_scale)),
        "--neutral-action-dropout", str(float(neutral_action_dropout)),
        "--action-contrast-weight", str(float(action_contrast_weight)),
        "--action-contrast-margin", str(float(action_contrast_margin)),
        "--contrast-every", str(max(1, int(contrast_every))),
        "--contrast-samples", str(max(2, int(contrast_samples))),
        "--chunk-size", str(max(0, int(chunk_size))),
        "--chunk-mode", chunk_mode,
        "--chunk-offset", str(max(0, int(chunk_offset))),
        "--replay-older-percent", str(max(0.0, float(replay_older_percent))),
        "--chunk-seed", str(int(seed)),
        "--benchmark-batches", str(max(1, int(benchmark_batches))),
        "--recovery-minutes", str(max(0, int(recovery_minutes))),
        "--save-every", str(int(save_every)),
        "--best-checkpoint-min-improvement", str(max(0.0, float(best_checkpoint_min_improvement)) / 100.0),
        "--preview-every", str(int(preview_every) if preview_enabled else int(epochs) + 1),
        "--preview-steps", str(int(preview_steps)),
        "--mixed-precision", str(mixed_precision),
    ]
    if resume:
        command.extend(["--continue-action-model", str(resume)])
    if base:
        command.extend(["--base-model", str(base)])
    if gradient_checkpointing:
        command.append("--gradient-checkpointing")
    if tf32:
        command.append("--tf32")
    if balance_actions:
        command.append("--balance-actions")
    process = subprocess.Popen(
        command, cwd=str(root), stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
        text=True, encoding="utf-8", errors="replace", shell=False,
        creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0,
    )
    lines: queue.Queue[str | None] = queue.Queue()

    def read_output() -> None:
        assert process.stdout is not None
        for line in process.stdout:
            lines.put(line.rstrip())
        lines.put(None)

    threading.Thread(target=read_output, daemon=True).start()
    context.progress(1, "Starting Oasis trainer")
    stopped = False
    stop_requested_at: float | None = None
    force_stop_sent = False
    suspended = False
    stop_file = output / "stop_action_training.flag"
    while True:
        should_pause = not context.run_event.is_set()
        if should_pause != suspended and set_process_tree_paused(process, should_pause):
            suspended = should_pause
            context.log("Oasis trainer paused safely." if suspended else "Oasis trainer resumed.")
        if context.cancel_event.is_set() and not stopped:
            if suspended:
                set_process_tree_paused(process, False)
                suspended = False
            output.mkdir(parents=True, exist_ok=True)
            stop_file.touch(exist_ok=True)
            stopped = True
            stop_requested_at = time.monotonic()
            context.log("Safe stop requested; waiting for Oasis to save a usable checkpoint.")
        if stop_requested_at and not force_stop_sent and time.monotonic() - stop_requested_at > FORCE_STOP_TIMEOUT_SECONDS:
            terminate_process_tree(process, timeout=3)
            force_stop_sent = True
            context.log("Oasis did not stop in time; terminating the trainer process.")
        try:
            line = lines.get(timeout=0.15)
            if line and line.startswith("ACTION_FLOW_EVENT:"):
                try:
                    _parse_event(context, json.loads(line.split(":", 1)[1]), int(epochs), output, bool(preview_enabled), int(preview_every), int(preview_steps))
                except json.JSONDecodeError:
                    context.log(line)
            elif line:
                context.log(line)
        except queue.Empty:
            pass
        if process.poll() is not None and lines.empty():
            break
    if stopped:
        raise ToolCancelled("Oasis training stopped by user.")
    if process.returncode != 0:
        raise ToolExecutionError(f"Oasis trainer exited with code {process.returncode}. See the job log for details.")
    if not _valid_action_model(output):
        raise ToolExecutionError("Oasis training ended, but no valid action model metadata was written.")
    context.progress(100, "Oasis training completed")
    return {
        "output_folder": str(output),
        "model_name": model_name,
        "assets": [{
            "kind": "model",
            "name": model_name,
            "path": str(output),
            "trainer": "oasis",
            "dataset_path": dataset_value,
            "checkpoint": str(output),
            "epochs": int(epochs),
        }],
    }


def launch_oasis_player(
    context: ToolContext,
    model_path: str,
    model_name: str = "",
    starting_frame: str = "",
    seed: int = 0,
) -> dict[str, object]:
    root = _oasis_root(context)
    model = Path(model_path).expanduser().resolve()
    if not _valid_action_model(model):
        raise ToolExecutionError("Choose a valid Oasis action model folder.")
    best_checkpoint = model / "best_validation_checkpoint"
    if _valid_action_model(best_checkpoint):
        context.log("Using the model's held-out best-validation checkpoint for playback.")
        model = best_checkpoint
    if starting_frame:
        start = Path(starting_frame).expanduser()
        if not start.is_file():
            raise ToolExecutionError("The selected Oasis starting frame does not exist.")
    command = [sys.executable, str(root / "roblox_action_flow_app.py")]
    env = os.environ.copy()
    env["ADAM_OASIS_MODEL_PATH"] = str(model)
    env["ADAM_OASIS_STARTING_FRAME"] = starting_frame
    env["ADAM_OASIS_SEED"] = str(int(seed))
    process = subprocess.Popen(
        command, cwd=str(root), env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
        shell=False, creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0,
    )
    context.log(f"Launched Oasis playable inference for {model_name or model.name}.")
    context.progress(100, "Oasis player launched")
    return {"player_pid": process.pid, "model_path": str(model)}