from flask import Flask, request, jsonify
import requests
import time

app = Flask(__name__)

def get_token():
    url = "https://fluxaiweb.com/flux/getToken"
    response = requests.get(url)
    if response.status_code == 200:
        response_json = response.json()
        return response_json.get("data", {}).get("token")
    return None

def req_flux(token, prompt_value, aspect_ratio="1:1", output_format="webp", num_outputs=1, output_quality=90):
    url = "https://fluxaiweb.com/flux/generateImage"
    payload = {
        "prompt": prompt_value,
        "aspectRatio": aspect_ratio,
        "outputFormat": output_format,
        "numOutputs": num_outputs,
        "outputQuality": output_quality
    }
    headers = {
        'Content-Type': 'application/json',
        'token': token
    }
    try:
        response = requests.post(url, headers=headers, json=payload)
        response.raise_for_status()
        data = response.json()
        return data.get("data", {}).get("image")
    except requests.exceptions.RequestException as e:
        print(f"Error making request: {e}")
        return None

@app.route('/hf/v1/chat/completions', methods=['POST'])
def chat_completions():
    data = request.json
    messages = data.get('messages', [])
    
    # Extract the prompt from the last user message
    prompt = next((msg['content'] for msg in reversed(messages) if msg['role'] == 'user'), None)
    
    if not prompt:
        return jsonify({'error': 'No valid prompt provided'}), 400
    
    token = get_token()
    if not token:
        return jsonify({'error': 'Failed to get token'}), 500
    
    image_url = req_flux(token, prompt)
    if not image_url:
        return jsonify({'error': 'Failed to generate image'}), 500
    
    # Construct response in ChatCompletion format
    response = {
        "id": f"chatcmpl-{int(time.time())}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": "flux-ai-image-generator",
        "choices": [
            {
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": f"I've generated an image based on your prompt. You can view it here: {image_url}"
                },
                "finish_reason": "stop"
            }
        ],
        "usage": {
            "prompt_tokens": len(prompt.split()),
            "completion_tokens": 20,  # Approximate
            "total_tokens": len(prompt.split()) + 20
        }
    }
    
    return jsonify(response)

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