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

@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 under Space-Share namespace
        space_name = generate_space_name()
        full_repo_id = f"Space-Share/{space_name}"

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

        # 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
        except Exception as e:
            return jsonify({"error": f"Failed to verify repository {full_repo_id}: {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=HF_TOKEN
                )
                print(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:
                upload_file(
                    path_or_fileobj="temp_requirements.txt",
                    path_in_repo="requirements.txt",
                    repo_id=full_repo_id,
                    repo_type="space",
                    token=HF_TOKEN
                )
                print("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:
                upload_file(
                    path_or_fileobj="temp_Dockerfile",
                    path_in_repo="Dockerfile",
                    repo_id=full_repo_id,
                    repo_type="space",
                    token=HF_TOKEN
                )
                print("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/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"}), 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)