# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Contract tests for final-layout BF16 Host Weight Runtime artifacts."""

from __future__ import annotations

import dataclasses
import gc
from collections.abc import Sequence
from pathlib import Path
from types import SimpleNamespace

import pytest
import torch
from torch import nn

from vllm_omni.diffusion.model_loader.host_weights import (
    FINAL_LAYOUT_BF16_POLICY,
    FINAL_LAYOUT_BF16_SPEC,
    FinalLayoutBF16Producer,
    FinalLayoutIdentityContext,
    FinalLayoutLoaderIdentity,
    FinalLayoutParallelIdentity,
    FinalLayoutRequest,
    FinalLayoutTensorRestorer,
    ImplementationIdentity,
    PreparedWeightSource,
    WeightSourceKind,
    build_final_layout_identity,
)
from vllm_omni.diffusion.model_loader.host_weights.contracts import (
    FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA,
    FinalLayoutArtifactSpec,
    FinalLayoutContractCode,
    FinalLayoutContractError,
    FinalLayoutTensorPolicy,
    implementation_abi_fingerprint,
)
from vllm_omni.diffusion.model_loader.host_weights.tensor_layout import (
    RuntimeTensorTarget,
    collect_final_layout_targets,
)
from vllm_omni.diffusion.models.host_weight_contract import (
    FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA,
    FinalLayoutModelContract,
)
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import MiniMaxH3DiTModel
from vllm_omni.host_weight_runtime import (
    AdaptationIdentity,
    CanonicalJson,
    CoordinationScope,
    FailureCode,
    HostWeightRuntime,
    HostWeightRuntimeConfig,
    LookupPhase,
    PostLoadPublicationOutcome,
    ProducerIdentity,
    ProductionPolicy,
    ProductionSourceMode,
    ResolutionOutcome,
    ResolutionStage,
    RuntimeMode,
    StorageDomainPolicy,
    StorageScope,
    TensorKind,
    WaitPolicy,
    WeightRepresentation,
)
from vllm_omni.host_weight_runtime.filesystem import FilesystemHostWeightStore, detect_storage_class

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


class _TinyDiT(nn.Module):
    host_weight_restore_contract = FinalLayoutModelContract(
        implementation_id="test-tiny-dit",
        version="1",
    )

    def __init__(self, *, device: torch.device | str = "cpu") -> None:
        super().__init__()
        target = torch.device(device)
        self.restore_validations = 0
        self.proj = nn.Linear(3, 2, dtype=torch.bfloat16, device=target)
        self.fp32_gain = nn.Parameter(torch.empty(1, dtype=torch.float32, device=target))
        self.register_buffer("scale", torch.empty(1, dtype=torch.float32, device=target))
        self.register_buffer(
            "derived",
            torch.tensor([7.0], dtype=torch.float32),
            persistent=False,
        )

    def validate_restored_host_weights(self) -> None:
        self.restore_validations += 1
        assert not self.proj.weight.is_meta
        assert self.proj.weight.dtype is torch.bfloat16
        assert self.fp32_gain.dtype is torch.float32
        assert self.scale.dtype is torch.float32


class _TinyPipeline(nn.Module):
    def __init__(self, *, dit_device: torch.device | str = "cpu") -> None:
        super().__init__()
        self.transformer = _TinyDiT(device=dit_device)
        self.text_encoder = nn.Linear(2, 2, dtype=torch.bfloat16)
        self.vae = nn.Module()
        self.vae.register_buffer("gain", torch.tensor([9.0], dtype=torch.float32))


class _AlternateTinyPipeline(_TinyPipeline):
    pass


class _MultiTinyPipeline(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.transformer_a = _TinyDiT()
        self.transformer_b = _TinyDiT()


class _Float32DiT(nn.Module):
    host_weight_restore_contract = FinalLayoutModelContract(
        implementation_id="test-float32-dit",
        version="1",
    )

    def __init__(self) -> None:
        super().__init__()
        self.proj = nn.Linear(2, 2, dtype=torch.float32)

    def validate_restored_host_weights(self) -> None:
        assert self.proj.weight.dtype is torch.float32


class _Float32Pipeline(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.transformer = _Float32DiT()


_SYNTHETIC_FP32_ABI = CanonicalJson.from_value(
    {
        "artifact_identity": "diffusion-final-layout-identity-v1",
        "model_contract": FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA,
        "producer": "test-fp32-producer-v1",
        "representation_policy": "test-all-fp32-v1",
        "restorer": "exact-final-layout-tensor-rebind-v1",
        "source_identity": "prepared-diffusion-weight-source-v2",
        "tensor_contract": "complete-strided-tensor-ownership-v1",
    }
)
_SYNTHETIC_FP32_SPEC = FinalLayoutArtifactSpec(
    representation=WeightRepresentation(
        name="test-final-layout-fp32",
        dtype=str(torch.float32),
        metadata=CanonicalJson.from_value({"policy": "all-fp32"}),
    ),
    producer=ProducerIdentity(
        producer_id="test.final-layout-fp32",
        version="1",
        implementation_fingerprint=implementation_abi_fingerprint(_SYNTHETIC_FP32_ABI),
        manifest_schema="test-final-layout-fp32-manifest-v1",
        restorer_schema=FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA,
    ),
    implementation_abi=_SYNTHETIC_FP32_ABI,
)


class _SyntheticFP32Policy:
    @property
    def spec(self) -> FinalLayoutArtifactSpec:
        return _SYNTHETIC_FP32_SPEC

    def validate_request(self, request: FinalLayoutRequest) -> None:
        if request.load_format != "default":
            raise ValueError("synthetic policy requires default loading")

    def tensor_role(self, name: str, tensor: torch.Tensor, kind: TensorKind) -> str:
        del name, tensor
        return "weight" if kind is TensorKind.PARAMETER else "persistent_buffer"

    def validate_target(self, target: RuntimeTensorTarget) -> None:
        if target.tensor.dtype is not torch.float32:
            raise ValueError("synthetic policy requires FP32 tensors")

    def validate_collection(self, targets: Sequence[RuntimeTensorTarget]) -> None:
        if not targets:
            raise ValueError("synthetic policy requires tensors")

    def build_format_metadata(
        self,
        *,
        component_names: tuple[str, ...],
        tensor_contract_digest: str,
        tensor_count: int,
    ) -> CanonicalJson:
        return CanonicalJson.from_value(
            {
                "component_names": list(component_names),
                "format": "test-final-layout-fp32",
                "tensor_contract_sha256": tensor_contract_digest,
                "tensor_count": tensor_count,
            }
        )

    def validate_format_metadata(
        self,
        metadata: CanonicalJson,
        *,
        component_names: tuple[str, ...],
        tensor_contract_digest: str,
        tensor_count: int,
    ) -> None:
        assert metadata == self.build_format_metadata(
            component_names=component_names,
            tensor_contract_digest=tensor_contract_digest,
            tensor_count=tensor_count,
        )


_SYNTHETIC_FP32_POLICY = _SyntheticFP32Policy()


class _PolicyWithSpec:
    def __init__(self, base: FinalLayoutTensorPolicy, spec: FinalLayoutArtifactSpec) -> None:
        self._base = base
        self._spec = spec

    @property
    def spec(self) -> FinalLayoutArtifactSpec:
        return self._spec

    def validate_request(self, request: FinalLayoutRequest) -> None:
        self._base.validate_request(request)

    def tensor_role(self, name: str, tensor: torch.Tensor, kind: TensorKind) -> str:
        return self._base.tensor_role(name, tensor, kind)

    def validate_target(self, target: RuntimeTensorTarget) -> None:
        self._base.validate_target(target)

    def validate_collection(self, targets: Sequence[RuntimeTensorTarget]) -> None:
        self._base.validate_collection(targets)

    def build_format_metadata(
        self,
        *,
        component_names: tuple[str, ...],
        tensor_contract_digest: str,
        tensor_count: int,
    ) -> CanonicalJson:
        return self._base.build_format_metadata(
            component_names=component_names,
            tensor_contract_digest=tensor_contract_digest,
            tensor_count=tensor_count,
        )

    def validate_format_metadata(
        self,
        metadata: CanonicalJson,
        *,
        component_names: tuple[str, ...],
        tensor_contract_digest: str,
        tensor_count: int,
    ) -> None:
        self._base.validate_format_metadata(
            metadata,
            component_names=component_names,
            tensor_contract_digest=tensor_contract_digest,
            tensor_count=tensor_count,
        )


class _FakeH3Linear(nn.Module):
    def __init__(
        self,
        *args: object,
        bias: bool = False,
        params_dtype: torch.dtype | None = None,
        total_num_heads: int | None = None,
        total_num_kv_heads: int | None = None,
        **kwargs: object,
    ) -> None:
        del args, kwargs
        super().__init__()
        dtype = params_dtype or torch.get_default_dtype()
        self.weight = nn.Parameter(torch.empty(1, dtype=dtype))
        if bias:
            self.bias = nn.Parameter(torch.empty(1, dtype=dtype))
        else:
            self.register_parameter("bias", None)
        self.num_heads = total_num_heads
        self.num_kv_heads = total_num_kv_heads


class _FakeH3Attention(nn.Module):
    def __init__(self, **kwargs: object) -> None:
        del kwargs
        super().__init__()


class _H3Pipeline(nn.Module):
    def __init__(self, transformer: nn.Module) -> None:
        super().__init__()
        self.transformer = transformer


def _small_h3_config() -> SimpleNamespace:
    return SimpleNamespace(
        tf_model_config={
            "num_layers": 1,
            "token_refiner_num_layers": 1,
            "hidden_size": 8,
            "num_attention_heads": 2,
            "attention_head_dim": 4,
            "ffn_hidden_size": 16,
            "latents_dim": 2,
            "audio_latents_dim": 2,
            "patch_size": (1, 2, 2),
            "text_dim": 6,
            "timestep_input_dim": 4,
            "time_embed_hidden_size": 8,
            "time_embed_dim": 4,
            "adaln_out_features": 18 * 8,
            "final_adaln_out_features": 2 * 8,
            "rope_inv_freq_len": 2,
        },
        parallel_config=SimpleNamespace(ulysses_degree=1),
    )


def _prepared_source(
    tmp_path: Path,
    *,
    content: bytes = b"canonical-source-for-identity",
    requested_revision: str | None = None,
    resolved_revision: str | None = None,
    directory: str = "canonical",
    prefix: str = "transformer.",
) -> PreparedWeightSource:
    source_root = tmp_path / directory
    if resolved_revision is not None:
        source_root = source_root / "snapshots" / resolved_revision
    source_root.mkdir(parents=True, exist_ok=True)
    weight_file = source_root / "model.safetensors"
    if not weight_file.exists():
        weight_file.write_bytes(content)
    elif weight_file.read_bytes() != content:
        weight_file.write_bytes(content)
    return PreparedWeightSource(
        model_or_path="test-org/tiny-diffusion",
        subfolder=None,
        requested_revision=requested_revision,
        prefix=prefix,
        resolved_root=source_root,
        weight_files=(weight_file,),
        use_safetensors=True,
    )


def _request(**changes: object) -> FinalLayoutRequest:
    request = FinalLayoutRequest(
        model_id="test-org/tiny-diffusion",
        loader=FinalLayoutLoaderIdentity(
            implementation=ImplementationIdentity(
                implementation_id="test-diffusion-loader",
                version="1",
                fingerprint="test-loader-implementation-v1",
            ),
            model_config_fingerprint="test-model-config-v1",
            weight_transform_fingerprint="test-weight-transform-v1",
        ),
    )
    return dataclasses.replace(request, **changes)


def _identity(
    model: nn.Module,
    source: PreparedWeightSource,
    *,
    request: FinalLayoutRequest | None = None,
    policy: FinalLayoutTensorPolicy = FINAL_LAYOUT_BF16_POLICY,
) -> FinalLayoutIdentityContext:
    transformer = model.get_submodule("transformer")
    return build_final_layout_identity(
        model,
        dit_modules=(("transformer", transformer),),
        prepared_sources=(source,),
        request=request or _request(),
        policy=policy,
    )


def _multi_identity(
    model: _MultiTinyPipeline,
    sources: tuple[PreparedWeightSource, ...],
) -> FinalLayoutIdentityContext:
    return build_final_layout_identity(
        model,
        dit_modules=(
            ("transformer_a", model.transformer_a),
            ("transformer_b", model.transformer_b),
        ),
        prepared_sources=sources,
        request=_request(),
        policy=FINAL_LAYOUT_BF16_POLICY,
    )


def _runtime(root: Path) -> HostWeightRuntime:
    return HostWeightRuntime.from_config(
        HostWeightRuntimeConfig(
            mode=RuntimeMode.PREFERRED,
            domain=StorageDomainPolicy(
                root=root,
                scope=StorageScope.NODE,
                domain_id="node",
                storage_class=detect_storage_class(root.parent),
            ),
            production=ProductionPolicy(
                allow_local_build=False,
                allow_post_load_publish=True,
            ),
            wait=WaitPolicy(coordination_timeout_seconds=5.0),
        )
    )


def _dit_modules(model: _TinyPipeline) -> tuple[tuple[str, nn.Module], ...]:
    return (("transformer", model.transformer),)


def _fill_final_weights(model: _TinyPipeline) -> None:
    with torch.no_grad():
        model.transformer.proj.weight.copy_(torch.arange(6, dtype=torch.float32).to(torch.bfloat16).reshape(2, 3))
        model.transformer.proj.bias.copy_(torch.tensor([3.0, 4.0], dtype=torch.bfloat16))
        model.transformer.fp32_gain.copy_(torch.tensor([6.5]))
        model.transformer.scale.copy_(torch.tensor([2.5]))
        model.text_encoder.weight.fill_(11)
        model.text_encoder.bias.fill_(12)


def _pointer_snapshot(model: nn.Module) -> dict[str, tuple[int, int, str]]:
    snapshot: dict[str, tuple[int, int, str]] = {}
    for name, tensor in list(model.named_parameters()) + list(model.named_buffers()):
        pointer = 0 if tensor.is_meta else tensor.untyped_storage().data_ptr()
        snapshot[name] = (id(tensor), pointer, tensor.device.type)
    return snapshot


def test_preload_identity_is_stable_for_cpu_and_meta_skeletons(tmp_path: Path) -> None:
    source = _prepared_source(tmp_path)
    cpu_model = _TinyPipeline()
    meta_model = _TinyPipeline(dit_device="meta")

    cpu = _identity(cpu_model, source)
    meta = _identity(meta_model, source)

    assert cpu.identity == meta.identity
    assert cpu.tensor_names == {
        "transformer.fp32_gain",
        "transformer.proj.bias",
        "transformer.proj.weight",
        "transformer.scale",
    }
    assert "transformer.derived" not in cpu.tensor_names
    assert all("text_encoder" not in name and "vae" not in name for name in cpu.tensor_names)

    request_fields = {field.name for field in dataclasses.fields(FinalLayoutRequest)}
    assert not any("data_parallel" in name for name in request_fields)
    assert not any("metadata" in name for name in request_fields)
    identity_bytes = cpu.identity.canonical_bytes
    assert b"data_parallel" not in identity_bytes
    assert b"dlo_use_allgather" not in identity_bytes
    assert b"registration" not in identity_bytes
    with pytest.raises(TypeError, match="unexpected keyword argument"):
        dataclasses.replace(_request(), loader_metadata=CanonicalJson.empty())


def test_generic_identity_and_tensor_discovery_accept_a_second_policy(tmp_path: Path) -> None:
    model = _Float32Pipeline()
    source = _prepared_source(tmp_path, directory="fp32-source")

    context = _identity(model, source, policy=_SYNTHETIC_FP32_POLICY)
    targets = collect_final_layout_targets(
        model,
        (("transformer", model.transformer),),
        policy=_SYNTHETIC_FP32_POLICY,
        require_materialized=True,
    )

    assert context.spec is _SYNTHETIC_FP32_SPEC
    assert context.identity.representation.name == "test-final-layout-fp32"
    assert context.identity.representation.dtype == str(torch.float32)
    assert all(target.tensor.dtype is torch.float32 for target in targets)
    assert context.identity.producer.implementation_fingerprint != (
        FINAL_LAYOUT_BF16_SPEC.producer.implementation_fingerprint
    )


def test_source_identity_requires_complete_multi_component_coverage(tmp_path: Path) -> None:
    model = _MultiTinyPipeline()
    source_a = _prepared_source(
        tmp_path,
        directory="source-a",
        prefix="transformer_a.",
    )

    with pytest.raises(FinalLayoutContractError, match="transformer_b") as exc_info:
        _multi_identity(model, (source_a,))
    assert exc_info.value.code is FinalLayoutContractCode.SOURCE_COVERAGE_INVALID

    source_b = _prepared_source(
        tmp_path,
        directory="source-b",
        prefix="transformer_b.",
    )
    context = _multi_identity(model, (source_a, source_b))
    metadata = context.identity.source.metadata.to_value()
    assert isinstance(metadata, dict)
    bindings = metadata["target_bindings"]
    assert isinstance(bindings, list)
    assert {binding["source_prefix"] for binding in bindings} == {
        "transformer_a.",
        "transformer_b.",
    }
    assert {binding["target_name"].split(".", 1)[0] for binding in bindings} == {
        "transformer_a",
        "transformer_b",
    }


def test_source_identity_resolves_longest_prefix_and_rejects_ties(tmp_path: Path) -> None:
    model = _MultiTinyPipeline()
    root_source = _prepared_source(
        tmp_path,
        directory="root-source",
        prefix="",
    )
    source_a = _prepared_source(
        tmp_path,
        directory="specific-a",
        prefix="transformer_a.",
    )

    context = _multi_identity(model, (root_source, source_a))
    reordered = _multi_identity(model, (source_a, root_source))
    assert context.identity == reordered.identity
    metadata = context.identity.source.metadata.to_value()
    assert isinstance(metadata, dict)
    bindings = metadata["target_bindings"]
    assert isinstance(bindings, list)
    by_target = {binding["target_name"]: binding["source_prefix"] for binding in bindings}
    assert all(prefix == "transformer_a." for name, prefix in by_target.items() if name.startswith("transformer_a."))
    assert all(prefix == "" for name, prefix in by_target.items() if name.startswith("transformer_b."))

    conflicting_a = _prepared_source(
        tmp_path,
        directory="conflicting-a",
        prefix="transformer_a.",
    )
    with pytest.raises(FinalLayoutContractError, match="equally specific") as exc_info:
        _multi_identity(model, (root_source, source_a, conflicting_a))
    assert exc_info.value.code is FinalLayoutContractCode.SOURCE_COVERAGE_INVALID


def test_local_hex_named_symlink_target_is_content_hashed_across_startups(tmp_path: Path) -> None:
    model = _TinyPipeline()
    source_root = tmp_path / "local-symlink-source"
    blobs_root = tmp_path / "local-blobs"
    source_root.mkdir()
    blobs_root.mkdir()
    blob_path = blobs_root / ("a" * 64)
    blob_path.write_bytes(b"a" * 32)
    weight_path = source_root / "model.safetensors"
    weight_path.symlink_to(blob_path)
    source = PreparedWeightSource(
        model_or_path=str(source_root),
        subfolder=None,
        requested_revision=None,
        prefix="transformer.",
        resolved_root=source_root,
        weight_files=(weight_path,),
        use_safetensors=True,
    )

    first = _identity(model, source)
    first_metadata = first.identity.source.metadata.to_value()
    assert isinstance(first_metadata, dict)
    first_sources = first_metadata["sources"]
    assert isinstance(first_sources, list)
    assert first_sources[0]["files"][0]["content_id"].startswith("sha256:")

    blob_path.write_bytes(b"b" * 32)
    second = _identity(model, source)

    assert second.identity != first.identity


def test_hf_blob_shortcut_requires_validated_snapshot_topology(tmp_path: Path) -> None:
    model = _TinyPipeline()
    model_id = "test-org/tiny-diffusion"
    repo_root = tmp_path / "models--test-org--tiny-diffusion"
    blobs_root = repo_root / "blobs"
    revision = "1" * 40
    snapshot_root = repo_root / "snapshots" / revision
    blobs_root.mkdir(parents=True)
    snapshot_root.mkdir(parents=True)
    blob_name = "2" * 64
    blob_path = blobs_root / blob_name
    blob_path.write_bytes(b"trusted-hf-blob")
    weight_path = snapshot_root / "model.safetensors"
    weight_path.symlink_to(Path("..") / ".." / "blobs" / blob_name)
    source = PreparedWeightSource(
        model_or_path=model_id,
        subfolder=None,
        requested_revision="main",
        prefix="transformer.",
        resolved_root=snapshot_root,
        weight_files=(weight_path,),
        use_safetensors=True,
        source_kind=WeightSourceKind.HUGGING_FACE_HUB,
    )

    context = _identity(model, source)
    metadata = context.identity.source.metadata.to_value()
    assert isinstance(metadata, dict)
    sources = metadata["sources"]
    assert isinstance(sources, list)
    assert sources[0]["files"][0]["content_id"] == f"immutable-blob:{blob_name}"


def test_identity_uses_resolved_revision_and_exact_semantics(tmp_path: Path) -> None:
    model = _TinyPipeline()
    commit = "0123456789abcdef0123456789abcdef01234567"
    source = _prepared_source(tmp_path, requested_revision="main", resolved_revision=commit)
    by_alias = _identity(model, source)
    by_commit = _identity(model, dataclasses.replace(source, requested_revision=commit))

    assert by_alias.identity == by_commit.identity

    base = _request()
    variants = (
        dataclasses.replace(
            base,
            loader=dataclasses.replace(base.loader, model_config_fingerprint="test-model-config-v2"),
        ),
        dataclasses.replace(
            base,
            loader=dataclasses.replace(
                base.loader,
                implementation=dataclasses.replace(base.loader.implementation, version="2"),
            ),
        ),
        dataclasses.replace(
            base,
            parallel=FinalLayoutParallelIdentity(tensor_parallel_size=2, tensor_parallel_rank=0),
        ),
        dataclasses.replace(
            base,
            parallel=FinalLayoutParallelIdentity(tensor_parallel_size=2, tensor_parallel_rank=1),
        ),
        dataclasses.replace(
            base,
            parallel=FinalLayoutParallelIdentity(
                sequence_parallel_size=2,
                ulysses_degree=2,
                ulysses_mode="strict",
            ),
        ),
    )
    base_identity = _identity(model, source, request=base).identity
    assert all(_identity(model, source, request=variant).identity != base_identity for variant in variants)

    model_v2 = _TinyPipeline()
    model_v2.transformer.host_weight_restore_contract = dataclasses.replace(
        model_v2.transformer.host_weight_restore_contract,
        version="2",
    )
    assert _identity(model_v2, source).identity != base_identity

    changed_source = _prepared_source(
        tmp_path,
        directory="changed",
        content=b"different-canonical-source",
    )
    assert _identity(model, changed_source).identity != base_identity
    source_with_adapter = dataclasses.replace(
        source,
        checkpoint_adapter=ImplementationIdentity(
            implementation_id="test-checkpoint-adapter",
            version="1",
            fingerprint="test-checkpoint-adapter-v1",
        ),
    )
    assert _identity(model, source_with_adapter).identity != base_identity

    abi_v2 = CanonicalJson.from_value(
        {
            "artifact_identity": "diffusion-final-layout-identity-v1",
            "model_contract": FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA,
            "producer": "finalized-bf16-tensor-writer-v1",
            "representation_policy": "bf16-with-preserved-fp32-v1",
            "restorer": "exact-final-layout-tensor-rebind-v2",
            "source_identity": "prepared-diffusion-weight-source-v2",
            "tensor_contract": "complete-strided-tensor-ownership-v1",
        }
    )
    with pytest.raises(ValueError, match="fingerprint differs"):
        dataclasses.replace(FINAL_LAYOUT_BF16_SPEC, implementation_abi=abi_v2)
    spec_v2 = dataclasses.replace(
        FINAL_LAYOUT_BF16_SPEC,
        producer=dataclasses.replace(
            FINAL_LAYOUT_BF16_SPEC.producer,
            version="2",
            implementation_fingerprint=implementation_abi_fingerprint(abi_v2),
        ),
        implementation_abi=abi_v2,
    )
    identity_v2 = _identity(
        model,
        source,
        policy=_PolicyWithSpec(FINAL_LAYOUT_BF16_POLICY, spec_v2),
    ).identity
    assert identity_v2.representation == base_identity.representation
    assert identity_v2.key != base_identity.key


@pytest.mark.parametrize(
    ("semantic_request", "message"),
    [
        (_request(load_format="diffusers"), "load_format='default'"),
        (
            _request(parallel=FinalLayoutParallelIdentity(pipeline_parallel_size=2)),
            "pipeline-parallel",
        ),
        (
            _request(parallel=FinalLayoutParallelIdentity(cfg_parallel_size=2)),
            "CFG-parallel",
        ),
        (_request(parallel=FinalLayoutParallelIdentity(use_hsdp=True)), "HSDP"),
        (
            _request(parallel=FinalLayoutParallelIdentity(enable_expert_parallel=True)),
            "expert-parallel",
        ),
        (
            _request(adaptation=AdaptationIdentity(kind="merged-lora", fingerprint="adapter-sha256")),
            "LoRA",
        ),
    ],
)
def test_bf16_policy_fails_closed_for_unsupported_semantics(
    tmp_path: Path,
    semantic_request: FinalLayoutRequest,
    message: str,
) -> None:
    model = _TinyPipeline()
    with pytest.raises(ValueError, match=message):
        _identity(model, _prepared_source(tmp_path), request=semantic_request)


def test_tensor_ownership_is_complete_mixed_precision_and_alias_free(tmp_path: Path) -> None:
    model = _TinyPipeline()
    records = collect_final_layout_targets(
        model,
        _dit_modules(model),
        policy=FINAL_LAYOUT_BF16_POLICY,
        require_materialized=True,
    )
    by_name = {record.name: record for record in records}

    assert set(by_name) == {
        "transformer.fp32_gain",
        "transformer.proj.bias",
        "transformer.proj.weight",
        "transformer.scale",
    }
    assert by_name["transformer.proj.weight"].tensor.dtype is torch.bfloat16
    assert by_name["transformer.fp32_gain"].tensor.dtype is torch.float32
    assert by_name["transformer.scale"].role == "persistent_buffer"

    model.transformer.register_parameter("alias", model.transformer.proj.weight)
    with pytest.raises(FinalLayoutContractError, match="aliases tensor object"):
        _identity(model, _prepared_source(tmp_path))


def test_identity_requires_explicit_model_restore_contract(tmp_path: Path) -> None:
    model = _TinyPipeline()
    model.transformer.host_weight_restore_contract = "unsupported-contract"

    with pytest.raises(FinalLayoutContractError, match="does not declare"):
        _identity(model, _prepared_source(tmp_path))


def test_publication_is_warm_only_and_restore_is_transactional(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    source = _prepared_source(tmp_path)
    cold_model = _TinyPipeline()
    _fill_final_weights(cold_model)
    context = _identity(cold_model, source)
    runtime = _runtime(tmp_path / "store")
    producer = FinalLayoutBF16Producer(context, cold_model, _dit_modules(cold_model), max_shard_bytes=16)
    assert producer.spec.lookup_phase is LookupPhase.POST_LOAD_ONLY
    assert producer.spec.source_mode is ProductionSourceMode.FINALIZED_MODEL
    assert producer.spec.coordination_scope is CoordinationScope.SINGLE_PROCESS

    miss = runtime.resolve(context.identity)
    assert miss.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
    cold_pointers = _pointer_snapshot(cold_model)
    cold_values = {name: tensor.detach().clone() for name, tensor in cold_model.named_parameters()}

    publication = runtime.publish_after_load(context.identity, producer=producer)

    assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED
    assert publication.failure is None
    assert _pointer_snapshot(cold_model) == cold_pointers
    assert all(torch.equal(dict(cold_model.named_parameters())[name], value) for name, value in cold_values.items())
    assert cold_model.transformer.restore_validations == 1

    def unexpected_produce(_writer: object) -> object:
        raise AssertionError("already-present publication must not invoke the producer")

    monkeypatch.setattr(producer, "produce", unexpected_produce)
    already_present = runtime.publish_after_load(context.identity, producer=producer)
    assert already_present.outcome is PostLoadPublicationOutcome.ALREADY_PRESENT

    warm_model = _TinyPipeline(dit_device="meta")
    warm_context = _identity(warm_model, source)
    assert warm_context.identity == context.identity
    hit = runtime.resolve(warm_context.identity)
    assert hit.report.outcome is ResolutionOutcome.LOCAL_HIT
    assert hit.lease is not None

    before_plan = _pointer_snapshot(warm_model)
    text_weight = warm_model.text_encoder.weight
    vae_gain = warm_model.vae.gain
    plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease)
    assert _pointer_snapshot(warm_model) == before_plan
    assert warm_model.transformer.restore_validations == 0

    plan.commit()
    with pytest.raises(RuntimeError, match="already committed"):
        plan.commit()

    expected = {
        "transformer.proj.weight": torch.arange(6, dtype=torch.float32).to(torch.bfloat16).reshape(2, 3),
        "transformer.proj.bias": torch.tensor([3.0, 4.0], dtype=torch.bfloat16),
        "transformer.fp32_gain": torch.tensor([6.5]),
        "transformer.scale": torch.tensor([2.5]),
    }
    restored = dict(warm_model.named_parameters()) | dict(warm_model.named_buffers())
    assert all(torch.equal(restored[name], value) for name, value in expected.items())
    assert all(not restored[name].is_meta for name in warm_context.tensor_names)
    assert warm_model.transformer.restore_validations == 1
    assert warm_model.text_encoder.weight is text_weight
    assert warm_model.vae.gain is vae_gain
    assert warm_model.transformer.derived.item() == 7.0
    for name in warm_context.tensor_names:
        assert restored[name].untyped_storage().data_ptr() == hit.lease.tensors[name].untyped_storage().data_ptr()

    del restored, warm_model, plan
    gc.collect()
    hit.lease.close()
    assert isinstance(runtime.store, FilesystemHostWeightStore)
    assert runtime.store.cleanup(context.identity) is None


def test_independent_publications_are_content_deterministic(tmp_path: Path) -> None:
    source = _prepared_source(tmp_path)
    model = _TinyPipeline()
    _fill_final_weights(model)
    context = _identity(model, source)
    first_runtime = _runtime(tmp_path / "first-store")
    second_runtime = _runtime(tmp_path / "second-store")

    first = first_runtime.publish_after_load(
        context.identity,
        producer=FinalLayoutBF16Producer(context, model, _dit_modules(model), max_shard_bytes=16),
    )
    second = second_runtime.publish_after_load(
        context.identity,
        producer=FinalLayoutBF16Producer(context, model, _dit_modules(model), max_shard_bytes=1024),
    )
    assert first.outcome is PostLoadPublicationOutcome.PUBLISHED
    assert second.outcome is PostLoadPublicationOutcome.PUBLISHED

    first_hit = first_runtime.resolve(context.identity)
    second_hit = second_runtime.resolve(context.identity)
    assert first_hit.lease is not None and second_hit.lease is not None
    assert first_hit.lease.manifest.artifact_content_sha256 == second_hit.lease.manifest.artifact_content_sha256
    assert first_hit.lease.manifest.format_metadata == second_hit.lease.manifest.format_metadata
    assert {entry.name: entry.sha256 for entry in first_hit.lease.manifest.tensors} == {
        entry.name: entry.sha256 for entry in second_hit.lease.manifest.tensors
    }
    first_hit.lease.close()
    second_hit.lease.close()


def test_source_replacement_prevents_publication(tmp_path: Path) -> None:
    source = _prepared_source(tmp_path)
    model = _TinyPipeline()
    _fill_final_weights(model)
    context = _identity(model, source)
    runtime = _runtime(tmp_path / "store")
    source.weight_files[0].write_bytes(b"replacement-source")

    publication = runtime.publish_after_load(
        context.identity,
        producer=FinalLayoutBF16Producer(context, model, _dit_modules(model)),
    )

    assert publication.outcome is PostLoadPublicationOutcome.FAILED
    assert publication.failure is not None
    assert publication.failure.stage is ResolutionStage.CANONICAL_LOADING
    assert publication.failure.code is FailureCode.CANONICAL_SOURCE_FAILED
    assert publication.failure.details.to_value() == {"final_layout_code": "source_changed"}
    assert not context.sources_unchanged()
    assert runtime.resolve(context.identity).report.outcome is ResolutionOutcome.CANONICAL_FALLBACK


def test_source_replacement_between_restore_plan_and_commit_causes_no_mutation(tmp_path: Path) -> None:
    source = _prepared_source(tmp_path)
    cold_model = _TinyPipeline()
    _fill_final_weights(cold_model)
    context = _identity(cold_model, source)
    runtime = _runtime(tmp_path / "store")
    assert (
        runtime.publish_after_load(
            context.identity,
            producer=FinalLayoutBF16Producer(context, cold_model, _dit_modules(cold_model)),
        ).outcome
        is PostLoadPublicationOutcome.PUBLISHED
    )
    hit = runtime.resolve(context.identity)
    assert hit.lease is not None

    warm_model = _TinyPipeline()
    warm_context = _identity(warm_model, source)
    plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease)
    before = _pointer_snapshot(warm_model)
    source.weight_files[0].write_bytes(b"source-changed-before-commit")

    with pytest.raises(FinalLayoutContractError, match="canonical source changed"):
        plan.commit()
    assert _pointer_snapshot(warm_model) == before
    del plan
    gc.collect()
    hit.lease.close()


def test_restore_rejects_wrong_identity_and_coverage_without_mutation(tmp_path: Path) -> None:
    source = _prepared_source(tmp_path)
    cold_model = _TinyPipeline()
    _fill_final_weights(cold_model)
    context = _identity(cold_model, source)
    runtime = _runtime(tmp_path / "store")
    assert (
        runtime.publish_after_load(
            context.identity,
            producer=FinalLayoutBF16Producer(context, cold_model, _dit_modules(cold_model)),
        ).outcome
        is PostLoadPublicationOutcome.PUBLISHED
    )
    hit = runtime.resolve(context.identity)
    assert hit.lease is not None

    alternate = _AlternateTinyPipeline()
    _fill_final_weights(alternate)
    with pytest.raises(ValueError, match="producer model implementation differs"):
        FinalLayoutBF16Producer(context, alternate, _dit_modules(alternate))
    with pytest.raises(ValueError, match="restore model implementation differs"):
        FinalLayoutTensorRestorer(context).plan_restore(alternate, hit.lease)

    model = _TinyPipeline()
    wrong_context = _identity(
        model,
        source,
        request=_request(
            parallel=FinalLayoutParallelIdentity(
                tensor_parallel_size=2,
                tensor_parallel_rank=1,
            )
        ),
    )
    with pytest.raises(ValueError, match="semantic identity differs"):
        FinalLayoutTensorRestorer(wrong_context).plan_restore(model, hit.lease)

    exact_context = _identity(model, source)
    before = _pointer_snapshot(model)
    model.transformer.register_buffer("new_persistent_state", torch.tensor([1.0]))
    with pytest.raises(FinalLayoutContractError, match="tensor contract differs"):
        FinalLayoutTensorRestorer(exact_context).plan_restore(model, hit.lease)
    after = _pointer_snapshot(model)
    assert all(after[name] == value for name, value in before.items())
    hit.lease.close()


def test_reduced_minimax_h3_satisfies_real_tensor_ownership_contract(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3

    monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeH3Linear)
    monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeH3Linear)
    monkeypatch.setattr(h3, "QKVParallelLinear", _FakeH3Linear)
    monkeypatch.setattr(h3, "RowParallelLinear", _FakeH3Linear)
    monkeypatch.setattr(h3, "Attention", _FakeH3Attention)
    monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1)

    transformer = h3.MiniMaxH3DiTModel(_small_h3_config(), quant_config=None)
    pipeline = _H3Pipeline(transformer)
    source = _prepared_source(tmp_path, directory="minimax-h3-source")
    context = _identity(pipeline, source)
    targets = collect_final_layout_targets(
        pipeline,
        (("transformer", transformer),),
        policy=FINAL_LAYOUT_BF16_POLICY,
        require_materialized=True,
    )
    by_name = {target.name: target for target in targets}

    assert context.model_contracts == (MiniMaxH3DiTModel.host_weight_restore_contract,)
    assert "transformer.video_patch_proj.weight" in by_name
    assert "transformer.blocks.0.attn.qkv_proj.weight" in by_name
    assert "transformer.rope.inv_freq" in by_name
    assert by_name["transformer.video_patch_proj.weight"].tensor.dtype is torch.float32
    assert by_name["transformer.blocks.0.attn.qkv_proj.weight"].tensor.dtype is torch.bfloat16
    assert by_name["transformer.rope.inv_freq"].role == "persistent_buffer"
    transformer.validate_restored_host_weights()


def test_minimax_h3_declares_the_final_layout_restore_contract() -> None:
    contract = MiniMaxH3DiTModel.host_weight_restore_contract
    assert isinstance(contract, FinalLayoutModelContract)
    assert contract.schema == FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA
    assert contract.implementation_id == "minimax-h3-dit"
    assert callable(MiniMaxH3DiTModel.validate_restored_host_weights)
    assert "data_parallel" not in FinalLayoutRequest.__dataclass_fields__