src/utils.py

import requests
from typing import Optional
from bs4 import BeautifulSoup
import time
import re
import os
from dotenv import load_dotenv

load_dotenv()

def get_url_list():
    # Liste d'URLs à pinger
    default_urls = os.getenv("DEFAULT_URLS", "")
    url_list = [url.strip() for url in default_urls.split(",") if url.strip()]

    return url_list

def add_url(url: str) -> None:
    """
    Add a new URL to the list of URLs.
    """
    if url not in get_url_list():
        os.environ["DEFAULT_URLS"] = ",".join([os.environ["DEFAULT_URLS"], url])
        print(f"Added URL: {url}")
    else:
        print(f"URL already exists: {url}")

def remove_url(url: str) -> None:
    """
    Remove a URL from the list of URLs.
    """
    url_list = get_url_list()
    if url in url_list:
        url_list.remove(url)
        os.environ["DEFAULT_URLS"] = ",".join(url_list)
        print(f"Removed URL: {url}")
    else:
        print(f"URL not found: {url}")

# Fonction Gradio pour ajouter / retirer des URLs
def manage_urls(action, url):
    url_list = get_url_list()

    if action == "Ajouter":
        if url and url not in url_list:
            add_url(url)
            message = f"✅ Ajouté: {url}"
        else:
            message = "⚠️ URL déjà présente ou vide"
    elif action == "Supprimer":
        if url in url_list:
            remove_url(url)
            message = f"❌ Supprimé: {url}"
        else:
            message = "⚠️ URL non trouvée"
    else:
        message = "❓ Action invalide"

    # Toujours retourner la liste actuelle formatée
    return message, "\n".join(get_url_list())

def wake_up_space(space_url: str):
    action_name = get_wake_up_action_name(space_url)

    if not action_name:
        print(f"[wake-up] ❌ Invalid URL: {space_url}")
        return
    
    try:
        session = requests.Session()
        response = session.get(space_url, timeout=10)
        response.raise_for_status()
        
        # Vérifier si le formulaire de réveil est présent
        soup = BeautifulSoup(response.text, "html.parser")
        form = soup.find("form", action=action_name)

        if form:
            print("[wake-up] 💤 Space is sleeping, trying to restart it...")
            csrf_token_input = form.find("input", {"name": "csrf"})
            csrf_token = csrf_token_input["value"] if csrf_token_input else ""

            post_response = session.post(
                f"https://huggingface.co{action_name}",
                headers={
                    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
                },
                cookies=session.cookies,
                data={"csrf": csrf_token},
                timeout=10
            )

            if post_response.status_code == 200:
                print("[wake-up] 🚀 Restart triggered successfully. Waiting for the Space to be ready...")
                time.sleep(60)  # attendre que le Space redémarre
            else:
                print(f"[wake-up] ❌ Failed to restart Space: {post_response.status_code}")
        else:
            print("[wake-up] ✅ Space is already awake.")

    except Exception as e:
        if "403 Forbidden" in str(e):
            # This is ok, the space is already awake
            print("[wake-up] ✅ Space is already awake.")
        else:
            print(f"[wake-up] ❌ Error waking up space: {e}")

def get_wake_up_action_name(space_url: str) -> Optional[str]:
    space_pattern = 'https://([\\w]+)-([\\w-]+).hf.space'
    pattern = re.compile(space_pattern, flags=re.IGNORECASE)

    if match := pattern.match(space_url):
        # Extract the owner and the space name
        owner = match.group(1)
        space_name = match.group(2)

        return f'/spaces/{owner}/{space_name}/start'