#!/usr/bin/env python3
"""
Test script to verify the Odoo AI Model Trainer setup
"""

import sys
import torch

def test_cuda():
    """Test CUDA availability"""
    print("Testing CUDA availability...")
    if torch.cuda.is_available():
        print(f"✅ CUDA available: {torch.cuda.get_device_name(0)}")
        print(f"   Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")
        return True
    else:
        print("❌ CUDA not available")
        return False

def test_imports():
    """Test if all required packages can be imported"""
    print("\nTesting imports...")

    required_packages = [
        'requests',
        'beautifulsoup4',
        'pandas',
        'numpy',
        'datasets',
        'transformers',
        'unsloth'
    ]

    failed_imports = []

    for package in required_packages:
        try:
            if package == 'beautifulsoup4':
                import bs4
            elif package == 'unsloth':
                try:
                    import unsloth
                except ImportError:
                    print(f"⚠️  {package} not available (this is normal if not installed)")
                    continue
            else:
                __import__(package)
            print(f"✅ {package}")
        except ImportError:
            print(f"❌ {package}")
            failed_imports.append(package)

    return len(failed_imports) == 0

def test_file_structure():
    """Test if all required files exist"""
    print("\nTesting file structure...")

    required_files = [
        'main.py',
        'data_scraper.py',
        'data_preprocessor.py',
        'train_model.py',
        'requirements.txt',
        'README.md'
    ]

    missing_files = []

    for file in required_files:
        try:
            with open(file, 'r') as f:
                pass
            print(f"✅ {file}")
        except FileNotFoundError:
            print(f"❌ {file}")
            missing_files.append(file)

    return len(missing_files) == 0

def main():
    """Run all tests"""
    print("🚀 Odoo AI Model Trainer - Setup Test")
    print("=" * 50)

    tests = [
        ("CUDA Setup", test_cuda),
        ("Package Imports", test_imports),
        ("File Structure", test_file_structure)
    ]

    results = []

    for test_name, test_func in tests:
        print(f"\n{test_name}:")
        print("-" * 30)
        result = test_func()
        results.append((test_name, result))

    print("\n" + "=" * 50)
    print("📋 Test Summary:")

    all_passed = True
    for test_name, result in results:
        status = "✅ PASS" if result else "❌ FAIL"
        print(f"   {test_name}: {status}")
        if not result:
            all_passed = False

    print("\n" + "=" * 50)
    if all_passed:
        print("🎉 All tests passed! Your setup is ready.")
        print("\nNext steps:")
        print("1. Run 'python main.py' to start the full pipeline")
        print("2. Or run 'python data_scraper.py' to start with data collection")
    else:
        print("❌ Some tests failed. Please fix the issues above.")
        print("\nCommon solutions:")
        print("1. Install missing packages: pip install -r requirements.txt")
        print("2. Check CUDA installation")
        print("3. Ensure all files are in the correct directory")

    return 0 if all_passed else 1

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