app.py

from flask import Flask, request, jsonify, Response
import os
import json
from huggingface_hub import HfApi, create_repo, upload_file, login
import random
import string
import logging
import black

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Import documentation from assembler_docs.py
from assembler_docs import DOCUMENTATION

app = Flask(__name__)

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

# Log in to Hugging Face Hub with the token
try:
    login(token=HF_TOKEN)
    logger.info("Successfully logged in to Hugging Face Hub")
except Exception as e:
    raise ValueError(f"Failed to log in to Hugging Face Hub: {str(e)}")

hf_api = HfApi()

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}"
    return "auto_space_1"
def verify_organization_access(organization: str, token: str) -> bool:
    """Verify if the token has write access to the organization."""
    try:
        # Try to list repositories or check organization membership
        repos = hf_api.list_repos(author=organization, token=token)
        logger.info(f"Verified access to organization {organization}")
        return True
    except Exception as e:
        logger.error(f"No write access to organization {organization}: {str(e)}")
        return False

@app.route('/create-space', methods=['POST'])
def create_hf_space():
    try:
        # Parse JSON input
        data = request.get_json()
        if not data:
            return jsonify({"error": "No JSON data provided"}), 400

        # 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"}), 400

        # 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}"}), 400

        # Organization namespace for new Spaces
        ORGANIZATION = "broadfield-dev"
        space_name = generate_space_name()
        full_repo_id = f"{ORGANIZATION}/{space_name}"

        # Verify token has access to Space-Share organization
        #if not verify_organization_access(ORGANIZATION, HF_TOKEN):
        #    return jsonify({"error": f"No write access to organization {ORGANIZATION}. Check your HF_TOKEN permissions."}), 403
        api=HfApi(endpoint="https://huggingface.co", token=HF_TOKEN)
        try:
            # Attempt to create the repository with explicit organization context
            repo_info = api.create_repo(
                repo_id=full_repo_id,
                repo_type="space",
                space_sdk=space_type,
                private=False,
                exist_ok=True  # Allow creation even if the repo might exist
            )
            logger.info(f"Repository created: {full_repo_id}")
        except Exception as e:
            #return jsonify({"error": f"Failed to create repository {full_repo_id}: {str(e)}"}), 500
            print(f"Failed to create repository {full_repo_id}: {str(e)}")

        # Verify repository existence before uploading
        try:
            repo_exists = hf_api.repo_exists(repo_id=full_repo_id, repo_type="space")
            if not repo_exists:
                return jsonify({"error": f"Repository {full_repo_id} does not exist after creation attempt."}), 500
            logger.info(f"Verified repository exists: {full_repo_id}")
        except Exception as e:
            return jsonify({"error": f"Failed to verify repository {full_repo_id}: {str(e)}"}), 500

        # Handle multi-file uploads
        print(str(files))
        for filename, content in files.items():
            #content_lines = content.split("\n")
            #content = ""
            #for line in content_lines:
            #    content += line + "\n"
            # Write content to a temporary file
            output=""
            with open(f"temp_{filename}", "w") as f:
                '''content=""
                if filename.endswith(".py"):
                    # Inject parameters into Python files if present
                    #content = f"PARAMS = {json.dumps(params)}\n\n{content}"
                    content_box = list(content.split("\n"))
                    for line in content_box:
                        f.write(line + "\n")
                else:'''
                print(content)
                content_box = list(content.replace("\\n", "\n").split("\n"))
                print(content_box)
                for line in content_box:
                    output+=line + "\n"
                    
                    
                f.write(output)
            f.close()
            
            #if filename.endswith(".py"):
            #    os.system(f"black temp_{filename}")
            # Upload to the new Space
            try:
                api.upload_file(
                    path_or_fileobj=f"temp_{filename}",
                    path_in_repo=filename,
                    repo_id=full_repo_id,
                    repo_type="space",
                )
                logger.info(f"Uploaded file: {filename}")
            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:
                api.upload_file(
                    path_or_fileobj="temp_requirements.txt",
                    path_in_repo="requirements.txt",
                    repo_id=full_repo_id,
                    repo_type="space",
                    token=HF_TOKEN
                )
                logger.info("Uploaded requirements.txt")
            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:
                api.upload_file(
                    path_or_fileobj="temp_Dockerfile",
                    path_in_repo="Dockerfile",
                    repo_id=full_repo_id,
                    repo_type="space",
                    token=HF_TOKEN
                )
                logger.info("Uploaded Dockerfile")
            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/{full_repo_id}"
        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"}), 400
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/docs', methods=['GET'])
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)