#!/usr/bin/env python3
"""
Model loader that handles joblib compatibility issues
"""

import os
import joblib
import pickle
import sys
from typing import Optional, Any

# Import our models module
from models import DataPreprocessor


def load_models_safely(models_dir: str = 'models') -> tuple[Optional[Any], Optional[DataPreprocessor]]:
    """
    Safely load models with proper class resolution for joblib
    
    Returns:
        tuple: (model, preprocessor) or (None, None) if failed
    """
    
    # Set up the module namespace for joblib
    current_module = sys.modules[__name__]
    current_module.DataPreprocessor = DataPreprocessor
    
    # Also set it in the main module
    import __main__
    __main__.DataPreprocessor = DataPreprocessor
    
    # And in the models module
    import models
    models.DataPreprocessor = DataPreprocessor
    
    try:
        # Try to load preprocessor
        preprocessor_path = os.path.join(models_dir, 'preprocessor.pkl')
        if not os.path.exists(preprocessor_path):
            print(f"❌ Preprocessor file not found: {preprocessor_path}")
            return None, None
        
        print(f"📥 Loading preprocessor from {preprocessor_path}")
        preprocessor = joblib.load(preprocessor_path)
        print("✅ Preprocessor loaded successfully")
        
        # Try to load Random Forest model
        rf_path = os.path.join(models_dir, 'random_forest_model.pkl')
        if not os.path.exists(rf_path):
            print(f"❌ Random Forest file not found: {rf_path}")
            return None, None
            
        print(f"📥 Loading Random Forest model from {rf_path}")
        model = joblib.load(rf_path)
        print("✅ Random Forest model loaded successfully")
        
        return model, preprocessor
        
    except Exception as e:
        print(f"❌ Error loading models: {e}")
        
        # Try alternative loading method
        try:
            print("🔄 Trying alternative loading method...")
            
            # Load with explicit class mapping
            with open(preprocessor_path, 'rb') as f:
                # Create a custom unpickler that knows about our classes
                import pickle
                
                class CustomUnpickler(pickle.Unpickler):
                    def find_class(self, module, name):
                        if name == 'DataPreprocessor':
                            return DataPreprocessor
                        return super().find_class(module, name)
                
                preprocessor = CustomUnpickler(f).load()
            
            print("✅ Preprocessor loaded with custom unpickler")
            
            # Load model normally (should work since it's sklearn)
            model = joblib.load(rf_path)
            print("✅ Random Forest model loaded successfully")
            
            return model, preprocessor
            
        except Exception as e2:
            print(f"❌ Alternative loading also failed: {e2}")
            return None, None


def test_model_loading():
    """Test function to verify model loading works"""
    print("🧪 Testing model loading...")
    model, preprocessor = load_models_safely()
    
    if model is not None and preprocessor is not None:
        print("✅ Model loading test passed!")
        
        # Test basic functionality
        try:
            print("🔍 Testing preprocessor functionality...")
            print(f"   - Label encoders available: {list(preprocessor.label_encoders.keys()) if hasattr(preprocessor, 'label_encoders') else 'None'}")
            print(f"   - Scaler available: {preprocessor.scaler is not None if hasattr(preprocessor, 'scaler') else 'Unknown'}")
            
            print("🔍 Testing model functionality...")
            print(f"   - Model type: {type(model).__name__}")
            
            return True
        except Exception as e:
            print(f"❌ Functionality test failed: {e}")
            return False
    else:
        print("❌ Model loading test failed!")
        return False


if __name__ == "__main__":
    test_model_loading()