app.py

from flask import Flask, request, jsonify, Response
import os
import json
from huggingface_hub import HfApi, create_repo, upload_file
import random
import string
from flask_httpauth import HTTPTokenAuth

# Import documentation from assembler_docs.py
from assembler_docs import DOCUMENTATION

app = Flask(__name__)

# Set up HTTP Token Authentication
auth = HTTPTokenAuth(scheme='Bearer')

# Hugging Face API token (set in Space settings or provided in request)
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("HF_TOKEN not set. Add it in Space settings.")

hf_api = HfApi()

# Authentication function to validate the token
@auth.verify_token
def verify_token(token):
    if not token:
        return False
    try:
        # Attempt to verify the token by checking if it can access the user's profile or Space-Share
        user_info = hf_api.whoami(token=token)
        # Optionally, check if the token has write access to Space-Share
        # This is a basic check; you might want to add more specific permission checks
        return bool(user_info and "Space-Share" in user_info.get('orgs', [])) or token == HF_TOKEN
    except Exception:
        return False

def generate_space_name():
    """Generate a unique Space name."""
    random_suffix = ''.join(random.choices(string.ascii_lowercase + string.digits, k=6))
    return f"GeneratedSpace-{random_suffix}"

@app.route('/create-space', methods=['POST'])
@auth.login_required
def create_hf_space():
    try:
        # Use the authenticated token for Hugging Face operations
        token = auth.current_user()
        hf_api = HfApi(token=token)

        # Parse JSON input
        data = request.get_json()
        if not data:
            return jsonify({"error": "No JSON data provided"}), 401

        # Extract parameters
        space_type = data.get("space_type", "gradio")  # Default to gradio if not specified
        files = data.get("files", {})  # Dictionary of filename: content
        params = data.get("parameters", {})  # Optional parameters

        if not files:
            return jsonify({"error": "No files provided in JSON"}), 401

        # Validate space_type
        valid_space_types = ["gradio", "static", "docker", "streamlit"]
        if space_type not in valid_space_types:
            return jsonify({"error": f"Invalid space_type. Must be one of {valid_space_types}"}), 401

        # Create a unique Space name and repo under Space-Share namespace
        space_name = generate_space_name()
        full_repo_id = f"Space-Share/{space_name}"

        try:
            # Attempt to create the repository
            repo_info = create_repo(
                repo_id=space_name,
                repo_type="space",
                space_sdk=space_type,
                token=token,
                private=False,
                exist_ok=True  # Allow creation even if the repo might exist (rare case)
            )
        except Exception as e:
            return jsonify({"error": f"Failed to create repository: {str(e)}"}), 500

        # Handle multi-file uploads
        for filename, content in files.items():
            # Write content to a temporary file
            with open(f"temp_{filename}", "w") as f:
                if filename.endswith(".py"):
                    # Inject parameters into Python files if present
                    content = f"PARAMS = {json.dumps(params)}\n\n{content}"
                f.write(content)

            # Upload to the new Space
            try:
                upload_file(
                    path_or_fileobj=f"temp_{filename}",
                    path_in_repo=filename,
                    repo_id=full_repo_id,
                    repo_type="space",
                    token=token
                )
            except Exception as e:
                os.remove(f"temp_{filename}")
                return jsonify({"error": f"Failed to upload file {filename}: {str(e)}"}), 500

            os.remove(f"temp_{filename}")

        # Add requirements.txt if not provided (basic defaults)
        if "requirements.txt" not in files:
            default_requirements = {
                "gradio": "gradio",
                "static": "",
                "docker": "flask",  # Example; adjust based on needs
                "streamlit": "streamlit"
            }.get(space_type, "")
            with open("temp_requirements.txt", "w") as f:
                f.write(default_requirements)
            try:
                upload_file(
                    path_or_fileobj="temp_requirements.txt",
                    path_in_repo="requirements.txt",
                    repo_id=full_repo_id,
                    repo_type="space",
                    token=token
                )
            except Exception as e:
                os.remove("temp_requirements.txt")
                return jsonify({"error": f"Failed to upload requirements.txt: {str(e)}"}), 500
            os.remove("temp_requirements.txt")

        # Special handling for Docker Spaces
        if space_type == "docker" and "Dockerfile" not in files:
            default_dockerfile = """
FROM python:3.10-slim
WORKDIR /app
COPY . .
RUN pip install -r requirements.txt
EXPOSE 7860
CMD ["python", "app.py"]
"""
            with open("temp_Dockerfile", "w") as f:
                f.write(default_dockerfile)
            try:
                upload_file(
                    path_or_fileobj="temp_Dockerfile",
                    path_in_repo="Dockerfile",
                    repo_id=full_repo_id,
                    repo_type="space",
                    token=token
                )
            except Exception as e:
                os.remove("temp_Dockerfile")
                return jsonify({"error": f"Failed to upload Dockerfile: {str(e)}"}), 500
            os.remove("temp_Dockerfile")

        space_url = f"https://huggingface.co/spaces/Space-Share/{space_name}"
        return jsonify({
            "message": "New Space created",
            "url": space_url,
            "note": "It may take a few minutes to build and deploy."
        }), 200

    except json.JSONDecodeError:
        return jsonify({"error": "Invalid JSON format"}), 401
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/docs', methods=['GET'])
@auth.login_required
def get_docs():
    """Return the API documentation as plain text."""
    return Response(DOCUMENTATION, mimetype='text/plain'), 200

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)