"""
PRISM Checkpoint Exporter & Serializer.
Packages the Custom DiT model into standard PRISM-1.safetensors with embedded metadata.
"""

import argparse
import os
import time
from typing import Dict, Any
import torch
from safetensors.torch import save_file, load_file

from model import PrismDiT


def export_prism_safetensors(
    output_path: str = "PRISM-1.safetensors",
    precision: str = "fp16",
    hidden_size: int = 1152,
    depth: int = 28,
    num_heads: int = 16,
    context_dim: int = 1024,
    source_weights_path: str = None,
) -> str:
    """
    Initializes or loads PRISM DiT weights and serializes to .safetensors with complete metadata.
    """
    print(f"=== Initializing PRISM-1 Checkpoint Generator ===")
    print(f"Configuration:")
    print(f"  - Hidden Size: {hidden_size}")
    print(f"  - Depth (Layers): {depth}")
    print(f"  - Attention Heads: {num_heads}")
    print(f"  - Context Dimension: {context_dim}")
    print(f"  - Precision: {precision}")
    print(f"  - Target File: {output_path}")

    # 1. Instantiate DiT Model
    model = PrismDiT(
        input_size=64,
        patch_size=2,
        in_channels=16,
        hidden_size=hidden_size,
        depth=depth,
        num_heads=num_heads,
        context_dim=context_dim,
    )

    # 2. Optionally load existing source weights
    if source_weights_path and os.path.exists(source_weights_path):
        print(f"Loading source weights from: {source_weights_path}")
        if source_weights_path.endswith(".safetensors"):
            state_dict = load_file(source_weights_path)
        else:
            state_dict = torch.load(source_weights_path, map_location="cpu")
            if "state_dict" in state_dict:
                state_dict = state_dict["state_dict"]
            elif "model" in state_dict:
                state_dict = state_dict["model"]
        model.load_state_dict(state_dict, strict=False)
        print("Successfully merged source weights into PRISM DiT.")
    else:
        print("No external source weights provided; generating clean base initialized checkpoint...")

    # 3. Cast precision
    dtype_map = {
        "fp32": torch.float32,
        "fp16": torch.float16,
        "bf16": torch.bfloat16,
    }
    target_dtype = dtype_map.get(precision.lower(), torch.float16)
    model = model.to(dtype=target_dtype)

    # 4. Extract contiguous state dict
    state_dict = {k: v.contiguous() for k, v in model.state_dict().items()}
    total_params = sum(p.numel() for p in state_dict.values())
    print(f"Total Parameters: {total_params / 1e6:.2f}M")

    # 5. Metadata for Hugging Face / Civitai / ComfyUI loader inspection
    metadata: Dict[str, str] = {
        "model_name": "PRISM-1",
        "architecture": "PrismDiT (Diffusion Transformer)",
        "version": "1.0.0",
        "format": "safetensors",
        "precision": precision,
        "in_channels": "16",
        "patch_size": "2",
        "hidden_size": str(hidden_size),
        "depth": str(depth),
        "num_heads": str(num_heads),
        "context_dim": str(context_dim),
        "diffusion_type": "Flow-Matching / Rectified Flow",
        "recommended_sampler": "Euler",
        "recommended_steps": "12-25",
        "recommended_cfg": "5.0-6.5",
        "generator": "PRISM Checkpoint Exporter v1.0",
    }

    # 6. Save safetensors
    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
    t0 = time.time()
    save_file(state_dict, output_path, metadata=metadata)
    elapsed = time.time() - t0

    file_size_mb = os.path.getsize(output_path) / (1024 * 1024)
    print(f"\n[SUCCESS] Checkpoint saved successfully!")
    print(f"  File: {output_path}")
    print(f"  Size: {file_size_mb:.2f} MB")
    print(f"  Time taken: {elapsed:.2f}s")

    return output_path


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Export PRISM-1.safetensors Checkpoint")
    parser.add_argument("--output", type=str, default="PRISM-1.safetensors")
    parser.add_argument("--precision", type=str, default="fp16", choices=["fp32", "fp16", "bf16"])
    parser.add_argument("--hidden-size", type=int, default=1152)
    parser.add_argument("--depth", type=int, default=28)
    parser.add_argument("--heads", type=int, default=16)
    parser.add_argument("--context-dim", type=int, default=1024)
    parser.add_argument("--source-weights", type=str, default=None)
    args = parser.parse_args()

    export_prism_safetensors(
        output_path=args.output,
        precision=args.precision,
        hidden_size=args.hidden_size,
        depth=args.depth,
        num_heads=args.heads,
        context_dim=args.context_dim,
        source_weights_path=args.source_weights,
    )