from flask import Flask, jsonify, request, Response, stream_with_context
from flask_cors import CORS
import requests
from bs4 import BeautifulSoup
import os
import re
import urllib.parse
import time
import random
import base64
from io import BytesIO
from googlesearch import search
import logging
import queue
from huggingface_hub import HfApi


app = Flask(__name__)
# Enable CORS with specific settings
CORS(app, resources={
    r"/*": {
        "origins": "*",
        "methods": ["GET", "POST", "OPTIONS"],
        "allow_headers": ["Content-Type", "Authorization"]
    }
})


HF_TOKEN = os.getenv("HF_TOKEN")  # Make sure you set the HF_TOKEN in your environment

@app.route('/restart_space', methods=['POST'])
def api_restart_space():
    """API route to restart a Hugging Face Space."""
    space_id = 'Pamudu13/web-scraper'  
    factory_reboot = request.json.get('factory_reboot', False)  # Optional: Set to True if you want a factory reboot

    if not space_id:
        return jsonify({'error': 'space_id parameter is required'}), 400

    try:
        hfapi = HfApi()

        # Call the restart_space method
        res = hfapi.restart_space(
            space_id,
            token=HF_TOKEN,
            factory_reboot=factory_reboot
        )

        return jsonify({
            'success': True,
            'message': f"Successfully restarted Space: {space_id}",
            'response': res
        }), 200

    except Exception as e:
        return jsonify({
            'success': False,
            'message': f"Error: {str(e)}"
        }), 500    

@app.route('/get_live_space_status', methods=['GET'])
def get_live_space_status():
    """API route to stream live status of a Hugging Face Space."""
    space_id = request.args.get('space_id', 'Pamudu13/web-scraper')  # Default to 'Pamudu13/web-scraper' if not provided

    def generate():
        while True:
            try:
                # Fetch the current runtime status of the Space
                hf_api = HfApi()
                space_runtime = hf_api.get_space_runtime(repo_id=space_id)
                
                # Extract relevant details
                status = space_runtime.stage  # e.g., 'BUILDING', 'RUNNING', etc.
                hardware = space_runtime.hardware  # e.g., 'cpu-basic', 't4-medium', etc.
                
                # Send the status as a Server-Sent Event
                yield f"data: {status}\n\n"
                yield f"data: {hardware}\n\n"

                # Delay before checking the status again
                time.sleep(5)  # Adjust polling interval as needed

            except Exception as e:
                # Handle errors and send an error message
                yield f"data: Error: {str(e)}\n\n"
                break  # Stop the stream in case of an error

    return Response(stream_with_context(generate()), mimetype='text/event-stream')


if __name__ == '__main__':
    logger.info("Starting Flask API server...")
    app.run(host='0.0.0.0', port=5001, debug=True)