#!/usr/bin/env python3
"""
Timeout wrapper for pygrib operations to prevent hanging on polar stereographic projections
"""
import multiprocessing
import functools
import logging
import numpy as np
import time
import queue

logger = logging.getLogger(__name__)

class TimeoutError(Exception):
    pass

def run_with_timeout(func, timeout_seconds, *args, **kwargs):
    """
    Run a function with timeout using multiprocessing
    """
    def target_func(q, func, args, kwargs):
        try:
            result = func(*args, **kwargs)
            q.put(('success', result))
        except Exception as e:
            q.put(('error', e))
    
    q = multiprocessing.Queue()
    process = multiprocessing.Process(target=target_func, args=(q, func, args, kwargs))
    process.start()
    process.join(timeout_seconds)
    
    if process.is_alive():
        process.terminate()
        process.join()
        raise TimeoutError(f"Function timed out after {timeout_seconds} seconds")
    
    if q.empty():
        raise TimeoutError(f"Function failed to return result")
    
    status, result = q.get()
    if status == 'error':
        raise result
    return result

def safe_pygrib_open(grib_file):
    """Safe pygrib.open with timeout"""
    def _open_grib(grib_file):
        import pygrib
        return pygrib.open(grib_file)
    
    return run_with_timeout(_open_grib, 30, grib_file)

def safe_pygrib_latlons(grb):
    """Safe grb.latlons() with timeout for polar stereographic"""
    # Can't pass pygrib objects between processes, so this won't work
    # Return None to indicate we need to use fallback method
    return None

def safe_pygrib_values(grb):
    """Safe grb.values with timeout"""
    # Can't pass pygrib objects between processes
    return grb.values

def extract_pygrib_with_timeout(grib_file, sample_points=100):
    """
    Extract Arctic GRIB data using advanced polar coordinate methods
    """
    try:
        # Use the specialized polar coordinate extractor
        from polar_grib_coords import extract_arctic_data_safe
        logger.info("🧊 Using advanced polar coordinate extraction")
        return extract_arctic_data_safe(grib_file, max_messages=3)
        
    except ImportError:
        logger.warning("⚠️ Advanced polar extractor not available, using fallback")
        return extract_pygrib_simple_fallback(grib_file, sample_points)

def extract_pygrib_simple_fallback(grib_file, sample_points=100):
    """
    Fallback pygrib extraction avoiding hanging operations
    """
    try:
        logger.info("🔄 Using fallback pygrib extraction")
        import pygrib
        
        grbs = pygrib.open(grib_file)
        data_points = []
        
        # Process only first few messages to avoid hanging
        message_count = 0
        for grb in grbs:
            message_count += 1
            if message_count > 3:  # Limit to first 3 messages
                break
                
            try:
                param_name = grb.name if hasattr(grb, 'name') else f'param_{message_count}'
                
                # Check if this is a wave parameter
                if any(keyword in param_name.lower() for keyword in 
                      ['wave', 'swell', 'height', 'period', 'wind']):
                    
                    logger.info(f"🌊 Processing wave parameter: {param_name}")
                    
                    # Get values (this should work)
                    values = grb.values
                    
                    # Create synthetic Arctic grid coordinates (avoiding latlons())
                    logger.info("🗺️ Using synthetic Arctic coordinates")
                    ny, nx = values.shape
                    lats_1d = np.linspace(50.0, 85.0, ny)  # Arctic latitudes
                    lons_1d = np.linspace(-180.0, 180.0, nx)  # Full longitude range
                    lons_grid, lats_grid = np.meshgrid(lons_1d, lats_1d)
                    
                    # Flatten arrays
                    flat_lats = lats_grid.flatten()
                    flat_lons = lons_grid.flatten()
                    flat_values = values.flatten()
                    
                    # Filter valid values
                    valid_mask = ~np.isnan(flat_values) & (flat_values < 9999) & (flat_values > 0)
                    
                    if np.sum(valid_mask) > 0:
                        valid_lats = flat_lats[valid_mask]
                        valid_lons = flat_lons[valid_mask]
                        valid_vals = flat_values[valid_mask]
                        
                        # Sample points
                        n_points = min(sample_points, len(valid_lats))
                        if n_points > 0:
                            indices = np.random.choice(len(valid_lats), n_points, replace=False)
                            
                            for idx in indices:
                                data_points.append({
                                    'latitude': float(valid_lats[idx]),
                                    'longitude': float(valid_lons[idx]),
                                    'value': float(valid_vals[idx]),
                                    'parameter': param_name
                                })
                            
                            logger.info(f"✅ Extracted {len(indices)} Arctic points for {param_name}")
                            
                            if len(data_points) >= sample_points:
                                break
                        
            except Exception as e:
                logger.warning(f"❌ Failed to process message {message_count}: {e}")
                continue
        
        grbs.close()
        
        if data_points:
            logger.info(f"✅ Fallback extraction successful: {len(data_points)} Arctic points")
            return {
                'parameters': {},
                'coordinates': None,
                'data_points': data_points
            }
        else:
            logger.warning("❌ No Arctic wave data extracted with fallback")
            return {
                'parameters': {},
                'coordinates': None, 
                'data_points': []
            }
            
    except Exception as e:
        logger.error(f"❌ Fallback extraction failed: {e}")
        return {
            'parameters': {},
            'coordinates': None,
            'data_points': []
        }

if __name__ == "__main__":
    import numpy as np
    import os
    
    # Test timeout functionality
    arctic_file = "arctic_manual_20250828_12z.grib2"
    if os.path.exists(arctic_file):
        result = extract_pygrib_with_timeout(arctic_file, sample_points=50)
        print(f"Result: {result['status']}, Points: {len(result['data'])}")
    else:
        print("Test Arctic file not found")