import os
import json
import shutil
import tempfile
import uvicorn
from datetime import datetime
from flask import Flask, request, jsonify
from huggingface_hub import login, HfApi
from werkzeug.utils import secure_filename

app = Flask(__name__)

# Load environment variables
HF_TOKEN = os.getenv("HF_TOKEN")
REPO_ID = os.getenv("HF_REPO_ID")

# Initialize Hugging Face API
login(token=HF_TOKEN)
api = HfApi()

@app.route('/upload', methods=['POST'])
def upload_image():
    try:
        # Check if required data is present
        if 'image' not in request.files or 'guid' not in request.form or 'ip' not in request.form:
            return jsonify({"error": "Missing image, guid, or ip"}), 400

        image = request.files['image']
        guid = secure_filename(request.form['guid'])  # Sanitize GUID
        ip = request.form['ip']

        # Validate image
        if not image or not image.filename:
            return jsonify({"error": "No valid image provided"}), 400

        # Create temporary directory
        with tempfile.TemporaryDirectory() as temp_dir:
            # Create images folder
            temp_image_dir = os.path.join(temp_dir, "images")
            os.makedirs(temp_image_dir, exist_ok=True)

            # Save image with guid.jpg extension
            temp_image_path = os.path.join(temp_image_dir, f"{guid}.jpg")
            image.save(temp_image_path)

            # Create metadata
            path_in_repo = f"data/{guid}.jpg"
            metadata = {
                "file_name": path_in_repo,
                "guid": guid,
                "ip": ip,
                "upload_timestamp": datetime.utcnow().isoformat(),
            }

            # Write metadata to metadata.jsonl
            metadata_file = os.path.join(temp_dir, "metadata.jsonl")
            with open(metadata_file, "a") as f:
                f.write(json.dumps(metadata) + "\n")

            # Upload to Hugging Face
            api.upload_folder(
                folder_path=temp_dir,
                path_in_repo="",
                repo_id=REPO_ID,
                repo_type="dataset"
            )

        return jsonify({"message": "Image and metadata uploaded successfully"}), 200

    except Exception as e:
        return jsonify({"error": str(e)}), 500

def start_server(port: int = 7860):
    uvicorn.run(
        "app:app",
        host="0.0.0.0",
        port=port,
        workers=1,
        factory=False,
        # Use UvicornWorker to handle WSGI app
        worker_class="uvicorn.workers.UvicornWorker"
    )

if __name__ == "__main__":
    start_server()