from flask import Flask, request, jsonify
import os
import requests # To forward requests to your original host

app = Flask(__name__)

# Retrieve secrets from environment variables
# This is the API key your *website* will send to this Hugging Face Space API
SPACE_API_KEY = os.environ.get("SPACE_API_KEY")
# This is the API key for your *original host*
ORIGINAL_HOST_API_KEY = os.environ.get("ORIGINAL_API_KEY")

@app.route("/chat/completions", methods=["POST"])
def chat_completions_proxy():
    # 1. Authenticate Request from Your Website
    client_api_key = request.headers.get("X-API-Key") # Or Authorization: Bearer <key>

    if not client_api_key or client_api_key != SPACE_API_KEY:
        return jsonify({"detail": "Unauthorized: Invalid Space API Key"}), 401

    # 2. Get data from your website's request
    try:
        data = request.json
        prompt = data.get("prompt")
        model = data.get("model")
        if not prompt or not model:
            return jsonify({"detail": "Missing 'prompt' or 'model' in request"}), 400
    except Exception as e:
        return jsonify({"detail": f"Invalid JSON or request body: {str(e)}"}), 400

    # 3. Forward Request to Your Original Host
    original_host_url = "https://your-original-host.com/api/chat" # Replace with your original API endpoint
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {ORIGINAL_HOST_API_KEY}" # Or whatever auth your original API uses
    }
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}] # Adapt to your original API's expected payload
    }

    try:
        response = requests.post(original_host_url, json=payload, headers=headers)
        response.raise_for_status() # Raise an exception for HTTP errors
        original_api_response = response.json()
        return jsonify(original_api_response), 200
    except requests.exceptions.RequestException as e:
        return jsonify({"detail": f"Error communicating with original host: {str(e)}"}), 500
    except Exception as e:
        return jsonify({"detail": f"An unexpected error occurred: {str(e)}"}), 500

if __name__ == "__main__":
    app.run(debug=os.environ.get("DEBUG", "False").lower() == "true", host="0.0.0.0", port=int(os.environ.get("PORT", 7860))) # Default Gradio port