"""Single-file upload security, strict graph reconstruction, and shared warm-start seam."""

import copy
import errno
import hashlib
import http.client
import importlib.util
import io
import json
import os
import socket
import sys
import tempfile
import threading
import unittest
from contextlib import contextmanager
from dataclasses import asdict
from http.server import ThreadingHTTPServer
from pathlib import Path
from unittest.mock import patch

ROOT = Path(__file__).resolve().parents[1]
for root in (ROOT, ROOT / "rl"):
    sys.path.insert(0, str(root))

from pretrained_sources import PretrainedSources, SourceError  # noqa: E402
from server import TrainingManager, TrainingRequestHandler  # noqa: E402


def fixture_onnx(state):
    import onnx
    from onnx import helper, numpy_helper
    from pretrained_upload import SEMANTICS

    tensors = [
        numpy_helper.from_array(value.numpy(), key)
        for key, value in state.items()
        if key.startswith("mlp.") or key == "obs_normalizer._mean"
    ]
    tensors.append(
        numpy_helper.from_array((state["obs_normalizer._std"] + 0.01).numpy(), "onnx::Div_24")
    )
    nodes = [
        helper.make_node("Sub", ["obs", "obs_normalizer._mean"], ["centered"]),
        helper.make_node("Div", ["centered", "onnx::Div_24"], ["normalized"]),
    ]
    previous = "normalized"
    for i in (0, 2, 4, 6):
        output = "actions" if i == 6 else f"linear{i}"
        nodes.append(
            helper.make_node(
                "Gemm", [previous, f"mlp.{i}.weight", f"mlp.{i}.bias"], [output], transB=1
            )
        )
        previous = output
        if i < 6:
            previous = f"elu{i}"
            nodes.append(helper.make_node("Elu", [output], [previous], alpha=1.0))
    model = helper.make_model(
        helper.make_graph(
            nodes,
            "actor",
            [helper.make_tensor_value_info("obs", onnx.TensorProto.FLOAT, [1, 47])],
            [helper.make_tensor_value_info("actions", onnx.TensorProto.FLOAT, [1, 12])],
            tensors,
        ),
        opset_imports=[helper.make_opsetid("", 18)],
        ir_version=8,
    )
    helper.set_model_props(
        model, {key: ",".join(map(str, value)) for key, value in SEMANTICS.items()}
    )
    return model


@unittest.skipUnless(
    all(importlib.util.find_spec(m) is not None for m in ("torch", "mjlab", "onnx", "onnxruntime")),
    "installed CPU RL/ONNX stack required",
)
class UploadTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        import torch
        from pretrained import make_reference_actor

        torch.set_num_threads(1)
        torch.manual_seed(123)
        actor = make_reference_actor()
        actor.obs_normalizer.count.fill_(983138304)
        cls.state = actor.state_dict()
        stream = io.BytesIO()
        torch.save({"actor_state_dict": cls.state, "iter": 10000}, stream)
        cls.pt = stream.getvalue()
        cls.onnx = fixture_onnx(cls.state).SerializeToString()

    def setUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.root = Path(self.temp.name)
        self.registry = PretrainedSources(None, self.root / "store", sys.executable, ROOT / "rl")

    def tearDown(self):
        self.temp.cleanup()

    def upload(self, data=None, fmt="pt", name="policy.pt"):
        data = self.pt if data is None else data
        return self.registry.receive_upload(
            io.BytesIO(data), len(data), fmt, "go2-legacy47-v1", name
        )

    def test_single_file_persistence_dedup_and_binding_cli(self):
        record = self.upload(name="../../etc/evil\n.pt")
        self.assertEqual(record["label"], "evil_.pt")
        bound = record["initialization"]
        directory = self.registry.verify(bound)
        self.assertEqual(
            {p.name for p in directory.iterdir()},
            {"upload.pt", "actor.pt", "upload.json", "label.json"},
        )
        self.assertEqual(self.upload(name="renamed.pt"), record)
        restored = PretrainedSources(None, self.root / "store", sys.executable, ROOT / "rl")
        self.assertEqual(restored.catalog(), [record])
        self.assertEqual(restored.bind(record["id"], "Unitree-Go2-Flat"), bound)
        args = restored.arguments(bound)
        self.assertIn("--pretrained-upload-manifest", args)
        self.assertNotIn("--resume-checkpoint", args)
        manager = TrainingManager(
            ROOT / "rl",
            sys.executable,
            ("Unitree-Go2-Flat",),
            check_environment=False,
            sources=restored,
        )
        cfg = manager.parse_config(
            {
                "taskId": "Unitree-Go2-Flat",
                "numEnvs": 4096,
                "maxIterations": 1,
                "seed": 42,
                "device": "cpu",
                "pretrainedSourceId": record["id"],
            }
        )
        self.assertEqual(cfg.pretrained, bound)
        self.assertIn("--pretrained-upload-manifest", manager.command_for(cfg))
        from scripts.train import TrainConfig, _load_pretrained
        from src.tasks.velocity.config.go2.env_cfgs import unitree_go2_flat_env_cfg
        from src.tasks.velocity.config.go2.rl_cfg import unitree_go2_ppo_runner_cfg

        train = TrainConfig(
            env=unitree_go2_flat_env_cfg(),
            agent=unitree_go2_ppo_runner_cfg(),
            pretrained_checkpoint=str(directory / "actor.pt"),
            pretrained_upload_manifest=str(directory / "upload.json"),
            pretrained_allowed_roots=[str(directory)],
            pretrained_source_id=record["id"],
        )
        source = _load_pretrained(train)
        self.assertEqual(source.manifest, bound["manifest"])
        from pretrained import public_initialization_metadata

        public = public_initialization_metadata(source.manifest)
        self.assertEqual(public["uploadSha256"], hashlib.sha256(self.pt).hexdigest())
        self.assertEqual(public["templateId"], "go2-legacy47-v1")
        self.assertNotIn(str(self.root), json.dumps(public))
        # Changing upload content never changes the previously bound job.
        newer = self.upload(self.onnx, "onnx")
        self.assertNotEqual(newer["id"], record["id"])
        self.assertEqual(restored.arguments(cfg.pretrained), args)
        (directory / "actor.pt").chmod(0o644)
        (directory / "actor.pt").write_bytes(b"tamper")
        with self.assertRaisesRegex(SourceError, "SHA"):
            restored.arguments(bound)
        broken = PretrainedSources(None, self.root / "store", sys.executable, ROOT / "rl")
        self.assertFalse(next(r for r in broken.catalog() if r["id"] == record["id"])["ready"])
        with self.assertRaises(SourceError):
            broken.bind(record["id"], "Unitree-Go2-Flat")

    def test_tuning_same_binding_new_trial_and_own_rung_resume(self):
        from test_tuning_manager import BASE_METRICS, FakeAdvisor
        from tuning.manager import TuningManager
        from tuning.process import GpuLease
        from tuning.schema import BASE_REWARD_CONFIGURATION

        record = self.upload(self.onnx, "onnx")
        manager = TuningManager(
            ROOT / "rl",
            sys.executable,
            self.root / "tuning",
            GpuLease(),
            advisor=FakeAdvisor(),
            sources=self.registry,
        )
        _, config, objective, fallback = manager.parse_create(
            {
                "mode": "approval",
                "pretrainedSourceId": record["id"],
                "trialCount": 2,
                "numEnvs": 4096,
                "initialIterations": 1,
                "middleIterations": 2,
                "finalIterations": 3,
            }
        )
        session = manager.storage.create_session("approval", config, objective, fallback)
        self.assertEqual(config["pretrained"], record["initialization"])
        from pretrained import public_initialization_metadata

        public = public_initialization_metadata(config["pretrained"]["manifest"])
        self.assertEqual(public["sourceFormat"], "onnx")
        self.assertIsNone(public["sourceIteration"])
        self.assertEqual(public["derivedFields"]["normalizer_count"]["value"], 1_000_000)
        self.assertNotIn(
            "onnxSha256", public
        )  # Do not confuse generated checkpoint with user upload.
        manager.cancel_events[session["id"]] = threading.Event()
        commands = []

        def run(_session, command, _cwd, _env, _log):
            commands.append(command)
            if "--output-dir" in command:
                directory = Path(command[command.index("--output-dir") + 1])
                (directory / "model_0.pt").write_bytes(b"own checkpoint with updated statistics")
                (directory / "policy.onnx").write_bytes(b"trial policy")
                (directory / "initialization.json").write_text(
                    json.dumps(record["initialization"]["manifest"])
                )
            else:
                Path(command[command.index("--output") + 1]).write_text(
                    json.dumps({"metrics": BASE_METRICS})
                )
            return 0

        with patch.object(manager, "_run_command", run):
            parent = None
            for number, rung in ((0, 0), (1, 0), (1, 1)):
                trial = manager.storage.create_trial(
                    session["id"],
                    number,
                    rung,
                    rung + 1,
                    BASE_REWARD_CONFIGURATION,
                    None,
                    f"trial-{number}-{rung}",
                )
                result = manager._execute_trial(session, trial, parent if rung else None)
                parent = manager._session_root(session["id"]) / result["checkpointPath"]
        train = [c for c in commands if "--output-dir" in c]
        for command in train[:2]:
            self.assertIn("--pretrained-upload-manifest", command)
            self.assertEqual(command[command.index("--pretrained-source-id") + 1], record["id"])
        self.assertNotIn("--pretrained-upload-manifest", train[2])
        self.assertNotIn("--pretrained-checkpoint", train[2])
        self.assertIn("trial-1-0/model_0.pt", train[2][-1])
        restored_sources = PretrainedSources(None, self.registry.store, sys.executable, ROOT / "rl")
        restored = TuningManager(
            ROOT / "rl",
            sys.executable,
            self.root / "tuning",
            GpuLease(),
            advisor=FakeAdvisor(),
            sources=restored_sources,
        )
        self.assertEqual(
            restored.detail(session["id"])["config"]["pretrained"], record["initialization"]
        )
        self.assertFalse(restored.workers)

    def test_input_limits_disconnect_timeout_capacity_and_cleanup(self):
        for length, fmt, template in (
            (0, "pt", "go2-legacy47-v1"),
            (256 * 1024**2 + 1, "pt", "go2-legacy47-v1"),
            (64 * 1024**2 + 1, "onnx", "go2-legacy47-v1"),
            (1, "zip", "go2-legacy47-v1"),
            (1, "pt", ""),
        ):
            with (
                self.subTest(length=length, fmt=fmt, template=template),
                self.assertRaises(SourceError),
            ):
                self.registry.receive_upload(io.BytesIO(b"x"), length, fmt, template, "x")
        with self.assertRaisesRegex(SourceError, "中断"):
            self.registry.receive_upload(io.BytesIO(b"x"), 100, "pt", "go2-legacy47-v1", "x")
        with (
            patch("pretrained_sources.time.monotonic", side_effect=[0, 61]),
            self.assertRaisesRegex(SourceError, "超时"),
        ):
            self.upload()
        trickle = unittest.mock.Mock()
        trickle.read.side_effect = AssertionError("buffered read could bypass total deadline")
        trickle.read1.return_value = b"x"
        with (
            patch("pretrained_sources.time.monotonic", side_effect=[0, 1, 61]),
            self.assertRaisesRegex(SourceError, "超时"),
        ):
            self.registry.receive_upload(trickle, 10, "pt", "go2-legacy47-v1", "x")
        self.assertEqual(trickle.read1.call_count, 1)
        stream = unittest.mock.Mock()
        stream.read1.side_effect = TimeoutError()
        with self.assertRaises(SourceError):
            self.registry.receive_upload(stream, 10, "pt", "go2-legacy47-v1", "x")
        with self.registry.upload_slot(1), self.assertRaisesRegex(SourceError, "正在上传"):
            self.upload()
        self.assertEqual(list(self.registry.store.iterdir()), [])
        import subprocess

        with (
            patch(
                "pretrained_sources.subprocess.run",
                side_effect=subprocess.TimeoutExpired("validator", 60),
            ),
            self.assertRaisesRegex(SourceError, "超时"),
        ):
            self.upload()
        self.assertEqual(list(self.registry.store.iterdir()), [])
        huge = self.registry.store / "quota"
        with huge.open("wb") as stream:
            stream.truncate(2 * 1024**3)
        with self.assertRaisesRegex(SourceError, "上限"):
            self.upload()
        huge.unlink()
        for n in range(32):
            (self.registry.store / f"{n:064x}").mkdir()
        with self.assertRaisesRegex(SourceError, "上限"):
            self.upload()

    def test_unsafe_pickle_invalid_checkpoint_and_zip_never_publish(self):
        import torch

        marker = self.root / "unsafe-executed"

        class Unsafe:
            def __reduce__(self):
                return (os.system, (f"touch {marker}",))

        bad = dict(self.state)
        bad["mlp.0.weight"] = torch.full((512, 47), float("nan"))
        wrong = dict(self.state)
        wrong["mlp.0.weight"] = torch.zeros(512, 97)
        cases = [b"PK\x03\x04bad zip", b"not a model"]
        for payload in (
            {"actor_state_dict": bad},
            {"actor_state_dict": wrong},
            {"actor_state_dict": self.state, "metadata": {"joint_names": ["wrong"]}},
            {"actor_state_dict": self.state, "evil": Unsafe()},
        ):
            stream = io.BytesIO()
            torch.save(payload, stream)
            cases.append(stream.getvalue())
        for data in cases:
            with self.subTest(size=len(data)), self.assertRaises(SourceError):
                self.upload(data)
            self.assertEqual(self.registry.catalog(), [])
            self.assertEqual(list(self.registry.store.iterdir()), [])
        self.assertFalse((self.root / "unsafe-executed").exists())

    def test_strict_graph_rejects_bad_edges_ops_attributes_external_and_metadata(self):
        import onnx
        from pretrained import PretrainedError
        from pretrained_upload import _onnx

        base = onnx.load_model_from_string(self.onnx)

        def external(m):
            m.graph.initializer[0].data_location = onnx.TensorProto.EXTERNAL
            m.graph.initializer[0].external_data.add(key="location", value="/etc/passwd")

        def attrs(m):
            m.graph.node[2].attribute.add(name="transA", type=onnx.AttributeProto.INT, i=1)

        def metadata(m):
            next(p for p in m.metadata_props if p.key == "joint_names").value = "wrong"

        def extra(m):
            m.graph.node.append(onnx.helper.make_node("Identity", ["actions"], ["branch"]))

        def shape(m):
            m.graph.input[0].type.tensor_type.shape.dim[1].dim_value = 97

        edits = [
            external,
            attrs,
            metadata,
            extra,
            shape,
            lambda m: setattr(m.graph.node[2], "domain", "evil"),
            lambda m: setattr(m.graph.node[3], "op_type", "Relu"),
            lambda m: m.graph.node[4].input.__setitem__(0, "normalized"),
            lambda m: m.graph.node[0].input.reverse(),
            lambda m: m.graph.node[4].input.__setitem__(1, "mlp.0.weight"),
            lambda m: setattr(m.graph.output[0], "name", "linear0"),
            lambda m: setattr(m.graph.initializer[0], "data_type", onnx.TensorProto.DOUBLE),
        ]
        for edit in edits:
            model = copy.deepcopy(base)
            edit(model)
            with self.subTest(edit=edit), self.assertRaises((PretrainedError, ValueError)):
                _onnx(model.SerializeToString())
        bad = copy.deepcopy(base)
        external(bad)
        with self.assertRaisesRegex(SourceError, "external"):
            self.upload(bad.SerializeToString(), "onnx")
        self.assertEqual(self.registry.catalog(), [])

    def test_http_upload_storage_errors_are_503_but_read_failures_are_400(self):
        manager = TrainingManager(
            ROOT / "rl",
            sys.executable,
            ("Unitree-Go2-Flat",),
            check_environment=False,
            sources=self.registry,
        )
        reset_timeouts = []

        class Handler(TrainingRequestHandler):
            access_token = "upload-test-token"
            read_error = None

            def log_message(self, *_args):
                pass

            def _upload(self):
                original_stream = self.rfile
                if self.read_error is not None:
                    self.rfile = unittest.mock.Mock(wraps=original_stream)
                    self.rfile.read1.side_effect = self.read_error
                try:
                    super()._upload()
                finally:
                    self.rfile = original_stream
                    reset_timeouts.append(self.connection.gettimeout())

        Handler.manager = manager
        server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
        thread = threading.Thread(target=server.serve_forever, daemon=True)
        thread.start()
        original_open = Path.open
        url = (
            "/api/training/pretrained-sources/upload"
            "?format=pt&template=go2-legacy47-v1&name=policy.pt"
        )

        try:
            for phase, error in (
                ("open", OSError(errno.ENOSPC, "disk full", "/private/upload.pt")),
                ("open", OSError(errno.EACCES, "permission denied", "/private/upload.pt")),
                ("write", OSError(errno.ENOSPC, "disk full", "/private/upload.pt")),
                ("close", OSError(errno.ENOSPC, "flush failed", "/private/upload.pt")),
                ("read", ConnectionResetError(errno.ECONNRESET, "connection reset")),
                ("read", TimeoutError("read timeout")),
                ("truncated", None),
            ):
                with self.subTest(phase=phase, error=error):
                    Handler.read_error = error if phase == "read" else None

                    @contextmanager
                    def faulty_output(path, *args, phase=phase, error=error, **kwargs):
                        with original_open(path, *args, **kwargs) as output:
                            if phase == "write":
                                output = unittest.mock.Mock(wraps=output)
                                output.write.side_effect = error
                            yield output
                        if phase == "close":
                            raise error

                    def open_file(
                        path,
                        *args,
                        phase=phase,
                        error=error,
                        output_factory=faulty_output,
                        **kwargs,
                    ):
                        if path.name == "upload.pt" and args == ("xb",):
                            if phase == "open":
                                raise error
                            return output_factory(path, *args, **kwargs)
                        return original_open(path, *args, **kwargs)

                    conn = http.client.HTTPConnection(*server.server_address, timeout=10)
                    try:
                        with patch.object(Path, "open", open_file):
                            conn.request(
                                "POST",
                                url,
                                body=b"x",
                                headers={
                                    "Authorization": "Bearer upload-test-token",
                                    "Content-Type": "application/octet-stream",
                                    "Content-Length": "2" if phase == "truncated" else "1",
                                },
                            )
                            if phase == "truncated":
                                conn.sock.shutdown(socket.SHUT_WR)
                            response = conn.getresponse()
                            message = json.loads(response.read())["error"]
                    finally:
                        conn.close()
                    storage_error = phase in ("open", "write", "close")
                    self.assertEqual(response.status, 503 if storage_error else 400)
                    self.assertIn("磁盘空间/权限" if storage_error else "中断", message)
                    self.assertNotIn("/private", message)
                    self.assertNotIn(str(self.root), message)
                    self.assertEqual(reset_timeouts[-1], 10)
                    self.assertEqual(list(self.registry.store.iterdir()), [])
                    self.assertEqual(self.registry.catalog(), [])
                    self.assertEqual(manager.jobs, {})
        finally:
            server.shutdown()
            server.server_close()
            thread.join()

    def test_http_auth_json_limit_binary_headers_and_incomplete_body(self):
        manager = TrainingManager(
            ROOT / "rl",
            sys.executable,
            ("Unitree-Go2-Flat",),
            check_environment=False,
            sources=self.registry,
        )

        class Handler(TrainingRequestHandler):
            access_token = "upload-test-token"

            def log_message(self, *_args):
                pass

        Handler.manager = manager
        server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
        thread = threading.Thread(target=server.serve_forever, daemon=True)
        thread.start()
        url = (
            "/api/training/pretrained-sources/upload"
            "?format=pt&template=go2-legacy47-v1&name=..%2Ftest.pt"
        )
        headers = {
            "Authorization": "Bearer upload-test-token",
            "Content-Type": "application/octet-stream",
        }

        def request(path=url, data=b"x", extra=None):
            conn = http.client.HTTPConnection(*server.server_address, timeout=30)
            conn.request("POST", path, body=data, headers={**headers, **(extra or {})})
            response = conn.getresponse()
            status, body = response.status, response.read()
            conn.close()
            return status, json.loads(body)

        try:
            self.assertEqual(request(extra={"Authorization": "Bearer wrong"})[0], 401)
            self.assertEqual(request(extra={"Origin": "https://evil.example"})[0], 403)
            self.assertEqual(request(extra={"Host": "evil.example"})[0], 403)
            self.assertEqual(request(extra={"Content-Type": "application/json"})[0], 400)
            self.assertEqual(request(extra={"Content-Length": str(256 * 1024**2 + 1)})[0], 413)
            self.assertEqual(request(extra={"Content-Encoding": "gzip"})[0], 400)
            self.assertEqual(request(extra={"Transfer-Encoding": "chunked"})[0], 400)
            self.assertEqual(
                request(path=url.replace("template=go2-legacy47-v1", "template="))[0], 400
            )
            self.assertEqual(
                request(path="/api/training/jobs", data=b" " * (128 * 1024 + 1))[0], 413
            )
            # A blocked decoder is isolated from health/catalog requests.
            import subprocess
            from types import SimpleNamespace

            Handler.tuning_manager = SimpleNamespace(capability=lambda: {})
            started, release = threading.Event(), threading.Event()
            original_run = subprocess.run

            def blocked(*args, **kwargs):
                started.set()
                release.wait(10)
                return original_run(*args, **kwargs)

            outcomes = []
            with patch("pretrained_sources.subprocess.run", side_effect=blocked):
                worker = threading.Thread(target=lambda: outcomes.append(request(data=self.pt)))
                worker.start()
                self.assertTrue(started.wait(5))
                health = http.client.HTTPConnection(*server.server_address, timeout=3)
                health.request("GET", "/api/training/health", headers=headers)
                response = health.getresponse()
                self.assertEqual(response.status, 200)
                self.assertTrue(json.loads(response.read())["pretrainedUpload"]["enabled"])
                health.close()
                release.set()
                worker.join(30)
            status, record = outcomes[0]
            self.assertEqual(status, 201)
            self.assertEqual(record["label"], "test.pt")
            status, onnx_record = request(
                path=url.replace("format=pt", "format=onnx"), data=self.onnx
            )
            self.assertEqual(status, 201)
            self.assertEqual(onnx_record["initialization"]["manifest"]["sourceFormat"], "onnx")
            raw = socket.create_connection(server.server_address, timeout=10)
            raw.sendall(
                (
                    f"POST {url} HTTP/1.0\r\nHost: localhost\r\n"
                    "Authorization: Bearer upload-test-token\r\n"
                    "Content-Type: application/octet-stream\r\nContent-Length: 1000\r\n\r\nx"
                ).encode()
            )
            raw.shutdown(socket.SHUT_WR)
            self.assertIn(b"400", raw.recv(4096))
            raw.close()
            self.assertFalse(list(self.registry.store.glob("upload-*")))
        finally:
            server.shutdown()
            server.server_close()
            thread.join()

    def test_onnx_expansion_and_checkpoint_preserve_updated_count(self):
        import torch
        from pretrained import (
            ValidatedSource,
            comparison_observations,
            make_reference_actor,
            warm_start_actor,
        )
        from pretrained_upload import _onnx
        from tensordict import TensorDict

        state, _, identity = _onnx(self.onnx)
        self.assertLess(identity["max_abs_error"], 2e-5)
        base = comparison_observations()
        reference = make_reference_actor().eval()
        reference.load_state_dict(state)
        expected = reference.mlp(reference.obs_normalizer(base)).detach()
        for dim in (81, 97):
            actor = make_reference_actor(dim)
            warm_start_actor(actor, ValidatedSource(state, {}))
            self.assertEqual(actor.mlp[0].weight[:, 47:].count_nonzero().item(), 0)
            values = torch.cat((base, torch.ones(len(base), dim - 47)), dim=1)
            torch.testing.assert_close(
                actor.mlp(actor.obs_normalizer(values)), expected, atol=2e-5, rtol=2e-5
            )
            batch = TensorDict({"actor": values}, batch_size=[len(base)])
            actor.update_normalization(batch)
            self.assertEqual(actor.obs_normalizer.count.item(), 1_000_048)
            self.assertGreater(actor.obs_normalizer._mean[:, 47:].min().item(), 0)
            actor(batch).square().mean().backward()
            grad = actor.mlp[0].weight.grad[:, 47:]
            self.assertTrue(torch.isfinite(grad).all())
            self.assertGreater(grad.abs().max().item(), 0)
            torch.optim.Adam(actor.parameters(), lr=1e-4).step()
            buffer = io.BytesIO()
            torch.save(actor.state_dict(), buffer)
            buffer.seek(0)
            restored = make_reference_actor(dim)
            restored.load_state_dict(torch.load(buffer, weights_only=True))
            for key, tensor in actor.state_dict().items():
                self.assertTrue(torch.equal(tensor, restored.state_dict()[key]), key)


@unittest.skipUnless(
    os.environ.get("GO2_UPLOAD_REAL_DIR"), "opt-in read-only real single-file probes"
)
class RealSingleFileUploadTest(UploadTest):
    """Inherited tests run against real uploads without ever requesting a sidecar."""

    @classmethod
    def setUpClass(cls):
        import torch

        torch.set_num_threads(1)
        root = Path(os.environ["GO2_UPLOAD_REAL_DIR"])
        # Explicit two files read separately as bytes; each upload receives only one.
        cls.pt = (root / "model_10000.pt").read_bytes()
        cls.onnx = (root / "policy.onnx").read_bytes()
        cls.state = torch.load(io.BytesIO(cls.pt), map_location="cpu", weights_only=True)[
            "actor_state_dict"
        ]

    def test_real_independent_uploads_ort_and_batch_drift(self):
        import torch
        from pretrained import comparison_observations, make_reference_actor, verify_onnx
        from pretrained_upload import read_uploaded_source
        from src.tasks.velocity.config.go2.env_cfgs import unitree_go2_flat_env_cfg
        from src.tasks.velocity.config.go2.rl_cfg import unitree_go2_ppo_runner_cfg

        evidence = {}
        for fmt, data in (("pt", self.pt), ("onnx", self.onnx)):
            with tempfile.TemporaryDirectory() as empty:
                registry = PretrainedSources(None, Path(empty), sys.executable, ROOT / "rl")
                record = registry.receive_upload(
                    io.BytesIO(data), len(data), fmt, "go2-legacy47-v1", f"single.{fmt}"
                )
                directory = registry.verify(record["initialization"])
                source = read_uploaded_source(
                    directory / "actor.pt",
                    allowed_roots=[directory],
                    manifest_path=directory / "upload.json",
                    target_env=asdict(unitree_go2_flat_env_cfg()),
                    target_agent=asdict(unitree_go2_ppo_runner_cfg()),
                )
                actor = make_reference_actor().eval()
                actor.load_state_dict(source.actor_state)
                identity = verify_onnx(self.onnx, actor)
                self.assertLess(identity["max_abs_error"], 2e-5)
                evidence[fmt] = {
                    "sourceId": record["id"],
                    "sha256": hashlib.sha256(data).hexdigest(),
                    "identity": identity,
                    "normalizer_count": actor.obs_normalizer.count.item(),
                    "updates": {},
                }
                probes = comparison_observations()[32:]
                before = actor.mlp(actor.obs_normalizer(probes)).detach()
                # Ordinary UI default=4096 and service maximum=16384 environments.
                for batch_size in (4096, 16384):
                    updated = copy.deepcopy(actor).train()
                    updated.obs_normalizer.update(probes.repeat(batch_size // len(probes), 1))
                    after = updated.mlp(updated.obs_normalizer(probes)).detach()
                    evidence[fmt]["updates"][str(batch_size)] = {
                        "rate": batch_size / updated.obs_normalizer.count.item(),
                        "max_action_drift": (after - before).abs().max().item(),
                        "max_mean_drift": (
                            updated.obs_normalizer._mean - actor.obs_normalizer._mean
                        )
                        .abs()
                        .max()
                        .item(),
                        "count_after": updated.obs_normalizer.count.item(),
                    }
                    self.assertTrue(torch.isfinite(after).all())
                self.assertEqual(
                    PretrainedSources(None, Path(empty), sys.executable, ROOT / "rl").catalog(),
                    [record],
                )
        print("REAL_UPLOAD_EVIDENCE=" + json.dumps(evidence, sort_keys=True))


if __name__ == "__main__":
    unittest.main()