#!/usr/bin/env python3
"""
Monitor and keep alive your Hugging Face Space
Run this on your local computer or a server
"""
import requests
import time
from datetime import datetime
import json

# Your Space details
SPACE_URL = "https://huggingface.co/spaces/MoraxCheng/Transeption_iGEM_BASISCHINA_2025"
EMBED_URL = "https://moraxcheng-transeption-igem-basischina-2025.hf.space"
API_URL = f"{EMBED_URL}/api/predict"

# Monitoring settings
CHECK_INTERVAL = 300  # 5 minutes
WAKE_UP_TIMEOUT = 180  # 3 minutes max wait for wake up

def check_space_status():
    """Check if Space is responding"""
    try:
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Checking Space status...")
        response = requests.get(EMBED_URL, timeout=10)
        
        if response.status_code == 200:
            print("✓ Space is online")
            return True
        else:
            print(f"⚠ Space returned status {response.status_code}")
            return False
    except requests.exceptions.Timeout:
        print("⚠ Request timed out - Space might be sleeping")
        return False
    except Exception as e:
        print(f"✗ Error checking status: {e}")
        return False

def wake_up_space():
    """Send a minimal request to wake up the Space"""
    try:
        print("Attempting to wake up Space...")
        
        # Send a minimal prediction request
        payload = {
            "fn_index": 0,
            "data": [
                "MSKGE",  # Minimal sequence
                1,        # Start position
                2,        # End position
                "Small",  # Model size
                False,    # No mirror scoring
                1         # Minimal batch size
            ]
        }
        
        headers = {
            "Content-Type": "application/json",
        }
        
        response = requests.post(
            API_URL,
            json=payload,
            headers=headers,
            timeout=30
        )
        
        if response.status_code in [200, 202]:
            print("✓ Wake-up request sent successfully")
            return True
        else:
            print(f"⚠ Wake-up request returned status {response.status_code}")
            return False
            
    except Exception as e:
        print(f"✗ Error sending wake-up request: {e}")
        return False

def wait_for_space_ready(max_wait=WAKE_UP_TIMEOUT):
    """Wait for Space to become ready"""
    print(f"Waiting for Space to become ready (max {max_wait}s)...")
    start_time = time.time()
    
    while time.time() - start_time < max_wait:
        if check_space_status():
            print(f"✓ Space is ready after {int(time.time() - start_time)}s")
            return True
        
        print(".", end="", flush=True)
        time.sleep(10)
    
    print(f"\n✗ Space did not become ready after {max_wait}s")
    return False

def monitor_loop():
    """Main monitoring loop"""
    print("="*60)
    print("Hugging Face Space Monitor")
    print(f"Space: {SPACE_URL}")
    print(f"Check interval: {CHECK_INTERVAL}s")
    print("Press Ctrl+C to stop")
    print("="*60)
    
    consecutive_failures = 0
    
    while True:
        try:
            # Check if Space is alive
            if check_space_status():
                consecutive_failures = 0
                print(f"Next check in {CHECK_INTERVAL}s...\n")
            else:
                consecutive_failures += 1
                print(f"Space appears to be down (failure #{consecutive_failures})")
                
                # Try to wake it up
                if wake_up_space():
                    # Wait for it to become ready
                    if wait_for_space_ready():
                        consecutive_failures = 0
                        print("✓ Space successfully revived!\n")
                    else:
                        print("✗ Failed to revive Space\n")
                else:
                    print("✗ Could not send wake-up request\n")
                
                if consecutive_failures >= 3:
                    print("⚠️  ATTENTION: Space has been down for multiple checks!")
                    print("You may need to manually restart it from the Hugging Face interface.")
                    print(f"Go to: {SPACE_URL}/settings\n")
            
            # Wait before next check
            time.sleep(CHECK_INTERVAL)
            
        except KeyboardInterrupt:
            print("\n\nMonitoring stopped by user")
            break
        except Exception as e:
            print(f"\nUnexpected error: {e}")
            print("Continuing monitoring...\n")
            time.sleep(60)

if __name__ == "__main__":
    # Test connection first
    print("Testing connection to Space...")
    if check_space_status():
        print("✓ Initial connection successful\n")
    else:
        print("⚠ Space appears to be sleeping, attempting to wake...")
        if wake_up_space() and wait_for_space_ready():
            print("✓ Space is now ready\n")
        else:
            print("✗ Could not wake Space. It may need manual restart.\n")
    
    # Start monitoring
    monitor_loop()