"""Strict native-component loader; supply the official decoder source explicitly."""
import hashlib,importlib.util,json,math
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
import torch
from torch import nn
from accelerate import init_empty_weights
from safetensors.torch import load_file
from diffusers import AutoencoderDC,SanaTransformer2DModel
from transformers import Qwen2Config,Qwen2Model,InternVLModel
from native.configuration_intern_vit import InternVisionConfig
from native.modeling_intern_vit import InternVisionModel

class VisualComponents(nn.Module):
    def __init__(self,root,decoder_source):
        super().__init__()
        root=Path(root)
        source=Path(decoder_source)
        expected=json.loads((root/'external_decoder_source.json').read_text())
        if hashlib.sha256(source.read_bytes()).hexdigest()!=expected['sha256']:
            raise ValueError('Decoder source differs from the pinned official source')
        spec=importlib.util.spec_from_file_location('_unilip_external_vae',source)
        external=importlib.util.module_from_spec(spec);spec.loader.exec_module(external)
        cfg=json.loads((root/'source_config.json').read_text())
        self.unilip_factor=cfg['unilip_factor']
        self.model=nn.Module()
        m=self.model
        m.vision_tower=InternVisionModel(InternVisionConfig(**json.loads((root/'intern_vision_config.json').read_text())))
        m.multi_modal_projector=nn.Sequential(nn.LayerNorm(4096),nn.Linear(4096,896),nn.GELU(),nn.Linear(896,896))
        m.latent_queries=nn.Parameter(torch.empty(1,cfg['n_query'],896))
        m.projector=nn.Linear(896,2304)
        connector_config=Qwen2Config(**cfg['text_config'])
        connector_config._attn_implementation='eager'
        m.llm_connector=Qwen2Model(connector_config)
        del m.llm_connector.layers[:-cfg['connect_layer']]
        del m.llm_connector.embed_tokens
        dc_cfg=json.loads((root/'dc_ae_config.json').read_text())
        # Build only configuration shapes; checkpoint supplies all actual parameters.
        with patch.object(AutoencoderDC,'from_pretrained',side_effect=lambda *a,**k:AutoencoderDC.from_config(dc_cfg)):
            m.vae_decoder=external.DCAE_Decoder(SimpleNamespace(model=SimpleNamespace(dc_ae_path='configuration-only')),896)
        m.dit=SanaTransformer2DModel.from_config(json.loads((root/'dit_config.json').read_text()))

    def image_features(self,pixels):
        v=self.model.vision_tower(pixels).last_hidden_state[:,1:]
        b,l,c=v.shape;s=math.isqrt(l)
        if s*s!=l or s%2:
            raise ValueError('Input must yield a square even-sized patch grid')
        v=v.reshape(b,s,s,c)
        # Use the licensed Transformers native space-to-depth ordering.
        v=InternVLModel.pixel_shuffle(self,v,scale_factor=0.5)
        return self.model.multi_modal_projector(v.reshape(b,-1,4*c))

    def encode(self,pixels):
        return self.model.vae_decoder.clip_down(self.image_features(pixels))

    def decode(self,latent):
        """Unscaled DC-AE latent. Divide SANA flow latents by unilip_factor first."""
        return self.model.vae_decoder.vae_decode(latent)

    def condition(self,language_hidden_states,attention_mask=None):
        b,n,_=language_hidden_states.shape
        if attention_mask is None:
            attention_mask=torch.ones(b,n,device=language_hidden_states.device,dtype=torch.bool)
        full=attention_mask[:,None,:,None].bool() & attention_mask[:,None,None,:].bool()
        additive=(1-full.to(language_hidden_states.dtype))*(-100000.0)
        out=self.model.llm_connector(inputs_embeds=language_hidden_states,attention_mask=additive,
            position_ids=(attention_mask.long().cumsum(-1)-1).clamp_min(0),use_cache=False,return_dict=True)
        return self.model.projector(out.last_hidden_state)

def load_model(root,decoder_source,device='cpu',dtype=torch.float32):
    with init_empty_weights(include_buffers=False):
        m=VisualComponents(root,decoder_source)
    m.load_state_dict(load_file(str(Path(root)/'model.safetensors')),strict=True,assign=True)
    return m.to(device=device,dtype=dtype).eval()