import os
import torch
from safetensors.torch import save_file
from huggingface_hub import HfApi

# Configuration
SOURCE_CKPT = "models/policy_sac (2).pt"
OUTPUT_SAFE = "models/model.safetensors"
REPO_ID = "S-Mithilesh/Energizer-V1"

def convert_and_push():
    print(f"--- Loading weights from {SOURCE_CKPT}...")
    if not os.path.exists(SOURCE_CKPT):
        print(f"ERROR: {SOURCE_CKPT} not found.")
        return

    # 1. Load weights
    ckpt = torch.load(SOURCE_CKPT, map_location="cpu", weights_only=False)
    
    if isinstance(ckpt, dict) and 'actor_state' in ckpt:
        state_dict = ckpt['actor_state']
        print("INFO: Extracted 'actor_state' from full SAC checkpoint.")
    else:
        state_dict = ckpt
        print("INFO: Checkpoint seems to be a raw state_dict.")

    # 2. Save as safetensors
    print(f"--- Saving to {OUTPUT_SAFE}...")
    save_file(state_dict, OUTPUT_SAFE)
    print("SUCCESS: Conversion successful.")

    # 3. Upload to Hugging Face
    print(f"--- Uploading to HF Model Hub: {REPO_ID}...")
    api = HfApi()
    
    try:
        api.upload_file(
            path_or_fileobj=OUTPUT_SAFE,
            path_in_repo="model.safetensors",
            repo_id=REPO_ID,
            repo_type="model",
            commit_message="Update SAC model weights (safetensors)"
        )
        print(f"SUCCESS: Successfully uploaded model.safetensors to {REPO_ID}")
    except Exception as e:
        print(f"ERROR: Upload failed: {e}")
        print("Make sure you are logged in via 'huggingface-cli login'.")

if __name__ == "__main__":
    convert_and_push()