"""Offline initialization helpers for DiffusionVLA backbones.

Diffusion-action checkpoints contain the trained vision, projector, LLM, and action decoder weights, but upstream model
construction still tries to download TIMM DINO/SigLIP pretrained weights and the LLaMA config/tokenizer before loading the
checkpoint. This patch prevents network access during construction and lets the local checkpoint provide the actual
weights.
"""
from __future__ import annotations

import os
from pathlib import Path

_MODEL_DIR = Path(__file__).resolve().parent
_TOKENIZER_DIR = Path(os.environ.get('VLA_TOKENIZER_DIR', _MODEL_DIR))
_LLAMA_CONFIG_PATH = Path(os.environ.get('VLA_LLAMA_CONFIG_PATH', _MODEL_DIR / 'llama2_7b_config.json'))
_LLAMA2_7B_ID = 'meta-llama/Llama-2-7b-hf'


def patch_offline_backbone_loading() -> None:
    """Patch TIMM and Transformers loaders so the diffusion-action policy can instantiate fully offline."""
    import timm
    from transformers import AutoConfig, AutoTokenizer, LlamaConfig, LlamaForCausalLM

    if not getattr(timm.create_model, '_local_vla_offline_patch', False):
        original_create_model = timm.create_model

        def create_model_offline(model_name, *args, **kwargs):
            # The uploaded diffusion-action checkpoint overwrites these weights immediately after construction.
            if kwargs.get('pretrained', False) and ('dinov2' in str(model_name) or 'siglip' in str(model_name)):
                kwargs = dict(kwargs)
                kwargs['pretrained'] = False
            return original_create_model(model_name, *args, **kwargs)

        create_model_offline._local_vla_offline_patch = True
        timm.create_model = create_model_offline

    def local_llama_config() -> LlamaConfig:
        if _LLAMA_CONFIG_PATH.exists():
            return LlamaConfig.from_json_file(str(_LLAMA_CONFIG_PATH))
        return LlamaConfig(
            vocab_size=32000,
            hidden_size=4096,
            intermediate_size=11008,
            num_hidden_layers=32,
            num_attention_heads=32,
            num_key_value_heads=32,
            max_position_embeddings=4096,
            rms_norm_eps=1e-5,
            rope_theta=10000.0,
            pad_token_id=0,
            bos_token_id=1,
            eos_token_id=2,
            tie_word_embeddings=False,
        )

    if not getattr(AutoConfig.from_pretrained, '_local_vla_offline_patch', False):
        original_config_from_pretrained = AutoConfig.from_pretrained

        def config_from_pretrained_offline(model_name_or_path, *args, **kwargs):
            if str(model_name_or_path) == _LLAMA2_7B_ID:
                return local_llama_config()
            return original_config_from_pretrained(model_name_or_path, *args, **kwargs)

        config_from_pretrained_offline._local_vla_offline_patch = True
        AutoConfig.from_pretrained = config_from_pretrained_offline

    if not getattr(LlamaConfig.from_pretrained, '_local_vla_offline_patch', False):
        original_llama_config_from_pretrained = LlamaConfig.from_pretrained

        def llama_config_from_pretrained_offline(model_name_or_path, *args, **kwargs):
            if str(model_name_or_path) == _LLAMA2_7B_ID:
                return local_llama_config()
            return original_llama_config_from_pretrained(model_name_or_path, *args, **kwargs)

        llama_config_from_pretrained_offline._local_vla_offline_patch = True
        LlamaConfig.from_pretrained = llama_config_from_pretrained_offline

    if not getattr(LlamaForCausalLM.from_pretrained, '_local_vla_offline_patch', False):
        original_llama_from_pretrained = LlamaForCausalLM.from_pretrained

        def llama_from_pretrained_offline(model_name_or_path, *args, **kwargs):
            if str(model_name_or_path) == _LLAMA2_7B_ID:
                return LlamaForCausalLM(local_llama_config())
            return original_llama_from_pretrained(model_name_or_path, *args, **kwargs)

        llama_from_pretrained_offline._local_vla_offline_patch = True
        LlamaForCausalLM.from_pretrained = llama_from_pretrained_offline

    if not getattr(AutoTokenizer.from_pretrained, '_local_vla_offline_patch', False):
        original_tokenizer_from_pretrained = AutoTokenizer.from_pretrained

        def tokenizer_from_pretrained_offline(model_name_or_path, *args, **kwargs):
            if str(model_name_or_path) == _LLAMA2_7B_ID:
                kwargs = dict(kwargs)
                kwargs.pop('token', None)
                return original_tokenizer_from_pretrained(str(_TOKENIZER_DIR), *args, **kwargs)
            return original_tokenizer_from_pretrained(model_name_or_path, *args, **kwargs)

        tokenizer_from_pretrained_offline._local_vla_offline_patch = True
        AutoTokenizer.from_pretrained = tokenizer_from_pretrained_offline