import requests
import io
import logging
from PIL import Image
from flask import Flask, request, jsonify, send_file
from functools import wraps
from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry

API_URL = "https://api-inference.huggingface.co/models/Paresh1879/stable-diffusion-xl-thumbsup-extend"
API_KEY = "your-api-key"

app = Flask(__name__)

# Storage for API key usage counts
api_key_usage = {}


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

# Configure HTTP Session with Retry Logic
session = requests.Session()
retry = Retry(
    total=10, 
    backoff_factor=2, 
    status_forcelist=[429, 500, 502, 503, 504],
    allowed_methods=["HEAD", "GET", "OPTIONS", "POST"]
)
adapter = HTTPAdapter(max_retries=retry)
session.mount("http://", adapter)
session.mount("https://", adapter)

def query(payload):
    try:
        response = session.post(API_URL, headers={"Authorization": f"Bearer {API_KEY}"}, json=payload)
        response.raise_for_status() 
        return response.content
    except requests.exceptions.RequestException as e:
        logger.error(f"Request failed: {e}")
        raise
# checking api_key
def require_api_key(f):
    @wraps(f)
    def check_api_key(*args, **kwargs):
        api_key = request.headers.get('Authorization')
        if api_key != f"Bearer {API_KEY}":
            return jsonify({"error": "Unauthorized"}), 401
        if api_key not in api_key_usage:
            api_key_usage[api_key] = 0
        api_key_usage[api_key] += 1
        return f(*args, **kwargs)
    return check_api_key

# /generate_image endpoint
@app.route('/generate_image', methods=['POST'])
@require_api_key
def generate_image():
    prompt = request.json.get('prompt')
    if not prompt:
        return jsonify({"error": "Missing prompt"}), 400

    try:
        image_bytes = query({"inputs": prompt})
        image = Image.open(io.BytesIO(image_bytes))
        img_io = io.BytesIO()
        image.save(img_io, 'PNG')
        img_io.seek(0)
        return send_file(img_io, mimetype='image/png')
    except requests.exceptions.RequestException as e:
        return jsonify({"error": "Service is temporarily unavailable. Please try again later."}), 503
    
# api_key_usage endpoint
@app.route('/api_key_usage', methods=['GET'])
@require_api_key
def get_api_key_usage():
    api_key = request.headers.get('Authorization')
    usage_count = api_key_usage.get(api_key, 0)
    return jsonify({"usage_count": usage_count})

if __name__ == '__main__':
    app.run(debug=True)