"""Script to train RL agent with RSL-RL."""

import json
import logging
import os
import sys
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Literal

# 训练器作为仓库内置子集直接从 scripts/ 启动，不要求额外执行 pip install -e。
TRAINER_ROOT = Path(__file__).resolve().parents[1]
SERVICE_ROOT = TRAINER_ROOT.parent
for source_root in (TRAINER_ROOT, SERVICE_ROOT):
  if str(source_root) not in sys.path:
    sys.path.insert(0, str(source_root))

import tyro
import warp as wp

# mjlab 1.2.0 的 GPU CUDA-graph 检查仍访问旧公开路径 wp.context；
# Warp 1.15 已把实现移到 warp._src.context，但保留了相同 runtime 契约。
if not hasattr(wp, "context"):
  from warp._src import context as warp_context

  wp.context = warp_context  # type: ignore[attr-defined]

from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
from mjlab.rl import MjlabOnPolicyRunner, RslRlBaseRunnerCfg, RslRlVecEnvWrapper
from mjlab.tasks.registry import list_tasks, load_env_cfg, load_rl_cfg, load_runner_cls
from mjlab.tasks.tracking.mdp import MotionCommandCfg
from mjlab.utils.gpu import select_gpus
from mjlab.utils.os import dump_yaml, get_checkpoint_path
from mjlab.utils.torch import configure_torch_backends
from mjlab.utils.wrappers import VideoRecorder
from task_config import OBSTACLE_TASK, deployment_metadata, validate_task_config
from tuning.schema import apply_reward_configuration, validate_configuration


@dataclass(frozen=True)
class TrainConfig:
  env: ManagerBasedRlEnvCfg
  agent: RslRlBaseRunnerCfg
  motion_file: str | None = None
  video: bool = False
  video_length: int = 200
  video_interval: int = 2000
  enable_nan_guard: bool = False
  torchrunx_log_dir: str | None = None
  gpu_ids: list[int] | Literal["all"] | None = field(default_factory=lambda: [0])
  output_dir: str | None = None
  resume_checkpoint: str | None = None
  pretrained_checkpoint: str | None = None
  pretrained_upload_manifest: str | None = None
  pretrained_onnx: str | None = None
  pretrained_source_id: str | None = None
  pretrained_allowed_roots: list[str] = field(default_factory=list)
  reward_config: str | None = None
  reward_config_json: str | None = None
  task_config: str | None = None

  @staticmethod
  def from_task(task_id: str) -> "TrainConfig":
    env_cfg = load_env_cfg(task_id)
    agent_cfg = load_rl_cfg(task_id)
    return TrainConfig(env=env_cfg, agent=agent_cfg)


def _load_reward_config(path: str | None, inline: str | None, task_id="Unitree-Go2-Flat") -> dict | None:
  if path is not None and inline is not None:
    raise ValueError("Use only one of reward_config and reward_config_json")
  if inline is not None:
    if len(inline.encode("utf-8")) > 64 * 1024:
      raise ValueError("Reward configuration is larger than 64 KiB")
    return validate_configuration(json.loads(inline), task_id)
  if path is None:
    return None
  source = Path(path).expanduser().resolve(strict=True)
  if source.stat().st_size > 64 * 1024:
    raise ValueError("Reward configuration is larger than 64 KiB")
  with source.open(encoding="utf-8") as stream:
    return validate_configuration(json.load(stream), task_id)


def _load_task_config(task_id: str, path: str | None, seed: int) -> dict | None:
  if path is None:
    return validate_task_config(task_id, {}, seed)
  source = Path(path).expanduser().resolve(strict=True)
  if source.stat().st_size > 128 * 1024:
    raise ValueError("Task configuration is larger than 128 KiB")
  payload = json.loads(source.read_text(encoding="utf-8"))
  allowed = {"terrainPreset", "terrainParams", "sensorCfg", "seed", "customTerrainBoxes"}
  if not isinstance(payload, dict) or payload.keys() - allowed:
    raise ValueError("Unknown task configuration fields")
  if (
    isinstance(payload.get("seed"), bool)
    or not isinstance(payload.get("seed"), int)
    or payload.get("seed") != seed
  ):
    raise ValueError("Task configuration seed must equal the agent seed")
  if payload.get("sensorCfg") is None:
    payload.pop("sensorCfg", None)
  return validate_task_config(task_id, payload, seed)


def _configure_task_and_rewards(task_id: str, cfg: TrainConfig):
  custom = _load_task_config(task_id, cfg.task_config, cfg.agent.seed)
  if custom is not None:
    from src.tasks.obstacle_avoidance.env_cfg import apply_obstacle_configuration
    from src.tasks.obstacle_avoidance.terrain import apply_terrain_configuration
    if task_id == OBSTACLE_TASK:
      apply_obstacle_configuration(cfg.env, custom)
    else:
      apply_terrain_configuration(cfg.env, custom)
  deployment = deployment_metadata(task_id, custom, cfg.agent.seed)
  reward_config = _load_reward_config(cfg.reward_config, cfg.reward_config_json, task_id)
  if reward_config is not None:
    apply_reward_configuration(cfg.env, reward_config, task_id)
    if task_id == OBSTACLE_TASK:
      deployment["navigation"]["speed"] = reward_config["params"]["target_velocity"]
      deployment["sensorCfg"]["avoidanceWeight"] = reward_config["weights"]["avoidance_weight"]

  return deployment, reward_config


def _load_pretrained(cfg: TrainConfig):
  if cfg.pretrained_checkpoint is None:
    if (cfg.pretrained_onnx is not None or cfg.pretrained_allowed_roots
        or cfg.pretrained_source_id or cfg.pretrained_upload_manifest):
      raise ValueError("Pretrained options require --pretrained-checkpoint (.pt), not ONNX alone")
    return None
  if cfg.resume_checkpoint is not None or cfg.agent.resume:
    raise ValueError("Pretrained warm-start and resume are mutually exclusive")
  from pretrained import read_pretrained_source

  options = {"onnx_path": cfg.pretrained_onnx}
  if cfg.pretrained_upload_manifest is not None:
    if cfg.pretrained_onnx is not None:
      raise ValueError("Uploaded actor cannot use adjacent ONNX sidecars")
    from pretrained_upload import read_uploaded_source
    read_pretrained_source = read_uploaded_source
    options = {"manifest_path": cfg.pretrained_upload_manifest}
  source = read_pretrained_source(
    cfg.pretrained_checkpoint,
    allowed_roots=cfg.pretrained_allowed_roots,
    **options,
    target_env=asdict(cfg.env),
    target_agent=asdict(cfg.agent),
  )
  if cfg.pretrained_source_id is not None and source.manifest["source_id"] != cfg.pretrained_source_id:
    raise ValueError("基础策略SHA身份变化，拒绝初始化")
  return source


def run_train(task_id: str, cfg: TrainConfig, log_dir: Path) -> None:
  deployment, reward_config = _configure_task_and_rewards(task_id, cfg)
  # Validate source identity/semantics before allocating a simulation or optimizer.
  pretrained = _load_pretrained(cfg)
  cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
  if cuda_visible == "":
    device = "cpu"
    seed = cfg.agent.seed
    rank = 0
  else:
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    rank = int(os.environ.get("RANK", "0"))
    # Set EGL device to match the CUDA device.
    os.environ["MUJOCO_EGL_DEVICE_ID"] = str(local_rank)
    device = f"cuda:{local_rank}"
    # Set seed to have diversity in different processes.
    seed = cfg.agent.seed + local_rank

  configure_torch_backends()

  cfg.agent.seed = seed
  cfg.env.seed = seed

  print(f"[INFO] Training with: device={device}, seed={seed}, rank={rank}")

  # Check if this is a tracking task by checking for motion command.
  is_tracking_task = "motion" in cfg.env.commands and isinstance(
    cfg.env.commands["motion"], MotionCommandCfg
  )

  if is_tracking_task:
    if not cfg.motion_file:
      raise ValueError("For tracking tasks, --motion-file must be set ...")
    motion_path = Path(cfg.motion_file).expanduser().resolve()
    if not motion_path.exists():
      raise FileNotFoundError(f"Motion file not found: {motion_path}")
    motion_cmd = cfg.env.commands["motion"]
    assert isinstance(motion_cmd, MotionCommandCfg)
    motion_cmd.motion_file = str(motion_path)
    print(f"[INFO] Using motion file: {motion_cmd.motion_file}")

    # Check if motion_file is already set (e.g., via CLI --env.commands.motion.motion-file).
    if motion_cmd.motion_file and Path(motion_cmd.motion_file).exists():
      print(f"[INFO] Using local motion file: {motion_cmd.motion_file}")

  # Enable NaN guard if requested.
  if cfg.enable_nan_guard:
    cfg.env.sim.nan_guard.enabled = True
    print(f"[INFO] NaN guard enabled, output dir: {cfg.env.sim.nan_guard.output_dir}")

  if rank == 0:
    print(f"[INFO] Logging experiment in directory: {log_dir}")

  env = ManagerBasedRlEnv(
    cfg=cfg.env, device=device, render_mode="rgb_array" if cfg.video else None
  )

  if pretrained is not None:
    from pretrained import validate_runtime_contract

    try:
      validate_runtime_contract(env)
    except Exception:
      env.close()
      raise

  # The ONNX runner exports this exact snapshot beside and inside policy.onnx.
  env.platform_deployment = deployment
  log_root_path = log_dir.parent  # Go up from specific run dir to experiment dir.

  resume_path: Path | None = None
  explicit_resume = cfg.resume_checkpoint is not None
  if explicit_resume:
    resume_path = Path(cfg.resume_checkpoint).expanduser().resolve(strict=True)
  elif cfg.agent.resume:
    # Load checkpoint from local filesystem.
    resume_path = get_checkpoint_path(
      log_root_path, cfg.agent.load_run, cfg.agent.load_checkpoint
    )

  # Only record videos on rank 0 to avoid multiple workers writing to the same files.
  if cfg.video and rank == 0:
    env = VideoRecorder(
      env,
      video_folder=Path(log_dir) / "videos" / "train",
      step_trigger=lambda step: step % cfg.video_interval == 0,
      video_length=cfg.video_length,
      disable_logger=True,
    )
    print("[INFO] Recording videos during training.")

  env = RslRlVecEnvWrapper(env, clip_actions=cfg.agent.clip_actions)

  agent_cfg = asdict(cfg.agent)
  env_cfg = asdict(cfg.env)

  runner_cls = load_runner_cls(task_id)
  if runner_cls is None:
    runner_cls = MjlabOnPolicyRunner

  runner_kwargs = {}
  runner = runner_cls(env, agent_cfg, str(log_dir), device, **runner_kwargs)

  runner.add_git_repo_to_log(__file__)
  if pretrained is not None:
    from pretrained import initialize_runner

    initialization = initialize_runner(runner, pretrained)
    env.unwrapped.platform_initialization = initialization
    if rank == 0:
      (log_dir / "initialization.json").write_text(
        json.dumps(initialization, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
      )
    print(f"[INFO] Warm-start actor from source {initialization['source_id']}; fresh critic/optimizer, iteration=0")
  if resume_path is not None:
    print(f"[INFO]: Loading model checkpoint from: {resume_path}")
    runner.load(str(resume_path), map_location=device)
    # Promotion copies the origin manifest, never the original actor weights.
    origin_path = resume_path.parent / "initialization.json"
    if origin_path.is_file():
      if origin_path.stat().st_size > 64 * 1024:
        raise ValueError("Initialization provenance exceeds 64 KiB")
      origin = json.loads(origin_path.read_text(encoding="utf-8"))
      if not isinstance(origin, dict) or origin.get("mode") != "pretrained-warm-start":
        raise ValueError("Invalid initialization provenance")
      env.unwrapped.platform_initialization = origin
      if rank == 0:
        (log_dir / "initialization.json").write_text(json.dumps(origin, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    if explicit_resume:
      # RSL-RL stores the last completed zero-based iteration and otherwise
      # repeats it after load. Explicit tuning promotion uses an absolute target.
      runner.current_learning_iteration += 1

  # Only write config files from rank 0 to avoid race conditions.
  if rank == 0:
    dump_yaml(log_dir / "params" / "env.yaml", env_cfg)
    dump_yaml(log_dir / "params" / "agent.yaml", agent_cfg)
    if reward_config is not None:
      reward_snapshot = log_dir / "params" / "reward_config.json"
      reward_snapshot.parent.mkdir(parents=True, exist_ok=True)
      reward_snapshot.write_text(
        json.dumps(reward_config, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
      )

  iterations = cfg.agent.max_iterations
  if explicit_resume:
    iterations = max(0, cfg.agent.max_iterations - runner.current_learning_iteration)
  print(
    f"[INFO] Learning target: current={runner.current_learning_iteration}, "
    f"additional={iterations}, target={cfg.agent.max_iterations}",
    flush=True,
  )
  runner.learn(num_learning_iterations=iterations, init_at_random_ep_len=True)

  env.close()


def launch_training(task_id: str, args: TrainConfig | None = None):
  args = args or TrainConfig.from_task(task_id)

  # Create log directory once before launching workers.
  if args.output_dir:
    log_dir = Path(args.output_dir).expanduser().resolve()
  else:
    log_root_path = (Path("logs") / "rsl_rl" / args.agent.experiment_name).resolve()
    log_dir_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    if args.agent.run_name:
      log_dir_name += f"_{args.agent.run_name}"
    log_dir = log_root_path / log_dir_name
  log_dir.mkdir(parents=True, exist_ok=True)

  # Select GPUs based on CUDA_VISIBLE_DEVICES and user specification.
  selected_gpus, num_gpus = select_gpus(args.gpu_ids)

  # Set environment variables for all modes.
  if selected_gpus is None:
    os.environ["CUDA_VISIBLE_DEVICES"] = ""
  else:
    os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, selected_gpus))
  os.environ["MUJOCO_GL"] = "egl"

  if num_gpus <= 1:
    # CPU or single GPU: run directly without torchrunx.
    run_train(task_id, args, log_dir)
  else:
    # Multi-GPU: use torchrunx.
    import torchrunx

    # torchrunx redirects stdout to logging.
    logging.basicConfig(level=logging.INFO)

    # Configure torchrunx logging directory.
    # Priority: 1) existing env var, 2) user flag, 3) default to {log_dir}/torchrunx.
    if "TORCHRUNX_LOG_DIR" not in os.environ:
      if args.torchrunx_log_dir is not None:
        # User specified a value via flag (could be "" to disable).
        os.environ["TORCHRUNX_LOG_DIR"] = args.torchrunx_log_dir
      else:
        # Default: put logs in training directory.
        os.environ["TORCHRUNX_LOG_DIR"] = str(log_dir / "torchrunx")

    print(f"[INFO] Launching training with {num_gpus} GPUs", flush=True)
    torchrunx.Launcher(
      hostnames=["localhost"],
      workers_per_host=num_gpus,
      backend=None,  # Let rsl_rl handle process group initialization.
      copy_env_vars=torchrunx.DEFAULT_ENV_VARS_FOR_COPY + ("MUJOCO*",),
    ).run(run_train, task_id, args, log_dir)


def main():
  # Parse first argument to choose the task.
  # Import tasks to populate the registry.
  import mjlab.tasks  # noqa: F401
  import src.tasks  # noqa: F401

  all_tasks = list_tasks()
  chosen_task, remaining_args = tyro.cli(
    tyro.extras.literal_type_from_choices(all_tasks),
    add_help=False,
    return_unknown_args=True,
    config=mjlab.TYRO_FLAGS,
  )

  args = tyro.cli(
    TrainConfig,
    args=remaining_args,
    default=TrainConfig.from_task(chosen_task),
    prog=sys.argv[0] + f" {chosen_task}",
    config=mjlab.TYRO_FLAGS,
  )
  del remaining_args

  launch_training(task_id=chosen_task, args=args)


if __name__ == "__main__":
  main()