#=====================================================================================================
# This script converts a PyTorch model to SafeTensors format and uploads it to Hugging Face Hub.
# It is used to fix the "safetensors not found" error.
# Usage: python scripts/convert_safetensors.py
# It need to Run only once if the "safetensors not found" error came otherwise don't run it.
# NOTE: U need write access to the repo to upload the model.
#=====================================================================================================

import sys
import os
import shutil

# Adjust path to find app module
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.append(parent_dir)

from app.config import settings
from transformers import AutoModel, AutoTokenizer
from huggingface_hub import HfApi

def convert():
    model_name = settings.MODEL_NAME
    token = settings.HUGGING_FACE_TOKEN
    temp_dir = os.path.join(parent_dir, "temp_safe_model")
    
    print(f"Loading original model: {model_name}...")
    try:
        # Load the PyTorch version explicitly
        model = AutoModel.from_pretrained(
            model_name, 
            use_safetensors=False, 
            token=token
        )
        tokenizer = AutoTokenizer.from_pretrained(
            model_name, 
            token=token
        )
    except Exception as e:
        print(f"Failed to load original model: {e}")
        return

    print("Saving model locally with SafeTensors format...")
    try:
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)
        os.makedirs(temp_dir, exist_ok=True)
            
        model.save_pretrained(temp_dir, safe_serialization=True)
        tokenizer.save_pretrained(temp_dir)
        
        # --- Verification Step ---
        print("Verifying converted model by loading it back...")
        try:
            # Try to load the model from the temporary directory using SafeTensors
            check_model = AutoModel.from_pretrained(temp_dir, use_safetensors=True)
            print("Verification successful! Model loaded correctly from SafeTensors.")
            # memory cleanup
            del check_model
        except Exception as e:
            print(f"Verification FAILED: {e}")
            print("Aborting upload.")
            return

        # --- Upload Step ---
        print("Model verified. Now uploading to Hub...")
        
        api = HfApi(token=token)
        api.upload_folder(
            folder_path=temp_dir,
            repo_id=model_name,
            repo_type="model"
        )
        
        print("Success! The model has been converted, verified, and pushed to your repository.")
        print("The auto-conversion error should now be resolved.")
        
    except Exception as e:
        print(f"An error occurred during the process: {e}")
    finally:
        # Cleanup
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)
            print("Cleaned up temporary files.")

if __name__ == "__main__":
    convert()