"""Administrator-registered, content-bound local training sources (no client paths)."""

from __future__ import annotations

import hashlib
import json
import os
import re
import shutil
import stat
import subprocess
import tempfile
import threading
import time
from contextlib import contextmanager
from copy import deepcopy
from pathlib import Path

SOURCE_ID = re.compile(r"^[A-Za-z0-9_-]{1,64}$")
DIGEST = re.compile(r"^[0-9a-f]{64}$")
LIMITS = {
    "checkpoint": 256 * 1024**2,
    "onnx": 64 * 1024**2,
    "env": 2 * 1024**2,
    "agent": 128 * 1024,
}
TASKS = ["Unitree-Go2-Flat", "Unitree-Go2-ObstacleAvoidance"]


class SourceError(ValueError):
    pass


def regular_bytes(path: Path, root: Path, limit: int) -> bytes:
    """Walk with dir_fd/O_NOFOLLOW, so symlink substitution cannot escape the root."""
    if ".." in path.parts or not path.is_absolute():
        raise SourceError("基础策略路径必须是允许根内的绝对路径，不能含 ..")
    try:
        parts = path.relative_to(root).parts
    except ValueError as error:
        raise SourceError("基础策略文件不在管理员配置的允许根内") from error
    if not parts:
        raise SourceError("基础策略必须是常规文件")
    try:
        fd = os.open(root, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
    except OSError as error:
        raise SourceError("基础策略快照目录缺失或无效，拒绝训练") from error
    try:
        for part in parts[:-1]:
            next_fd = os.open(part, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=fd)
            os.close(fd)
            fd = next_fd
        leaf = os.open(parts[-1], os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK, dir_fd=fd)
        with os.fdopen(leaf, "rb") as stream:
            info = os.fstat(stream.fileno())
            if not stat.S_ISREG(info.st_mode) or info.st_size > limit:
                raise SourceError("基础策略文件必须是常规文件且不能超过大小上限")
            data = stream.read(limit + 1)
            if len(data) > limit:
                raise SourceError("基础策略文件超过大小上限")
            return data
    except OSError as error:
        raise SourceError(
            "基础策略文件缺失或含symlink；请提供.pt、policy.onnx和params配置常规文件"
        ) from error
    finally:
        os.close(fd)


class PretrainedSources:
    def __init__(self, config_path: Path | None, store: Path, python: str, trainer_root: Path):
        self.store = store.expanduser().resolve()
        self.store.mkdir(parents=True, exist_ok=True, mode=0o700)
        self.python = python
        self.trainer_root = trainer_root.resolve()
        self.entries: dict[str, dict] = {}
        self._upload_lock = threading.Lock()
        self._catalog_lock = threading.RLock()
        self._restore_uploads()
        if config_path is None:
            return
        try:
            if config_path.stat().st_size > 128 * 1024:
                raise SourceError("基础策略注册配置不能超过128KiB")
            config = json.loads(config_path.read_text())
            if not isinstance(config, dict) or set(config) != {"allowedRoots", "sources"}:
                raise SourceError("注册配置只接受allowedRoots和sources")
            if (
                not isinstance(config["allowedRoots"], list)
                or not isinstance(config["sources"], list)
                or len(config["sources"]) > 32
            ):
                raise SourceError("允许根/sources必须为数组，最多32个注册条目")
            roots = [Path(p).expanduser().resolve(strict=True) for p in config["allowedRoots"]]
            if not roots or any(not p.is_dir() for p in roots):
                raise SourceError("请配置存在的本地允许根目录")
            for entry in config["sources"]:
                if not isinstance(entry, dict) or set(entry) != {
                    "id",
                    "label",
                    "checkpoint",
                    "onnx",
                }:
                    raise SourceError("注册条目必须包含id/label/checkpoint/onnx")
                key = entry["id"]
                if not isinstance(key, str) or not SOURCE_ID.fullmatch(key) or key in self.entries:
                    raise SourceError("注册id无效或重复")
                if not isinstance(entry["label"], str) or not 1 <= len(entry["label"]) <= 100:
                    raise SourceError("基础策略label必须为1–100字符")
                record = {"id": key, "label": entry["label"], "ready": False, "compatibleTasks": []}
                self.entries[key] = {"public": record}
                try:
                    checkpoint, onnx = (
                        Path(entry["checkpoint"]).expanduser(),
                        Path(entry["onnx"]).expanduser(),
                    )
                    if not re.fullmatch(r"[A-Za-z0-9_.-]+\.pt", checkpoint.name):
                        raise SourceError("请选择ONNX对应的.pt训练checkpoint，ONNX不能直接续训")
                    files = {
                        "checkpoint": checkpoint,
                        "onnx": onnx,
                        "env": checkpoint.parent / "params/env.yaml",
                        "agent": checkpoint.parent / "params/agent.yaml",
                    }
                    data = {}
                    for name, path in files.items():
                        root = next((r for r in roots if path.is_relative_to(r)), None)
                        if root is None:
                            raise SourceError("注册文件不在允许根内")
                        data[name] = regular_bytes(path, root, LIMITS[name])
                    names = {
                        "checkpoint": checkpoint.name,
                        "onnx": "policy.onnx",
                        "env": "params/env.yaml",
                        "agent": "params/agent.yaml",
                    }
                    with tempfile.TemporaryDirectory(dir=self.store, prefix="import-") as temporary:
                        directory = Path(temporary)
                        for name, content in data.items():
                            destination = directory / names[name]
                            destination.parent.mkdir(exist_ok=True)
                            destination.write_bytes(content)
                        manifest = self._validate(directory, checkpoint.name, TASKS[0], None)
                        digest = manifest["source_id"]
                        bound = {
                            "sourceId": digest,
                            "registeredId": key,
                            "label": entry["label"],
                            "manifest": manifest,
                        }
                        destination = self.store / digest
                        if not destination.exists():
                            shutil.copytree(directory, destination)
                            for file in destination.rglob("*"):
                                file.chmod(0o555 if file.is_dir() else 0o444)
                            destination.chmod(0o555)
                        self.verify(bound)
                    record.update(
                        id=digest,
                        ready=True,
                        compatibleTasks=TASKS,
                        observationSizes=[47, 81, 97],
                        initialization=bound,
                    )
                except (SourceError, OSError, ValueError, TypeError) as error:
                    record["error"] = f"{error}；请修正管理员注册配置并重启服务，不会退回随机初始化"
        except (OSError, TypeError, ValueError) as error:
            raise SourceError(f"基础策略注册配置无效：{error}") from error

    def _restore_uploads(self):
        # Only committed content directories are catalogued. Partial uploads are never sources.
        for directory in self.store.iterdir():
            if not DIGEST.fullmatch(directory.name) or not (directory / "upload.json").is_file():
                continue
            try:
                manifest = json.loads(
                    regular_bytes(directory / "upload.json", self.store, 64 * 1024)
                )
                label = json.loads(regular_bytes(directory / "label.json", self.store, 1024))[
                    "label"
                ]
                bound = {
                    "sourceId": directory.name,
                    "registeredId": directory.name,
                    "label": self._display_name(label),
                    "manifest": manifest,
                }
                if manifest["source_id"] != directory.name:
                    raise SourceError("上传来源身份损坏")
                self.verify(bound)
                self._publish_upload(bound)
            except (SourceError, ValueError, KeyError, TypeError, OSError):
                self.entries[directory.name] = {
                    "public": {
                        "id": directory.name,
                        "label": "失效的已上传策略",
                        "ready": False,
                        "compatibleTasks": [],
                        "error": "上传快照损坏，拒绝随机初始化；请重新上传原文件或恢复快照",
                    }
                }

    @staticmethod
    def _display_name(name):
        if not isinstance(name, str):
            return "uploaded-policy"
        name = name.replace("\\", "/").split("/")[-1]
        return (
            re.sub(r"[^\w .()\-]", "_", name, flags=re.UNICODE)[:100].strip(" .")
            or "uploaded-policy"
        )

    def _publish_upload(self, bound):
        record = {
            "id": bound["sourceId"],
            "label": bound["label"],
            "ready": True,
            "compatibleTasks": TASKS,
            "observationSizes": [47, 81, 97],
            "initialization": bound,
        }
        with self._catalog_lock:
            self.entries[bound["sourceId"]] = {"public": record}
        return deepcopy(record)

    @contextmanager
    def upload_slot(self, length):
        # At most one receiving/decoding upload, 32 sources and 2GiB cumulative store.
        if not self._upload_lock.acquire(blocking=False):
            raise SourceError("已有文件正在上传/验证，请稍后重试")
        try:
            files = list(self.store.rglob("*"))
            used = sum(p.lstat().st_size for p in files if not p.is_symlink() and p.is_file())
            count = sum(1 for p in self.store.iterdir() if DIGEST.fullmatch(p.name))
            if count >= 32 or used + length + 16 * 1024**2 > 2 * 1024**3:
                raise SourceError("基础策略存储已达32个来源或2GiB上限，请由本机维护者释放空间")
            with tempfile.TemporaryDirectory(dir=self.store, prefix="upload-") as temporary:
                yield Path(temporary)
        finally:
            self._upload_lock.release()

    def receive_upload(self, stream, length, fmt, template, display_name, *, set_timeout=None):
        if template != "go2-legacy47-v1":
            raise SourceError("必须明确确认go2-legacy47-v1观测与关节模板")
        limit = {"pt": 256 * 1024**2, "onnx": 64 * 1024**2}.get(fmt)
        if limit is None:
            raise SourceError("仅支持单个.pt或.onnx，不接受ZIP/路径/配套目录")
        if type(length) is not int or not 0 < length <= limit:
            raise SourceError("上传文件为空或超过.pt 256MiB / ONNX 64MiB上限")
        with self.upload_slot(length) as directory:
            path = directory / f"upload.{fmt}"
            deadline = time.monotonic() + 60
            remaining = length
            try:
                with path.open("xb") as output:
                    while remaining:
                        timeout = deadline - time.monotonic()
                        if timeout <= 0:
                            raise SourceError("上传超时，请重试")
                        try:
                            if set_timeout is not None:
                                set_timeout(min(10, timeout))
                            # BufferedReader.read(n) may perform many recv calls whose socket
                            # timeouts reset with each trickled byte. read1 returns after one
                            # raw read, allowing the total deadline to be checked every time.
                            read_chunk = getattr(stream, "read1", stream.read)
                            chunk = read_chunk(min(1024 * 1024, remaining))
                        except OSError as error:
                            raise SourceError("上传超时/连接中断，临时文件已清理") from error
                        if not chunk or len(chunk) > remaining:
                            raise SourceError("上传连接中断或实际长度不符")
                        output.write(chunk)
                        remaining -= len(chunk)
            finally:
                if set_timeout is not None:
                    set_timeout(10)
            request = {
                "path": str(path),
                "directory": str(directory),
                "format": fmt,
                "template": template,
            }
            env = {
                **os.environ,
                "CUDA_VISIBLE_DEVICES": "",
                "OMP_NUM_THREADS": "1",
                "OPENBLAS_NUM_THREADS": "1",
                "MKL_NUM_THREADS": "1",
            }
            try:
                result = subprocess.run(
                    [self.python, str(Path(__file__).parent / "rl/scripts/validate_upload.py")],
                    input=json.dumps(request),
                    capture_output=True,
                    text=True,
                    cwd=self.trainer_root,
                    env=env,
                    timeout=60,
                )
                manifest = json.loads(result.stdout)
            except (OSError, subprocess.TimeoutExpired, ValueError) as error:
                raise SourceError("模型验证超时/资源超限或训练Python依赖不可用") from error
            if result.returncode:
                raise SourceError(manifest.get("error", "不支持的模型文件"))
            digest = manifest["source_id"]
            if not DIGEST.fullmatch(digest):
                raise SourceError("模型验证返回无效身份")
            label = self._display_name(display_name)
            (directory / "upload.json").write_text(json.dumps(manifest), encoding="utf-8")
            (directory / "label.json").write_text(json.dumps({"label": label}), encoding="utf-8")
            destination = self.store / digest
            if destination.exists():
                # Dedup never replaces a committed artifact or changes an old job binding.
                old = next((r for r in self.catalog() if r["id"] == digest and r["ready"]), None)
                if old is None:
                    raise SourceError("同ID旧快照已损坏，请恢复快照后重试；不会覆盖旧任务来源")
                self.verify(old["initialization"])
                return old
            for file in directory.iterdir():
                file.chmod(0o444)
            directory.rename(destination)
            destination.chmod(0o555)
            bound = {
                "sourceId": digest,
                "registeredId": digest,
                "label": label,
                "manifest": manifest,
            }
            self.verify(bound)
            return self._publish_upload(bound)

    def _validate(self, directory: Path, checkpoint: str, task_id: str, task_config: dict | None):
        command = [self.python, str(Path(__file__).parent / "rl/scripts/validate_pretrained.py")]
        request = {
            "directory": str(directory),
            "checkpoint": checkpoint,
            "taskId": task_id,
            "taskConfig": task_config,
            "uploaded": (directory / "upload.json").is_file(),
        }
        try:
            result = subprocess.run(
                command,
                input=json.dumps(request),
                cwd=self.trainer_root,
                capture_output=True,
                text=True,
                timeout=60,
            )
        except (subprocess.TimeoutExpired, OSError) as error:
            raise SourceError("基础策略验证器超时或不可用，请检查训练Python环境") from error
        if result.returncode:
            # Validator emits only an actionable error, never a traceback or file contents.
            raise SourceError(
                result.stderr.strip()[-1500:] or "基础策略验证器失败，请检查训练Python依赖"
            )
        try:
            return json.loads(result.stdout)
        except ValueError as error:
            raise SourceError("基础策略验证器未返回有效身份描述") from error

    def catalog(self):
        with self._catalog_lock:
            sources = {}
            for entry in self.entries.values():
                sources.setdefault(entry["public"]["id"], deepcopy(entry["public"]))
            return list(sources.values())

    def bind(self, source_id, task_id: str, task_config=None):
        # Selection is content-addressed, so a stale UI cannot silently bind a newly
        # registered model under the same administrator-friendly registration name.
        entry = next(
            (entry for entry in self.catalog() if entry["id"] == source_id),
            None,
        )
        if not isinstance(source_id, str) or entry is None:
            raise SourceError("基础策略内容ID不存在或已变化；请刷新服务连接，不接受文件路径")
        if not entry["ready"]:
            raise SourceError(entry["error"])
        if task_id not in entry["compatibleTasks"]:
            raise SourceError("该基础策略仅兼容Flat47/Obstacle81或97；Rough高度扫描不支持迁移")
        bound = deepcopy(entry["initialization"])
        directory = self.verify(bound)
        manifest = self._validate(
            directory, bound["manifest"]["artifacts"]["checkpoint"]["name"], task_id, task_config
        )
        if manifest["source_id"] != bound["sourceId"]:
            raise SourceError("基础策略快照身份变化，拒绝初始化")
        return bound

    def verify(self, bound: dict) -> Path:
        try:
            digest = bound["sourceId"]
            if not isinstance(digest, str) or not DIGEST.fullmatch(digest):
                raise SourceError("基础策略快照身份无效")
            directory = self.store / digest
            artifacts = bound["manifest"]["artifacts"]
            if "sourceFormat" in bound["manifest"]:
                fmt = bound["manifest"]["sourceFormat"]
                if fmt not in ("pt", "onnx"):
                    raise SourceError("上传格式无效")
                stored = json.loads(regular_bytes(directory / "upload.json", self.store, 64 * 1024))
                if stored != bound["manifest"]:
                    raise SourceError("上传manifest SHA绑定变化，拒绝初始化")
                for name, relative, limit in (
                    (
                        "upload",
                        f"upload.{fmt}",
                        LIMITS["checkpoint"] if fmt == "pt" else LIMITS["onnx"],
                    ),
                    ("checkpoint", "actor.pt", 16 * 1024**2),
                ):
                    content = regular_bytes(directory / relative, self.store, limit)
                    if (
                        artifacts[name]["name"] != relative
                        or hashlib.sha256(content).hexdigest() != artifacts[name]["sha256"]
                    ):
                        raise SourceError("上传快照SHA不匹配，拒绝训练")
                return directory
            for name, relative in {
                "checkpoint": artifacts["checkpoint"]["name"],
                "onnx": "policy.onnx",
                "env": "params/env.yaml",
                "agent": "params/agent.yaml",
            }.items():
                content = regular_bytes(directory / relative, self.store, LIMITS[name])
                if hashlib.sha256(content).hexdigest() != artifacts[name]["sha256"]:
                    raise SourceError("基础策略快照SHA不匹配，拒绝训练；请恢复原快照")
            return directory
        except (KeyError, TypeError) as error:
            raise SourceError("持久化基础策略描述损坏，拒绝训练") from error

    def arguments(self, bound: dict) -> list[str]:
        directory = self.verify(bound)
        upload_args = (
            ["--pretrained-upload-manifest", str(directory / "upload.json")]
            if "sourceFormat" in bound["manifest"]
            else []
        )
        return upload_args + [
            "--pretrained-checkpoint",
            str(directory / bound["manifest"]["artifacts"]["checkpoint"]["name"]),
            "--pretrained-allowed-roots",
            json.dumps([str(directory)]),
            "--pretrained-source-id",
            bound["sourceId"],
        ]