main.py

import os
import logging
import requests
from datetime import datetime
from flask import Flask, jsonify
from apscheduler.schedulers.background import BackgroundScheduler

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

app = Flask(__name__)

# Global status store
ping_status = {
    "last_ping": None,
    "results": []
}

def get_spaces():
    """Parse SPACES_URLS from environment variable."""
    spaces = os.environ.get("SPACES_URLS", "")
    if not spaces:
        logger.warning("SPACES_URLS not set")
        return []
    return [url.strip() for url in spaces.split(",") if url.strip()]


def get_headers():
    """Build request headers including HF token if available."""
    headers = {}
    hf_token = os.environ.get("HF_TOKEN", "")
    if hf_token:
        headers["Authorization"] = f"Bearer {hf_token}"
    return headers


def ping_space(url):
    """Ping a single space and return result."""
    try:
        headers = get_headers()
        response = requests.get(url, headers=headers, timeout=30)
        return {
            "url": url,
            "status": "success",
            "status_code": response.status_code,
            "timestamp": datetime.now().isoformat()
        }
    except requests.exceptions.Timeout:
        return {
            "url": url,
            "status": "timeout",
            "error": "Request timed out",
            "timestamp": datetime.now().isoformat()
        }
    except requests.exceptions.RequestException as e:
        return {
            "url": url,
            "status": "error",
            "error": str(e),
            "timestamp": datetime.now().isoformat()
        }


def ping_all_spaces():
    """Ping all configured spaces."""
    spaces = get_spaces()
    if not spaces:
        logger.info("No spaces configured, skipping ping")
        return

    logger.info(f"Pinging {len(spaces)} space(s)...")
    results = []
    for url in spaces:
        result = ping_space(url)
        results.append(result)
        
        if result["status"] == "success":
            logger.info(f"✓ {url} - OK ({result['status_code']})")
        else:
            logger.error(f"✗ {url} - {result.get('status', 'unknown')}: {result.get('error', 'unknown error')}")

    ping_status["last_ping"] = datetime.now().isoformat()
    ping_status["results"] = results


@app.route("/")
def index():
    """Status page."""
    spaces = get_spaces()
    return jsonify({
        "status": "running",
        "configured_spaces": len(spaces),
        "spaces": spaces,
        "last_ping": ping_status["last_ping"],
        "last_results": ping_status["results"]
    })


@app.route("/health")
def health():
    """Health check endpoint."""
    return jsonify({"status": "healthy"})


if __name__ == "__main__":
    # Get ping interval from environment (default: 1 hour = 3600 seconds)
    ping_interval_seconds = int(os.environ.get("PING_INTERVAL_SECONDS", 3600))
    
    # Setup scheduler
    scheduler = BackgroundScheduler()
    scheduler.add_job(
        ping_all_spaces,
        'interval',
        seconds=ping_interval_seconds,
        id='pinger_job'
    )
    scheduler.start()
    logger.info(f"Pinger started. Will ping every {ping_interval_seconds} seconds")
    
    # Initial ping
    ping_all_spaces()
    
    # Run Flask app (required by HuggingFace Spaces)
    app.run(host="0.0.0.0", port=7860)