import joblib
from sentence_transformers import SentenceTransformer
import os

MODEL_DIR = "models"

# Initialize variables
tfidf_vectorizer = None
le = None
rf_model = None
sentence_model = None

try:
    tfidf_vectorizer = joblib.load(os.path.join(MODEL_DIR, "tfidf_vectorizer.pkl"))
    le = joblib.load(os.path.join(MODEL_DIR, "label_encoder.pkl"))
    rf_model = joblib.load(os.path.join(MODEL_DIR, "random_forest_model.pkl"))
    print("Classification models loaded.")

    # Set cache directory to a writable location for Hugging Face Spaces
    os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
    os.environ['HF_HOME'] = '/tmp/hf_home'

    sentence_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
    print("SentenceTransformer model loaded.")

except FileNotFoundError as e:
    print(f"MODEL LOADING ERROR: {e}")
    print("Make sure the .pkl files are in the 'models' directory.")
    raise e  # Re-raise to prevent the application from starting with None models

except Exception as e:
    print(f"An unexpected error occurred during model loading: {e}")
    raise e  # Re-raise to prevent the application from starting with None models