import argparse
from pathlib import Path

import torch
from safetensors.torch import save_file


def extract_state_dict(checkpoint):
    if isinstance(checkpoint, dict):
        for key in ("model_state_dict", "encoder_state_dict", "state_dict", "model"):
            value = checkpoint.get(key)
            if isinstance(value, dict):
                return value
    return checkpoint


def clean_state_dict(state_dict):
    cleaned = {}
    for key, value in state_dict.items():
        if not isinstance(value, torch.Tensor):
            continue

        clean_key = key
        for prefix in ("module.", "encoder."):
            if clean_key.startswith(prefix):
                clean_key = clean_key[len(prefix) :]

        cleaned[clean_key] = value.detach().cpu().contiguous()
    return cleaned


def convert_checkpoint(input_path, output_path):
    try:
        checkpoint = torch.load(input_path, map_location="cpu", weights_only=False)
    except TypeError:
        checkpoint = torch.load(input_path, map_location="cpu")
    state_dict = extract_state_dict(checkpoint)

    if not isinstance(state_dict, dict):
        raise TypeError(f"Could not extract a state_dict from {input_path}")

    tensors = clean_state_dict(state_dict)
    if not tensors:
        raise ValueError(f"No tensors found in {input_path}")

    metadata = {
        "format": "pt",
        "source_checkpoint": input_path.name,
        "model": "FLORO",
    }
    save_file(tensors, output_path, metadata=metadata)
    return len(tensors)


def main():
    parser = argparse.ArgumentParser(
        description="Convert a FLORO PyTorch checkpoint to safetensors."
    )
    parser.add_argument(
        "--input",
        default="checkpoints/floro_encoder_202603_ep150.pth",
        type=Path,
        help="Path to the source PyTorch checkpoint.",
    )
    parser.add_argument(
        "--output",
        default="checkpoints/floro_encoder_202603_ep150.safetensors",
        type=Path,
        help="Path for the converted safetensors checkpoint.",
    )
    args = parser.parse_args()

    args.output.parent.mkdir(parents=True, exist_ok=True)
    count = convert_checkpoint(args.input, args.output)
    print(f"Saved {count} tensors to {args.output}")


if __name__ == "__main__":
    main()