import argparse
import json
import os

import torch
from omegaconf import OmegaConf
from safetensors.torch import load_file, save_file

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Convert a PyTorch Lightning checkpoint to a safetensors file."
    )
    parser.add_argument("ckpt_path", type=str, help="Path to .ckpt file.")
    parser.add_argument(
        "--exclude_keys", type=str, nargs="*", default=[], help="Keys to exclude."
    )
    parser.add_argument(
        "--strip_prefix",
        type=str,
        default=None,
        help="Prefix to remove from keys (e.g., 'model.').",
    )
    parser.add_argument(
        "--verbose", action="store_true", help="Print keys being saved."
    )

    args = parser.parse_args()

    # Load checkpoint
    ckpt = torch.load(args.ckpt_path, map_location="cpu", weights_only=False)
    state_dict = ckpt["state_dict"]
    hparams = ckpt["hyper_parameters"]

    # Process: Exclude keys and strip prefixes
    parameters = {}
    for k, v in state_dict.items():
        if any(k.startswith(excl) for excl in args.exclude_keys):
            continue

        new_key = k
        if args.strip_prefix and k.startswith(f"{args.strip_prefix}."):
            new_key = k.replace(f"{args.strip_prefix}.", "", 1)

        parameters[new_key] = v

    if args.verbose:
        print("The following keys will be saved:")
        for key in parameters.keys():
            print(f" - {key}")

    # Save safetensors
    output_path = args.ckpt_path.replace(".ckpt", ".safetensors")
    save_file(parameters, output_path)
    print(f"Safetensors file saved to {output_path}")

    # Export config.json
    hparams_dict = OmegaConf.to_container(hparams, resolve=False)

    # We only save the 'model' key to keep the config clean
    config_data = hparams_dict.get("model", hparams_dict)

    config_path = os.path.join(os.path.dirname(output_path), "config.json")
    with open(config_path, "w", encoding="utf-8") as f:
        json.dump(config_data, f, indent=2)
    print(f"Configuration saved to {config_path}")

    # Verification
    try:
        loaded_params = load_file(output_path)
        assert len(parameters) == len(loaded_params), "Mismatch in parameter count!"
        for k in parameters:
            # We verify the tensor shape and dtype, as torch.equal can be slow/strict
            assert parameters[k].shape == loaded_params[k].shape, f"Shape mismatch: {k}"
            assert parameters[k].dtype == loaded_params[k].dtype, f"Dtype mismatch: {k}"
        print("Verification successful: File is valid.")
    except Exception as e:
        print(f"Verification failed: {e}")