#!/usr/bin/env python3
"""浏览器 MuJoCo 平台的本地 mjlab 训练桥接服务（仅绑定 loopback）。"""

from __future__ import annotations

import argparse
import hmac
import json
import os
import re
import secrets
import shutil
import signal
import subprocess
import sys
import threading
import uuid
from collections import deque
from contextlib import suppress
from dataclasses import dataclass, field
from datetime import UTC, datetime
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any
from urllib.parse import unquote, urlsplit

VERSION = "0.3.0"
# 浏览器当前 ONNX 运行时只实现 Go2 的 47→12 部署契约；其他任务须由服务启动参数显式放行。
DEFAULT_TASKS = ("Unitree-Go2-Flat",)
ACTIVE_STATES = {"queued", "running"}
MAX_JOBS = 20
ANSI_ESCAPE = re.compile(r"\x1b\[[0-?]*[ -/]*[@-~]")
ITERATION_PATTERNS = (
    re.compile(r"(?:learning\s+)?iteration\D+(\d+)\s*/\s*(\d+)", re.I),
    re.compile(r"(?:learning\s+)?iteration\D+(\d+)", re.I),
)
RUN_NAME = re.compile(r"^[A-Za-z0-9_.-]{1,64}$")
LOCAL_ORIGIN = re.compile(r"^https?://(?:localhost|127\.0\.0\.1)(?::\d+)?$")
LOCAL_HOST = re.compile(r"^(?:localhost|127\.0\.0\.1)(?::\d+)?$")


def now_iso() -> str:
    return datetime.now(UTC).isoformat()


def termination_signal_handler(_signum: int, _frame: Any) -> None:
    """将 SIGTERM 转成受控退出，使 main 的 finally 能清理训练子进程。"""
    raise KeyboardInterrupt


class ApiError(Exception):
    def __init__(self, status: int, message: str):
        super().__init__(message)
        self.status = status


@dataclass
class TrainingConfig:
    task_id: str
    num_envs: int
    max_iterations: int
    seed: int
    run_name: str
    device: str
    gpu_ids: list[int]
    wandb_mode: str


@dataclass
class TrainingJob:
    id: str
    config: TrainingConfig
    state: str = "queued"
    created_at: str = field(default_factory=now_iso)
    started_at: str | None = None
    ended_at: str | None = None
    iteration: int = 0
    message: str = "等待本地训练进程启动"
    logs: deque[str] = field(default_factory=lambda: deque(maxlen=200))
    artifact: Path | None = None
    process: subprocess.Popen[str] | None = None
    cancel_requested: bool = False

    def public(self) -> dict[str, Any]:
        progress = min(1.0, max(0.0, self.iteration / self.config.max_iterations))
        if self.state == "succeeded":
            progress = 1.0
        return {
            "id": self.id,
            "state": self.state,
            "taskId": self.config.task_id,
            "createdAt": self.created_at,
            "startedAt": self.started_at,
            "endedAt": self.ended_at,
            "iteration": self.iteration,
            "maxIterations": self.config.max_iterations,
            "progress": progress,
            "message": self.message,
            "logs": list(self.logs),
            "artifactReady": self.artifact is not None and self.artifact.is_file(),
            "artifactName": self.artifact.name if self.artifact else None,
        }


class TrainingManager:
    def __init__(
        self,
        trainer_root: Path,
        python: str,
        tasks: tuple[str, ...],
        check_environment: bool = True,
    ):
        self.trainer_root = trainer_root.expanduser().resolve()
        self.python = str(Path(python).expanduser()) if os.sep in python else python
        self.tasks = tasks
        self.jobs: dict[str, TrainingJob] = {}
        self.lock = threading.RLock()
        self.check_environment = check_environment
        self._environment_error: str | None | bool = False

    def readiness_error(self) -> str | None:
        if not self.trainer_root.is_dir():
            return f"训练工程目录不存在：{self.trainer_root}"
        if not (self.trainer_root / "scripts" / "train.py").is_file():
            return f"训练入口不存在：{self.trainer_root / 'scripts/train.py'}"
        executable = Path(self.python)
        if not executable.is_file() and shutil.which(self.python) is None:
            return f"Python 解释器不存在：{self.python}"
        if self.check_environment and self._environment_error is False:
            probe = (
                "import importlib.util,sys; "
                "missing=[m for m in ('mjlab','torch','tyro','wandb') "
                "if importlib.util.find_spec(m) is None]; "
                "print(','.join(missing)); sys.exit(bool(missing))"
            )
            try:
                result = subprocess.run(
                    [self.python, "-c", probe],
                    cwd=self.trainer_root,
                    capture_output=True,
                    text=True,
                    timeout=15,
                    check=False,
                )
                missing = result.stdout.strip()
                self._environment_error = (
                    f"训练 Python 缺少依赖：{missing}" if result.returncode else None
                )
            except (OSError, subprocess.TimeoutExpired) as error:
                self._environment_error = f"无法检查训练 Python 环境：{error}"
        return self._environment_error if isinstance(self._environment_error, str) else None

    def active_job_id(self) -> str | None:
        with self.lock:
            return next((job.id for job in self.jobs.values() if job.state in ACTIVE_STATES), None)

    def health(self) -> dict[str, Any]:
        error = self.readiness_error()
        return {
            "version": VERSION,
            "ready": error is None,
            "trainerRoot": str(self.trainer_root),
            "python": self.python,
            "tasks": list(self.tasks),
            "activeJobId": self.active_job_id(),
            "error": error,
        }

    def parse_config(self, payload: Any) -> TrainingConfig:
        if not isinstance(payload, dict):
            raise ApiError(HTTPStatus.BAD_REQUEST, "请求体必须是 JSON 对象")
        task_id = payload.get("taskId")
        if task_id not in self.tasks:
            raise ApiError(HTTPStatus.BAD_REQUEST, f"不允许的训练任务：{task_id}")

        def integer(name: str, minimum: int, maximum: int) -> int:
            value = payload.get(name)
            if (
                isinstance(value, bool)
                or not isinstance(value, int)
                or not minimum <= value <= maximum
            ):
                raise ApiError(HTTPStatus.BAD_REQUEST, f"{name} 必须在 {minimum}–{maximum} 之间")
            return value

        run_name = payload.get("runName", "web")
        if not isinstance(run_name, str) or not RUN_NAME.fullmatch(run_name):
            raise ApiError(
                HTTPStatus.BAD_REQUEST,
                "runName 只能包含字母、数字、点、下划线和连字符，最长 64 字符",
            )
        device = payload.get("device")
        if device not in ("cpu", "gpu"):
            raise ApiError(HTTPStatus.BAD_REQUEST, "device 必须是 cpu 或 gpu")
        raw_gpu_ids = payload.get("gpuIds", [])
        if not isinstance(raw_gpu_ids, list) or any(
            isinstance(value, bool) or not isinstance(value, int) or value < 0 or value > 255
            for value in raw_gpu_ids
        ):
            raise ApiError(HTTPStatus.BAD_REQUEST, "gpuIds 必须是非负整数数组")
        if device == "gpu" and not raw_gpu_ids:
            raise ApiError(HTTPStatus.BAD_REQUEST, "GPU 训练至少需要一个 GPU 编号")
        wandb_mode = payload.get("wandbMode", "offline")
        if wandb_mode not in ("offline", "disabled", "online"):
            raise ApiError(HTTPStatus.BAD_REQUEST, "wandbMode 必须是 offline、disabled 或 online")
        return TrainingConfig(
            task_id=task_id,
            num_envs=integer("numEnvs", 1, 16384),
            max_iterations=integer("maxIterations", 1, 1_000_000),
            seed=integer("seed", 0, 2_147_483_647),
            run_name=run_name,
            device=device,
            gpu_ids=raw_gpu_ids,
            wandb_mode=wandb_mode,
        )

    def start(self, payload: Any) -> dict[str, Any]:
        error = self.readiness_error()
        if error:
            raise ApiError(HTTPStatus.SERVICE_UNAVAILABLE, error)
        config = self.parse_config(payload)
        with self.lock:
            if self.active_job_id():
                raise ApiError(HTTPStatus.CONFLICT, "已有训练任务正在运行，请等待完成或先停止任务")
            while len(self.jobs) >= MAX_JOBS:
                completed = next(
                    (job_id for job_id, job in self.jobs.items() if job.state not in ACTIVE_STATES),
                    None,
                )
                if completed is None:
                    raise ApiError(HTTPStatus.CONFLICT, "训练任务历史已满，请稍后重试")
                del self.jobs[completed]
            job = TrainingJob(id=uuid.uuid4().hex, config=config)
            self.jobs[job.id] = job
            threading.Thread(
                target=self._run, args=(job,), name=f"training-{job.id[:8]}", daemon=True
            ).start()
            return job.public()

    def get(self, job_id: str) -> dict[str, Any]:
        with self.lock:
            job = self.jobs.get(job_id)
            if not job:
                raise ApiError(HTTPStatus.NOT_FOUND, "训练任务不存在或服务已重启")
            return job.public()

    def artifact(self, job_id: str) -> Path:
        with self.lock:
            job = self.jobs.get(job_id)
            if not job:
                raise ApiError(HTTPStatus.NOT_FOUND, "训练任务不存在或服务已重启")
            if not job.artifact or not job.artifact.is_file():
                raise ApiError(HTTPStatus.NOT_FOUND, "该训练任务尚未生成 policy.onnx")
            return job.artifact

    def cancel(self, job_id: str) -> dict[str, Any]:
        with self.lock:
            job = self.jobs.get(job_id)
            if not job:
                raise ApiError(HTTPStatus.NOT_FOUND, "训练任务不存在或服务已重启")
            if job.state not in ACTIVE_STATES:
                return job.public()
            job.cancel_requested = True
            job.message = "正在停止训练进程"
            process = job.process
        if process and process.poll() is None:
            with suppress(ProcessLookupError):
                os.killpg(process.pid, signal.SIGTERM)
            threading.Thread(target=self._kill_later, args=(process,), daemon=True).start()
        return self.get(job_id)

    @staticmethod
    def _kill_later(process: subprocess.Popen[str]) -> None:
        try:
            process.wait(timeout=5)
        except subprocess.TimeoutExpired:
            with suppress(ProcessLookupError):
                os.killpg(process.pid, signal.SIGKILL)

    def shutdown(self) -> None:
        """同步停止活动训练，避免服务退出后遗留子进程。"""
        active = self.active_job_id()
        if not active:
            return
        self.cancel(active)
        with self.lock:
            process = self.jobs[active].process
        if process and process.poll() is None:
            try:
                process.wait(timeout=6)
            except subprocess.TimeoutExpired:
                with suppress(ProcessLookupError):
                    os.killpg(process.pid, signal.SIGKILL)
                process.wait(timeout=2)

    def command_for(self, config: TrainingConfig) -> list[str]:
        command = [
            self.python,
            "-u",
            "scripts/train.py",
            config.task_id,
            f"--env.scene.num-envs={config.num_envs}",
            f"--agent.max-iterations={config.max_iterations}",
            f"--agent.seed={config.seed}",
            f"--agent.run-name={config.run_name}",
        ]
        if config.device == "cpu":
            command.extend(("--gpu-ids", "None"))
        else:
            # mjlab.TYRO_FLAGS 对 Union[list[int], Literal["all"], None] 使用 JSON 风格
            # list token；传成多个独立参数会被解析为错误的 Union 分支。
            command.extend(("--gpu-ids", json.dumps(config.gpu_ids, separators=(",", ":"))))
        return command

    def _update_from_log(self, job: TrainingJob, raw_line: str) -> None:
        line = ANSI_ESCAPE.sub("", raw_line).strip()
        if not line:
            return
        with self.lock:
            job.logs.append(line[-4000:])
            for pattern in ITERATION_PATTERNS:
                match = pattern.search(line)
                if match:
                    job.iteration = min(
                        job.config.max_iterations, max(job.iteration, int(match.group(1)))
                    )
                    break
            job.message = line[-240:]

    def _artifact_snapshot(self) -> dict[Path, int]:
        root = self.trainer_root / "logs" / "rsl_rl"
        if not root.is_dir():
            return {}
        return {
            path: path.stat().st_mtime_ns for path in root.glob("**/policy.onnx") if path.is_file()
        }

    def _find_artifact(self, before: dict[Path, int]) -> Path | None:
        root = self.trainer_root / "logs" / "rsl_rl"
        if not root.is_dir():
            return None
        changed = [
            path
            for path in root.glob("**/policy.onnx")
            if path.is_file() and before.get(path) != path.stat().st_mtime_ns
        ]
        return max(changed, key=lambda path: path.stat().st_mtime_ns) if changed else None

    def _run(self, job: TrainingJob) -> None:
        before = self._artifact_snapshot()
        command = self.command_for(job.config)
        environment = os.environ.copy()
        # 默认离线记录，保留本地 W&B 指标但不要求 API Key；只有前端明确选择
        # online 时才允许 wandb 发起登录/联网。
        environment["WANDB_MODE"] = job.config.wandb_mode
        environment["WANDB_SILENT"] = "true"
        try:
            # Popen 与 process 登记必须和取消检查处于同一个临界区：cancel() 要么在
            # 创建前标记取消，要么在创建后取得进程并终止，不能落入二者之间。
            with self.lock:
                if job.cancel_requested:
                    job.state, job.ended_at, job.message = "cancelled", now_iso(), "训练已取消"
                    return
                process = subprocess.Popen(
                    command,
                    cwd=self.trainer_root,
                    env=environment,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    text=True,
                    encoding="utf-8",
                    errors="replace",
                    bufsize=1,
                    start_new_session=True,
                )
                job.process = process
                job.state, job.started_at, job.message = (
                    "running",
                    now_iso(),
                    "本地训练进程已启动",
                )
            assert process.stdout is not None
            try:
                for line in process.stdout:
                    self._update_from_log(job, line)
            finally:
                process.stdout.close()
            return_code = process.wait()
            artifact = self._find_artifact(before)
            with self.lock:
                job.process = None
                job.ended_at = now_iso()
                if job.cancel_requested:
                    job.state, job.message = "cancelled", "训练已由用户取消"
                elif return_code != 0:
                    job.state, job.message = "failed", f"训练进程退出，返回码 {return_code}"
                elif artifact is None:
                    job.state, job.message = "failed", "训练结束，但没有找到本次生成的 policy.onnx"
                else:
                    job.state, job.artifact = "succeeded", artifact
                    job.iteration = job.config.max_iterations
                    job.message = f"训练完成：{artifact.relative_to(self.trainer_root)}"
        except Exception as error:  # 服务必须保留错误供前端诊断。
            with self.lock:
                job.process = None
                job.ended_at = now_iso()
                job.state = "cancelled" if job.cancel_requested else "failed"
                job.message = f"启动训练失败：{error}"
                job.logs.append(job.message)


class TrainingRequestHandler(BaseHTTPRequestHandler):
    manager: TrainingManager
    allowed_origins: tuple[str, ...] = ()
    access_token = ""
    server_version = "MuJoCoLocalTraining/0.3"

    def log_message(self, format: str, *args: Any) -> None:
        sys.stderr.write(f"[{self.log_date_time_string()}] {format % args}\n")

    def _origin_allowed(self) -> bool:
        origin = self.headers.get("Origin")
        return (
            origin is None or bool(LOCAL_ORIGIN.fullmatch(origin)) or origin in self.allowed_origins
        )

    def _host_allowed(self) -> bool:
        host = self.headers.get("Host", "")
        return bool(LOCAL_HOST.fullmatch(host))

    def _authorized(self) -> bool:
        authorization = self.headers.get("Authorization", "")
        prefix = "Bearer "
        return authorization.startswith(prefix) and hmac.compare_digest(
            authorization[len(prefix) :], self.access_token
        )

    def _cors(self) -> None:
        origin = self.headers.get("Origin")
        if origin and self._origin_allowed():
            self.send_header("Access-Control-Allow-Origin", origin)
            self.send_header("Vary", "Origin")

    def _json(self, status: int, payload: Any) -> None:
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self._cors()
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(body)

    def _error(self, error: Exception) -> None:
        if isinstance(error, ApiError):
            self._json(error.status, {"error": str(error)})
        else:
            self._json(
                HTTPStatus.INTERNAL_SERVER_ERROR, {"error": f"本地训练服务内部错误：{error}"}
            )

    def _ensure_origin(self) -> None:
        if not self._host_allowed():
            raise ApiError(HTTPStatus.FORBIDDEN, "不允许的 Host")
        if not self._origin_allowed():
            raise ApiError(HTTPStatus.FORBIDDEN, "不允许的浏览器来源")

    def _ensure_request(self) -> None:
        self._ensure_origin()
        if not self._authorized():
            raise ApiError(HTTPStatus.UNAUTHORIZED, "训练服务访问令牌无效")

    def _payload(self) -> Any:
        try:
            length = int(self.headers.get("Content-Length", "0"))
        except ValueError as error:
            raise ApiError(HTTPStatus.BAD_REQUEST, "Content-Length 无效") from error
        if length <= 0 or length > 32 * 1024:
            raise ApiError(
                HTTPStatus.REQUEST_ENTITY_TOO_LARGE, "训练请求体不能为空且不能超过 32 KiB"
            )
        try:
            return json.loads(self.rfile.read(length))
        except (UnicodeDecodeError, json.JSONDecodeError) as error:
            raise ApiError(HTTPStatus.BAD_REQUEST, "训练请求不是有效 JSON") from error

    @staticmethod
    def _route(path: str) -> tuple[str | None, bool]:
        match = re.fullmatch(r"/api/training/jobs/([0-9a-f]{32})(/artifacts/policy\.onnx)?", path)
        return (unquote(match.group(1)), bool(match.group(2))) if match else (None, False)

    def do_OPTIONS(self) -> None:
        try:
            self._ensure_origin()
            self.send_response(HTTPStatus.NO_CONTENT)
            self._cors()
            self.send_header("Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS")
            self.send_header("Access-Control-Allow-Headers", "Authorization, Content-Type")
            self.send_header("Access-Control-Max-Age", "600")
            self.end_headers()
        except Exception as error:
            self._error(error)

    def do_GET(self) -> None:
        try:
            self._ensure_request()
            path = urlsplit(self.path).path
            if path == "/api/training/health":
                self._json(HTTPStatus.OK, self.manager.health())
                return
            job_id, artifact = self._route(path)
            if not job_id:
                raise ApiError(HTTPStatus.NOT_FOUND, "接口不存在")
            if artifact:
                file_path = self.manager.artifact(job_id)
                size = file_path.stat().st_size
                self.send_response(HTTPStatus.OK)
                self._cors()
                self.send_header("Content-Type", "application/octet-stream")
                self.send_header("Content-Disposition", 'attachment; filename="policy.onnx"')
                self.send_header("Content-Length", str(size))
                self.send_header("Cache-Control", "no-store")
                self.end_headers()
                with file_path.open("rb") as source:
                    shutil.copyfileobj(source, self.wfile)
            else:
                self._json(HTTPStatus.OK, self.manager.get(job_id))
        except Exception as error:
            self._error(error)

    def do_POST(self) -> None:
        try:
            self._ensure_request()
            if urlsplit(self.path).path != "/api/training/jobs":
                raise ApiError(HTTPStatus.NOT_FOUND, "接口不存在")
            self._json(HTTPStatus.ACCEPTED, self.manager.start(self._payload()))
        except Exception as error:
            self._error(error)

    def do_DELETE(self) -> None:
        try:
            self._ensure_request()
            job_id, artifact = self._route(urlsplit(self.path).path)
            if not job_id or artifact:
                raise ApiError(HTTPStatus.NOT_FOUND, "接口不存在")
            self._json(HTTPStatus.ACCEPTED, self.manager.cancel(job_id))
        except Exception as error:
            self._error(error)


def default_trainer_root() -> Path:
    configured = os.environ.get("UNITREE_RL_MJLAB_ROOT")
    if configured:
        return Path(configured)
    return Path(__file__).resolve().parent / "rl"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="MuJoCo Web 平台本地强化学习训练服务")
    parser.add_argument(
        "--host",
        default="127.0.0.1",
        choices=("127.0.0.1", "localhost"),
        help="仅允许绑定本机回环地址",
    )
    parser.add_argument("--port", type=int, default=8765)
    parser.add_argument(
        "--trainer-root",
        type=Path,
        default=default_trainer_root(),
        help="训练工程目录；默认使用仓库内置的 Go2 训练器",
    )
    parser.add_argument(
        "--trainer-python", default=sys.executable, help="已安装 mjlab/torch 的 Python 解释器"
    )
    parser.add_argument(
        "--task", action="append", dest="tasks", help="允许前端启动的任务 ID；可重复"
    )
    parser.add_argument(
        "--allow-origin", action="append", default=[], help="额外允许的前端 Origin；可重复"
    )
    parser.add_argument(
        "--token",
        default=os.environ.get("MUJOCO_TRAINING_TOKEN"),
        help="访问令牌；默认随机生成，也可通过 MUJOCO_TRAINING_TOKEN 设置",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    token = args.token or secrets.token_urlsafe(24)
    if len(token) < 16:
        raise SystemExit("训练服务访问令牌至少需要 16 个字符")
    manager = TrainingManager(
        args.trainer_root, args.trainer_python, tuple(args.tasks or DEFAULT_TASKS)
    )
    TrainingRequestHandler.manager = manager
    TrainingRequestHandler.allowed_origins = tuple(args.allow_origin)
    TrainingRequestHandler.access_token = token
    server = ThreadingHTTPServer((args.host, args.port), TrainingRequestHandler)
    print(f"本地训练服务：http://{args.host}:{args.port}")
    print(f"访问令牌：{token}")
    print(f"训练工程：{manager.trainer_root}")
    print(f"Python：{manager.python}")
    if manager.readiness_error():
        print(f"警告：{manager.readiness_error()}", file=sys.stderr)
    signal.signal(signal.SIGTERM, termination_signal_handler)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\n正在停止本地训练服务…")
    finally:
        manager.shutdown()
        server.server_close()


if __name__ == "__main__":
    main()