wave_data_puller.py

import requests
import json
import os
import time
from datetime import datetime
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class WaveDataPuller:
    def __init__(self, space_url="https://huggingface.co/spaces/nakas/NWPS_SWAN"):
        self.space_url = space_url
        self.api_url = f"{space_url}/run/predict"
        self.output_dir = os.getenv('OUTPUT_DIR', '/tmp/wave_data')
        self.poll_interval = int(os.getenv('POLL_INTERVAL', '3600'))  # 1 hour default
        
        # Create output directory with error handling for restricted environments
        try:
            os.makedirs(self.output_dir, exist_ok=True)
        except PermissionError:
            # Fall back to /tmp if we can't create in the specified location
            self.output_dir = '/tmp/wave_data'
            os.makedirs(self.output_dir, exist_ok=True)
        

    def fetch_wave_data(self, lat=40.0, lon=-74.0):
        """Fetch wave data from the Hugging Face space for a specific location."""
        try:
            logger.info(f"Fetching wave data for lat={lat}, lon={lon} from Hugging Face space...")
            
            # Construct the payload for the Gradio API call
            # Assuming the NWPS_SWAN space has a function at fn_index 0 that takes lat and lon
            payload = {
                "fn_index": 0, 
                "data": [float(lat), float(lon)]
            }
            
            headers = {"Content-Type": "application/json"}
            response = requests.post(self.api_url, headers=headers, data=json.dumps(payload))
            response.raise_for_status()  # Raise an exception for HTTP errors
            
            result = response.json()
            
            # Gradio API responses have a specific structure, usually with 'data' key
            if "data" in result and len(result["data"]) > 0:
                # Assuming the first element of the data array is the relevant wave data
                return result["data"][0]
            else:
                logger.warning(f"Unexpected API response structure: {result}")
                logger.warning("No accessible endpoints found, returning mock data")
                return self.generate_mock_data(lat, lon)
            
        except requests.exceptions.RequestException as e:
            logger.error(f"Request to Hugging Face Space failed: {e}")
            logger.warning("No accessible endpoints found, returning mock data")
            return self.generate_mock_data(lat, lon)
        except Exception as e:
            logger.error(f"Error fetching wave data: {e}")
            return None

    def generate_mock_data(self, lat=40.0, lon=-74.0):
        """Generate mock SWAN wave data for testing"""
        import random
        
        return {
            "timestamp": datetime.utcnow().isoformat(),
            "location": {"lat": lat, "lon": lon},
            "wave_data": {
                "significant_wave_height": round(random.uniform(0.5, 3.0), 2),
                "peak_wave_period": round(random.uniform(4.0, 12.0), 2),
                "wave_direction": round(random.uniform(0, 360), 1),
                "wind_speed": round(random.uniform(2.0, 15.0), 2),
                "wind_direction": round(random.uniform(0, 360), 1)
            },
            "model": "NWPS_SWAN",
            "forecast_hours": [0, 6, 12, 18, 24]
        }

    def save_data(self, data):
        """Save wave data to file"""
        if not data:
            return
        
        timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
        filename = f"wave_data_{timestamp}.json"
        filepath = os.path.join(self.output_dir, filename)
        
        try:
            with open(filepath, 'w') as f:
                json.dump(data, f, indent=2)
            logger.info(f"Data saved to {filepath}")
        except Exception as e:
            logger.error(f"Error saving data: {e}")

    def run_continuous(self, lat=40.0, lon=-74.0):
        """Run continuous data pulling"""
        logger.info(f"Starting continuous wave data pulling every {self.poll_interval} seconds")
        
        while True:
            try:
                data = self.fetch_wave_data(lat, lon)
                if data:
                    self.save_data(data)
                    logger.info("Wave data fetched and saved successfully")
                else:
                    logger.warning("No data received")
                
                time.sleep(self.poll_interval)
                
            except KeyboardInterrupt:
                logger.info("Stopping wave data puller...")
                break
            except Exception as e:
                logger.error(f"Unexpected error: {e}")
                time.sleep(60)  # Wait 1 minute before retrying

    def run_once(self, lat=40.0, lon=-74.0):
        """Run a single data pull"""
        logger.info("Fetching wave data once...")
        data = self.fetch_wave_data(lat, lon)
        if data:
            self.save_data(data)
            logger.info("Wave data fetched and saved successfully")
        else:
            logger.error("Failed to fetch wave data")

if __name__ == "__main__":
    puller = WaveDataPuller()
    
    # Check if we should run once or continuously
    run_mode = os.getenv('RUN_MODE', 'once')
    
    if run_mode.lower() == 'continuous':
        puller.run_continuous()
    else:
        puller.run_once()