"""Frozen source-specific reference-storage contract; no native FP4 runtime."""
import math
from pathlib import PurePosixPath
from typing import Literal
from pydantic import BaseModel,ConfigDict,Field,model_validator

class Strict(BaseModel):
    model_config=ConfigDict(extra='forbid',frozen=True)

class FileDigest(Strict):
    sha256:str=Field(pattern=r'^[0-9a-f]{64}$')
    size_bytes:int=Field(ge=0)

class Tensor(Strict):
    shape:tuple[int,...]
    dtype:Literal['bfloat16','float32','float16','int64','int32','bool']
    file:str
    method:Literal['preserve','nvfp4']
    global_scale:float|None=Field(default=None,gt=0,allow_inf_nan=False)
    group_size:int|None=None

    @model_validator(mode='after')
    def layout(self):
        if any(x<0 for x in self.shape):raise ValueError('Negative shape')
        if self.method=='preserve':
            if self.global_scale is not None or self.group_size is not None:raise ValueError('Protected tensor has quantization fields')
        elif len(self.shape)!=2 or min(self.shape)<1 or self.shape[1]%16 or self.group_size!=16 or self.global_scale is None or self.dtype!='bfloat16':
            raise ValueError('Invalid NVFP4 BF16-source matrix')
        relative=PurePosixPath(self.file)
        if relative.is_absolute() or '..' in relative.parts or '\\' in self.file:raise ValueError('Unsafe tensor file')
        return self

class Component(Strict):
    class_name:Literal['QwenImage21Transformer2DModel','MiniMaxH3Transformer3DModel']
    tensors:dict[str,Tensor]
    quantized_matrices:int=Field(gt=0)
    preserved_tensors:int=Field(ge=0)

    @model_validator(mode='after')
    def inventory(self):
        selected=sum(x.method=='nvfp4' for x in self.tensors.values())
        if selected!=self.quantized_matrices or len(self.tensors)-selected!=self.preserved_tensors:raise ValueError('Inventory counts mismatch')
        return self

class DiffusionManifest(Strict):
    schema_version:Literal['axq-diffusion.reference/1']='axq-diffusion.reference/1'
    format:Literal['nvfp4']='nvfp4'
    status:Literal['development']='development'
    source_model:Literal['Qwen/Qwen-Image-2.1-Turbo','MiniMaxAI/MiniMax-H3']
    source_revision:str=Field(pattern=r'^[0-9a-f]{40}$')
    backend:Literal['numpy-reference']='numpy-reference'
    execution:Literal['reconstruct-bfloat16-diffusers']='reconstruct-bfloat16-diffusers'
    native_low_precision_verified:Literal[False]=False
    quality_certified:Literal[False]=False
    source_files:dict[str,FileDigest]
    components:dict[str,Component]
    files:dict[str,FileDigest]

    @model_validator(mode='after')
    def binding(self):
        source={
            'Qwen/Qwen-Image-2.1-Turbo':('fc13ba1986fcb905854dc8de2f52eee395f49b5b',{'transformer'},'QwenImage21Transformer2DModel'),
            'MiniMaxAI/MiniMax-H3':('42ed227ee7df40d41602854ae760620d6eb651fe',{'transformer','transformer_ref'},'MiniMaxH3Transformer3DModel'),
        }
        revision,components,class_name=source[self.source_model]
        if self.source_revision!=revision or set(self.components)!=components:raise ValueError('Source pin/component identity mismatch')
        if not self.files or not self.source_files:raise ValueError('Missing bound files')
        for name in [*self.files,*self.source_files]:
            relative=PurePosixPath(name)
            if relative.is_absolute() or '..' in relative.parts or '\\' in name:raise ValueError('Unsafe manifest member')
        for name,spec in self.components.items():
            if spec.class_name!=class_name:raise ValueError('Wrong component class')
            for tensor in spec.tensors.values():
                if tensor.file not in self.files or not tensor.file.startswith(name+'/'):raise ValueError('Unbound tensor member')
        return self