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

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.")

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}"

@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

        # Create a unique Space name and repo
        space_name = generate_space_name()
        full_repo_id = f"{hf_api.whoami(HF_TOKEN)['name']}/{space_name}"

        create_repo(
            repo_id=space_name,
            repo_type="space",
            space_sdk=space_type,
            token=HF_TOKEN,
            private=False
        )

        # 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
            upload_file(
                path_or_fileobj=f"temp_{filename}",
                path_in_repo=filename,
                repo_id=full_repo_id,
                repo_type="space",
                token=HF_TOKEN
            )
            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)
            upload_file(
                path_or_fileobj="temp_requirements.txt",
                path_in_repo="requirements.txt",
                repo_id=full_repo_id,
                repo_type="space",
                token=HF_TOKEN
            )
            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 5000
CMD ["python", "app.py"]
"""
            with open("temp_Dockerfile", "w") as f:
                f.write(default_dockerfile)
            upload_file(
                path_or_fileobj="temp_Dockerfile",
                path_in_repo="Dockerfile",
                repo_id=full_repo_id,
                repo_type="space",
                token=HF_TOKEN
            )
            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

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