#!/usr/bin/env python3
"""
Quick fix script to check and resolve import issues
"""

import sys
import importlib
from pathlib import Path

def check_imports():
    """Check if all required modules can be imported"""
    
    print("🔍 Checking imports...")
    
    # Core Python modules
    core_modules = [
        'os', 'sys', 'asyncio', 'logging', 'uuid', 'json', 'shutil',
        'pathlib', 'datetime', 'typing'
    ]
    
    # FastAPI modules
    fastapi_modules = [
        'fastapi', 'uvicorn', 'pydantic'
    ]
    
    # ML modules
    ml_modules = [
        'torch', 'transformers', 'datasets', 'safetensors'
    ]
    
    # Utility modules
    utility_modules = [
        'numpy', 'pillow', 'requests', 'psutil', 'cryptography'
    ]
    
    # Optional modules
    optional_modules = [
        'cv2', 'pydicom', 'SimpleITK', 'intel_extension_for_pytorch'
    ]
    
    all_good = True
    
    # Check core modules
    print("\n📦 Core Python modules:")
    for module in core_modules:
        try:
            importlib.import_module(module)
            print(f"  ✅ {module}")
        except ImportError as e:
            print(f"  ❌ {module}: {e}")
            all_good = False
    
    # Check FastAPI modules
    print("\n🌐 FastAPI modules:")
    for module in fastapi_modules:
        try:
            importlib.import_module(module)
            print(f"  ✅ {module}")
        except ImportError as e:
            print(f"  ❌ {module}: {e}")
            all_good = False
    
    # Check ML modules
    print("\n🤖 Machine Learning modules:")
    for module in ml_modules:
        try:
            importlib.import_module(module)
            print(f"  ✅ {module}")
        except ImportError as e:
            print(f"  ❌ {module}: {e}")
            all_good = False
    
    # Check utility modules
    print("\n🔧 Utility modules:")
    for module in utility_modules:
        try:
            if module == 'pillow':
                importlib.import_module('PIL')
            elif module == 'opencv-python':
                importlib.import_module('cv2')
            else:
                importlib.import_module(module)
            print(f"  ✅ {module}")
        except ImportError as e:
            print(f"  ❌ {module}: {e}")
            all_good = False
    
    # Check optional modules
    print("\n🔍 Optional modules:")
    for module in optional_modules:
        try:
            importlib.import_module(module)
            print(f"  ✅ {module}")
        except ImportError as e:
            print(f"  ⚠️  {module}: {e} (optional)")
    
    return all_good

def check_custom_modules():
    """Check if custom modules can be imported"""
    
    print("\n🏗️  Custom modules:")
    
    custom_modules = [
        'src.model_loader',
        'src.distillation', 
        'src.utils',
        'src.core.memory_manager',
        'src.core.chunk_loader',
        'src.core.cpu_optimizer',
        'src.core.token_manager',
        'src.medical.medical_datasets',
        'src.medical.dicom_handler',
        'src.medical.medical_preprocessing',
        'database.database',
        'database.models'
    ]
    
    all_good = True
    
    for module in custom_modules:
        try:
            importlib.import_module(module)
            print(f"  ✅ {module}")
        except ImportError as e:
            print(f"  ❌ {module}: {e}")
            all_good = False
        except Exception as e:
            print(f"  ⚠️  {module}: {e} (import error)")
            all_good = False
    
    return all_good

def check_files():
    """Check if required files exist"""
    
    print("\n📁 Required files:")
    
    required_files = [
        'app.py',
        'requirements.txt',
        'src/__init__.py',
        'src/model_loader.py',
        'src/distillation.py',
        'src/utils.py',
        'src/core/__init__.py',
        'src/medical/__init__.py',
        'database/__init__.py',
        'templates/index.html',
        'templates/token-management.html',
        'templates/medical-datasets.html',
        'static/css/style.css',
        'static/js/main.js'
    ]
    
    all_good = True
    
    for file_path in required_files:
        path = Path(file_path)
        if path.exists():
            print(f"  ✅ {file_path}")
        else:
            print(f"  ❌ {file_path}")
            all_good = False
    
    return all_good

def main():
    """Main function"""
    
    print("🚀 AI Knowledge Distillation Platform - Import Checker")
    print("=" * 60)
    
    # Check imports
    imports_ok = check_imports()
    
    # Check custom modules
    custom_ok = check_custom_modules()
    
    # Check files
    files_ok = check_files()
    
    print("\n" + "=" * 60)
    
    if imports_ok and custom_ok and files_ok:
        print("✅ All checks passed! The application should start successfully.")
        return 0
    else:
        print("❌ Some checks failed. Please fix the issues above.")
        
        if not imports_ok:
            print("\n💡 To fix import issues:")
            print("   pip install -r requirements.txt")
        
        if not custom_ok:
            print("\n💡 To fix custom module issues:")
            print("   Check that all Python files are properly created")
            print("   Ensure __init__.py files exist in all directories")
        
        if not files_ok:
            print("\n💡 To fix missing files:")
            print("   Ensure all required files are created")
            print("   Check templates and static directories")
        
        return 1

if __name__ == "__main__":
    sys.exit(main())