# Copyright (c) 2025 Deep Robotics
# SPDX-License-Identifier: BSD 3-Clause

"""checkpoint 解析与策略加载（IsaacLab + RSL-RL）。"""

from __future__ import annotations

import os
from typing import Any

import gymnasium as gym

from isaaclab.envs import (
    DirectMARLEnv,
    DirectMARLEnvCfg,
    DirectRLEnvCfg,
    ManagerBasedRLEnvCfg,
    multi_agent_to_single_agent,
)
from isaaclab_tasks.utils import get_checkpoint_path
from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper


def resolve_checkpoint(
    log_root_path: str,
    load_run: str | None = None,
    checkpoint: str | None = None,
) -> str:
    """解析出 checkpoint 的绝对路径。

    允许 ``load_run`` / ``checkpoint`` 传绝对路径或带目录的相对路径，从而跨实验目录
    取权重（例如从 Stage 1 的目录里取权重来评估），否则回退到 IsaacLab 的按名查找。
    """
    if checkpoint is not None:
        checkpoint_path = os.path.expanduser(checkpoint)
        if os.path.isabs(checkpoint_path) or os.path.dirname(checkpoint_path):
            checkpoint_path = os.path.abspath(checkpoint_path)
            if not os.path.isfile(checkpoint_path):
                raise FileNotFoundError(f"找不到 checkpoint：{checkpoint_path}")
            return checkpoint_path

    if load_run is not None:
        run_path = os.path.expanduser(load_run)
        if os.path.isabs(run_path) or os.path.dirname(run_path):
            run_path = os.path.abspath(run_path)
            if not os.path.isdir(run_path):
                raise FileNotFoundError(f"找不到运行目录：{run_path}")
            return get_checkpoint_path(os.path.dirname(run_path), os.path.basename(run_path), checkpoint)

    return get_checkpoint_path(log_root_path, load_run, checkpoint)


def make_env(
    task: str,
    env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg,
    clip_actions: float | None = None,
):
    """按已配置好的 ``env_cfg`` 建环境，并包成 RSL-RL 的向量环境。"""
    env = gym.make(task, cfg=env_cfg)
    if isinstance(env.unwrapped, DirectMARLEnv):
        env = multi_agent_to_single_agent(env)
    return RslRlVecEnvWrapper(env, clip_actions=clip_actions)


def load_policy(env, agent_cfg: Any, resume_path: str, device: str | None = None):
    """加载 checkpoint 并返回推理策略。"""
    from rsl_rl.runners import OnPolicyRunner

    runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=device or agent_cfg.device)
    runner.load(resume_path)
    return runner.get_inference_policy(device=device or env.unwrapped.device)