import dataclasses
import enum
import logging
import os
import pathlib
import socket

import numpy as np
import tyro

from openpi.policies import policy as _policy
from openpi.policies import policy_config as _policy_config
from openpi.serving import websocket_policy_server
from openpi.training import config as _config


def _interpolate_trajectory(traj: np.ndarray, target_len: int) -> np.ndarray:
    """将动作轨迹线性插值到目标帧数（逐维度 np.interp）"""
    n = len(traj)
    if n == target_len:
        return traj
    x_old = np.linspace(0.0, 1.0, n)
    x_new = np.linspace(0.0, 1.0, target_len)
    result = np.zeros((target_len, traj.shape[1]), dtype=np.float32)
    for i in range(traj.shape[1]):
        result[:, i] = np.interp(x_new, x_old, traj[:, i])
    return result


class SmoothedPolicy:
    """对模型输出动作做平滑后处理：
    - Trick 2: 首帧插入当前机器人状态，消除首帧跳变
    - Trick 3: 线性插值到 target_actions 帧，使轨迹更平滑连续
    """

    def __init__(self, policy, target_actions: int = 32, prepend_current_state: bool = True):
        self._policy = policy
        self._target_actions = target_actions
        self._prepend_current_state = prepend_current_state

    @property
    def metadata(self):
        return self._policy.metadata

    def infer(self, obs: dict) -> dict:
        result = self._policy.infer(obs)
        actions = result.get("actions")
        if actions is None:
            return result

        actions = np.array(actions, dtype=np.float32)  # (N, action_dim)

        # Trick 2: 首帧插入当前状态
        if self._prepend_current_state and "state" in obs:
            current_state = np.array(obs["state"], dtype=np.float32)
            if current_state.ndim == 1 and current_state.shape[0] == actions.shape[1]:
                actions = np.concatenate([current_state[None, :], actions], axis=0)

        # Trick 3: 插值到目标帧数
        if len(actions) != self._target_actions:
            actions = _interpolate_trajectory(actions, self._target_actions)

        result["actions"] = actions
        return result


class RTCPrefixPolicy:
    """方案 C 近似版：用上一 chunk 末尾动作状态混合当前观察状态，
    使模型感知到前一 chunk 执行到哪里，从而生成轨迹连续的新 chunk。

    obs 中若含 _rtc_mode=True，则提取 _prefix_actions 并做状态混合，
    再剥除所有 _rtc_* 键后调用下层 policy.infer()。
    """

    def __init__(self, policy, blend: float = 0.3):
        """
        blend: 机器人真实状态的权重（0=完全用前缀末动作, 1=完全用机器人状态）
               推荐 0.2~0.4：偏低减少 chunk 边界抖动，偏高减少状态偏差累积
        """
        self._policy = policy
        self._blend = blend

    @property
    def metadata(self):
        return self._policy.metadata

    def infer(self, obs: dict) -> dict:
        # 剥除 RTC 专用键，避免传入底层 policy
        rtc_mode = obs.pop("_rtc_mode", False)
        prefix_actions = obs.pop("_prefix_actions", None)
        obs.pop("_inference_delay", None)

        if rtc_mode and prefix_actions is not None and len(prefix_actions) > 0:
            last_action = np.array(prefix_actions[-1], dtype=np.float32).flatten()
            current_state = np.array(obs.get("state", []), dtype=np.float32).flatten()

            if last_action.shape == current_state.shape and len(last_action) >= 6:
                blend = self._blend
                blended = current_state.copy()
                # 混合 xyz（位置）和 yaw（偏航），roll/pitch 保留机器人真实值
                blended[0:3] = (1.0 - blend) * last_action[0:3] + blend * current_state[0:3]
                blended[5]   = (1.0 - blend) * last_action[5]   + blend * current_state[5]
                # gripper 用前缀末动作（更能反映当前意图状态）
                if len(last_action) >= 7:
                    blended[6] = last_action[6]
                obs["state"] = blended
                logging.debug(
                    "RTC prefix blend=%.2f | xyz_delta=[%.4f,%.4f,%.4f]",
                    blend,
                    last_action[0] - current_state[0],
                    last_action[1] - current_state[1],
                    last_action[2] - current_state[2],
                )

        return self._policy.infer(obs)


class EnvMode(enum.Enum):
    """Supported environments."""

    ALOHA = "aloha"
    ALOHA_SIM = "aloha_sim"
    DROID = "droid"
    LIBERO = "libero"


@dataclasses.dataclass
class Checkpoint:
    """Load a policy from a trained checkpoint."""

    # Training config name (e.g., "pi0_aloha_sim").
    config: str
    # Checkpoint directory (e.g., "checkpoints/pi0_aloha_sim/exp/10000").
    dir: str
    # 可选：指定 assets 目录（含 norm_stats.json）。
    # 若 checkpoint dir 下没有 assets/，脚本会自动软链接到此路径。
    # 例: checkpoints/pi05_franka_finetune/29999/assets
    assets_dir: str = ""


@dataclasses.dataclass
class Default:
    """Use the default policy for the given environment."""


@dataclasses.dataclass
class Args:
    """Arguments for the serve_policy script."""

    # Environment to serve the policy for. This is only used when serving default policies.
    env: EnvMode = EnvMode.ALOHA_SIM

    # If provided, will be used in case the "prompt" key is not present in the data, or if the model doesn't have a default
    # prompt.
    default_prompt: str | None = None

    # Port to serve the policy on.
    port: int = 8000
    # Record the policy's behavior for debugging.
    record: bool = False

    # Specifies how to load the policy. If not provided, the default policy for the environment will be used.
    policy: Checkpoint | Default = dataclasses.field(default_factory=Default)

    # ── 平滑参数 ──────────────────────────────────────────────────────────────
    # 是否启用动作平滑后处理
    smooth: bool = True
    # 插值后输出的动作帧数（模型原始输出经插值到此长度再发给客户端）
    target_actions: int = 32
    # 是否在轨迹首帧插入当前机器人状态（消除首帧跳变）
    prepend_current_state: bool = True

    # ── RTC 前缀条件化参数 ────────────────────────────────────────────────────
    # 是否启用 RTC 近似前缀条件化（需客户端推理节点配合发送 _prefix_actions）
    rtc_prefix: bool = True
    # 机器人真实状态的混合权重（0=完全用前缀末动作, 1=完全用机器人状态）
    rtc_blend: float = 0.3


# Default checkpoints that should be used for each environment.
DEFAULT_CHECKPOINT: dict[EnvMode, Checkpoint] = {
    EnvMode.ALOHA: Checkpoint(
        config="pi05_aloha",
        dir="gs://openpi-assets/checkpoints/pi05_base",
    ),
    EnvMode.ALOHA_SIM: Checkpoint(
        config="pi0_aloha_sim",
        dir="gs://openpi-assets/checkpoints/pi0_aloha_sim",
    ),
    EnvMode.DROID: Checkpoint(
        config="pi05_droid",
        dir="gs://openpi-assets/checkpoints/pi05_droid",
    ),
    EnvMode.LIBERO: Checkpoint(
        config="pi05_libero",
        dir="gs://openpi-assets/checkpoints/pi05_libero",
    ),
}


def create_default_policy(env: EnvMode, *, default_prompt: str | None = None) -> _policy.Policy:
    """Create a default policy for the given environment."""
    if checkpoint := DEFAULT_CHECKPOINT.get(env):
        return _policy_config.create_trained_policy(
            _config.get_config(checkpoint.config), checkpoint.dir, default_prompt=default_prompt
        )
    raise ValueError(f"Unsupported environment mode: {env}")


def _ensure_assets(ckpt_dir: str, assets_dir: str) -> None:
    """确保 checkpoint 目录下有 assets/，否则自动软链接。

    查找优先级：
    1. 显式传入的 assets_dir
    2. 同级目录下任意含 assets/ 的 checkpoint（按名称排序取第一个）
    """
    ckpt_path = pathlib.Path(ckpt_dir).resolve()
    assets_link = ckpt_path / "assets"

    # 清理 broken symlink
    if assets_link.is_symlink() and not assets_link.exists():
        logging.warning("检测到失效的 assets 软链接，已删除: %s", assets_link)
        assets_link.unlink()

    if assets_link.exists():
        return  # 已存在，无需处理

    # 候选来源列表
    candidates: list[pathlib.Path] = []
    if assets_dir:
        candidates.append(pathlib.Path(assets_dir).resolve())
    # 自动搜索同级 checkpoint 目录
    for sibling in sorted(ckpt_path.parent.iterdir()):
        if sibling.is_dir() and sibling != ckpt_path and (sibling / "assets").is_dir():
            candidates.append(sibling / "assets")

    if not candidates:
        logging.warning("未找到任何 assets 目录，将由 create_trained_policy 报错")
        return

    src = candidates[0]
    try:
        os.symlink(src, assets_link)
        logging.info("自动创建 assets 软链接: %s -> %s", assets_link, src)
    except OSError as e:
        logging.error(
            "创建 assets 软链接失败: %s\n"
            "请手动运行: ln -s %s %s",
            e, src, assets_link,
        )
        raise


def create_policy(args: Args) -> _policy.Policy:
    """Create a policy from the given arguments."""
    match args.policy:
        case Checkpoint():
            _ensure_assets(args.policy.dir, args.policy.assets_dir)
            return _policy_config.create_trained_policy(
                _config.get_config(args.policy.config), args.policy.dir, default_prompt=args.default_prompt
            )
        case Default():
            return create_default_policy(args.env, default_prompt=args.default_prompt)


def main(args: Args) -> None:
    policy = create_policy(args)
    policy_metadata = policy.metadata

    # Record the policy's behavior.
    if args.record:
        policy = _policy.PolicyRecorder(policy, "policy_records")

    # 平滑后处理：首帧插入当前状态 + 插值到 target_actions 帧
    if args.smooth:
        policy = SmoothedPolicy(
            policy,
            target_actions=args.target_actions,
            prepend_current_state=args.prepend_current_state,
        )
        logging.info(
            "SmoothedPolicy enabled: target_actions=%d, prepend_current_state=%s",
            args.target_actions,
            args.prepend_current_state,
        )

    # RTC 近似前缀条件化：用上一 chunk 末尾动作混合观察状态，减少 chunk 边界抖动
    # 注意：RTCPrefixPolicy 必须在 SmoothedPolicy 外层包装，确保 RTC 键在到达底层前被剥除
    if args.rtc_prefix:
        policy = RTCPrefixPolicy(policy, blend=args.rtc_blend)
        logging.info("RTCPrefixPolicy enabled: blend=%.2f", args.rtc_blend)

    hostname = socket.gethostname()
    try:
        local_ip = socket.gethostbyname(hostname)
    except socket.gaierror:
        local_ip = "127.0.0.1"
        logging.warning("Could not resolve hostname %s; using %s for logging only", hostname, local_ip)
    logging.info("Creating server (host: %s, ip: %s)", hostname, local_ip)

    server = websocket_policy_server.WebsocketPolicyServer(
        policy=policy,
        host="0.0.0.0",
        port=args.port,
        metadata=policy_metadata,
    )
    server.serve_forever()


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, force=True)
    main(tyro.cli(Args))