#!/usr/bin/env python3
"""
Test script to verify pygrib installation and Arctic GRIB file processing
"""
import sys
import os

def test_pygrib_import():
    """Test if pygrib can be imported successfully"""
    print("🔍 Testing pygrib import...")
    try:
        import pygrib
        print(f"✅ pygrib imported successfully, version: {pygrib.__version__}")
        return True
    except ImportError as e:
        print(f"❌ pygrib import failed: {e}")
        return False

def test_eccodes_import():
    """Test if eccodes can be imported successfully"""
    print("🔍 Testing eccodes import...")
    try:
        import eccodes
        print(f"✅ eccodes imported successfully")
        return True
    except ImportError as e:
        print(f"❌ eccodes import failed: {e}")
        return False

def test_arctic_grib_processing():
    """Test Arctic GRIB file processing with pygrib"""
    arctic_file = "arctic_manual_20250828_12z.grib2"
    
    if not os.path.exists(arctic_file):
        print(f"⚠️ Arctic test file not found: {arctic_file}")
        return False
    
    print(f"🧪 Testing Arctic GRIB processing with pygrib...")
    
    try:
        import pygrib
        
        # Open GRIB file
        grbs = pygrib.open(arctic_file)
        print(f"📂 Successfully opened Arctic GRIB file")
        
        # Count messages
        msg_count = grbs.messages
        print(f"📊 Total messages in file: {msg_count}")
        
        # Test reading first few messages
        grbs.rewind()
        wave_params_found = 0
        
        for i, grb in enumerate(grbs):
            if i >= 10:  # Test first 10 messages
                break
                
            param_name = grb.name
            short_name = grb.shortName if hasattr(grb, 'shortName') else 'unknown'
            
            print(f"  Message {i+1}: {param_name} ({short_name})")
            
            # Look for wave parameters
            if any(keyword in param_name.lower() for keyword in ['wave', 'swell', 'height', 'period']):
                wave_params_found += 1
                print(f"    🌊 Found wave parameter!")
                
                try:
                    # Test coordinate extraction
                    lats, lons = grb.latlons()
                    values = grb.values
                    
                    print(f"    📍 Grid shape: {values.shape}")
                    print(f"    📈 Value range: {values.min():.3f} to {values.max():.3f}")
                    print(f"    🗺️ Lat range: {lats.min():.2f}° to {lats.max():.2f}°")
                    print(f"    🗺️ Lon range: {lons.min():.2f}° to {lons.max():.2f}°")
                    
                    # Check for Arctic coverage
                    arctic_points = (lats >= 50.0).sum()
                    print(f"    🧊 Arctic points (≥50°N): {arctic_points:,}")
                    
                except Exception as coord_error:
                    print(f"    ❌ Coordinate extraction failed: {coord_error}")
                    continue
        
        grbs.close()
        print(f"✅ Arctic GRIB processing test completed")
        print(f"🌊 Wave parameters found: {wave_params_found}")
        
        return wave_params_found > 0
        
    except Exception as e:
        print(f"❌ Arctic GRIB processing failed: {e}")
        import traceback
        traceback.print_exc()
        return False

def main():
    """Run all pygrib tests"""
    print("🧪 PyGRIB Setup Test Suite")
    print("=" * 50)
    
    tests_passed = 0
    total_tests = 3
    
    # Test 1: Import pygrib
    if test_pygrib_import():
        tests_passed += 1
    
    print()
    
    # Test 2: Import eccodes
    if test_eccodes_import():
        tests_passed += 1
    
    print()
    
    # Test 3: Arctic GRIB processing (if file exists)
    if test_arctic_grib_processing():
        tests_passed += 1
    
    print()
    print("=" * 50)
    print(f"✅ Tests passed: {tests_passed}/{total_tests}")
    
    if tests_passed == total_tests:
        print("🎉 All tests passed! pygrib is working correctly.")
    else:
        print("⚠️ Some tests failed. Check error messages above.")
    
    return tests_passed == total_tests

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)