"""Load the actual verified Swift BF16 source lazily, without quantization."""

import argparse
import hashlib
import json
import platform
import resource
import time
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path

import mlx.core as mx
from mlx.utils import tree_flatten
from transformers import AutoProcessor, AutoTokenizer

from mlx_lm.utils import load_model


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", required=True, type=Path)
    parser.add_argument("--verification", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args()
    if args.output.exists():
        raise FileExistsError(args.output)
    source = args.source.resolve(strict=True)
    verified = json.loads(args.verification.read_text())
    assert verified["status"] == "PASS"
    assert Path(verified["source_root"]).resolve() == source
    manifest_hash = "0a00065b88ab003281853a7fb9bd5ce0086bc3781b36136d8c39da19933923ae"
    assert verified["manifest_sha256"] == manifest_hash
    assert hashlib.sha256((source / "EXPORT_MANIFEST.json").read_bytes()).hexdigest() == manifest_hash
    assert verified["shard_count"] == 18
    assert verified["shard_bytes"] == 55563006776
    assert verified["tensor_count"] == 1199
    for item in verified["files"]:
        assert (source / item["name"]).stat().st_size == item["bytes"]

    started = time.monotonic()
    model, config = load_model(source, lazy=True, strict=True)
    assert type(model).__module__ == "mlx_lm.models.qwen3_5_full"
    assert not config.get("quantization") and not config.get("quantization_config")
    rows = model.weight_mapping()
    parameters = dict(tree_flatten(model.parameters()))
    assert {row["source"] for row in rows} == set(verified["tensors"])
    assert {row["destination"] for row in rows} == set(parameters)
    for row in rows:
        header = verified["tensors"][row["source"]]
        value = parameters[row["destination"]]
        assert list(row["source_shape"]) == header["shape"]
        assert list(value.shape) == list(row["destination_shape"])
        assert value.dtype == mx.bfloat16 and header["dtype"] == "BF16"
    categories = dict(Counter(row["category"] for row in rows))
    assert categories == {"text": 851, "vision": 333, "MTP": 15}

    tokenizer = AutoTokenizer.from_pretrained(source, local_files_only=True, trust_remote_code=False)
    assert tokenizer.chat_template
    chats = []
    for options in [{"enable_thinking": False}, {"reasoning_effort": "low"}, {"reasoning_effort": "xhigh"}]:
        prompt = tokenizer.apply_chat_template(
            [{"role": "user", "content": "Say hello."}],
            tokenize=False, add_generation_prompt=True, **options,
        )
        ids = tokenizer.encode(prompt, add_special_tokens=False)
        assert ids and all(0 <= token < config["text_config"]["vocab_size"] for token in ids)
        chats.append({"options": options, "tokens": len(ids), "rendered": prompt})
    processor = AutoProcessor.from_pretrained(source, local_files_only=True, trust_remote_code=False)
    result = {
        "status": "PASS_ACTUAL_SOURCE_LAZY_STRUCTURAL_LOAD",
        "recorded_at": datetime.now(timezone.utc).isoformat(),
        "platform": platform.platform(),
        "source": str(source),
        "source_repo": "ukisai/Swift-1.5-Qwen3.8-27b",
        "source_revision": "00ccd14e006897d28cb0ed5bf26390e60d274251",
        "source_verification": str(args.verification.resolve()),
        "source_verification_sha256": hashlib.sha256(args.verification.read_bytes()).hexdigest(),
        "source_tensors": len(verified["tensors"]),
        "mapped_tensors": len(parameters),
        "categories": categories,
        "ignored_tensors": 0,
        "unexplained_tensors": 0,
        "weight_backing": "actual verified source safetensors; lazy loading",
        "full_parameter_evaluation": False,
        "tokenizer": type(tokenizer).__name__,
        "chat_templates": chats,
        "processor": type(processor).__name__,
        "generation": "NOT_RUN",
        "mtp_runtime": "component-only; no integrated speculative decoding",
        "vision_runtime": "encoder-only; no integrated multimodal generation",
        "quantization_executed": False,
        "elapsed_seconds": time.monotonic() - started,
        "mlx_active_memory_bytes": mx.get_active_memory(),
        "process_peak_rss_bytes": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * (1024 if platform.system() == "Linux" else 1),
    }
    with args.output.open("x") as stream:
        json.dump(result, stream, indent=2)
        stream.write("\n")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()