"""Reload packed FlatQuant W4A4 weights (Linear4bit + transforms) from the artifact."""
from __future__ import annotations

import json
import os
import sys
from argparse import Namespace

import torch
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

QUANT_AGENT_ADAPTER_API = 1


def load_model_and_tokenizer(model_path, model_id, dtype, device, trust_remote_code):
    method_repo = os.environ.get("QUANT_AGENT_METHOD_REPO")
    if method_repo:
        sys.path.insert(0, method_repo)
    import flatquant.utils as utils
    from flatquant.model_tools.phi3_utils import apply_flatquant_to_phi3
    import flatquant.pack_utils as pack_utils

    packed_path = os.path.join(model_path, "packed_w4a4.pt")
    if not os.path.isfile(packed_path):
        raise FileNotFoundError("packed W4A4 artifact missing: %s" % packed_path)
    dtypes = {"float16": torch.float16, "bfloat16": torch.bfloat16, "auto": "auto"}
    torch_dtype = dtypes.get(dtype, torch.float16)
    tokenizer = AutoTokenizer.from_pretrained(
        model_path, trust_remote_code=trust_remote_code, local_files_only=True
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    config = AutoConfig.from_pretrained(
        model_path, trust_remote_code=trust_remote_code, local_files_only=True
    )
    config._attn_implementation = "eager"
    config._attn_implementation_internal = "eager"
    utils.skip_initialization()
    model = AutoModelForCausalLM.from_config(
        config,
        torch_dtype=torch_dtype if device == "cuda" else torch.float32,
        trust_remote_code=trust_remote_code,
        attn_implementation="eager",
    )
    model.seqlen = 2048
    with open(os.path.join(model_path, "flatquant_args.json"), "r", encoding="utf-8") as handle:
        args = Namespace(**json.load(handle))
    model = apply_flatquant_to_phi3(args, model)
    for layer in model.model.layers:
        layer.self_attn.rep_matrix_only()
        layer.mlp.rep_matrix_only()
    pack_utils.prepare_packed_slots(model)
    payload = torch.load(packed_path, map_location="cpu", weights_only=False)
    missing, unexpected = model.load_state_dict(payload, strict=False)
    packed_missing = [key for key in missing if "linear.weight" in key or key.endswith("a4.clip_factor_a_max")]
    if packed_missing:
        raise RuntimeError(
            "packed_w4a4.pt did not restore Linear4bit weights: missing=%s unexpected=%s"
            % (packed_missing[:12], list(unexpected)[:12])
        )
    pack_utils.cast_linear4bit_kernel_dtypes(model)
    pack_utils.restore_packed_eval_runtime(model)
    model.to(device)
    pack_utils.cast_linear4bit_kernel_dtypes(model)
    model.eval()
    return model, tokenizer