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

app = Flask(__name__)

@app.route('/')
def index():
    return "text-to-image with siliconflow", 200
    
@app.route('/ai/v1/chat/completions', methods=['POST'])
def handle_request():
    try:
        body = request.json
        model = body.get('model')
        messages = body.get('messages')
        stream = body.get('stream', False)

        if not model or not messages or len(messages) == 0:
            return jsonify({"error": "Bad Request: Missing required fields"}), 400

        prompt = messages[-1]['content']
        new_url = f'https://api.siliconflow.cn/v1/{model}/text-to-image'

        new_request_body = {
            "prompt": prompt,
            "image_size": "1024x1024",
            "batch_size": 1,
            "num_inference_steps": 4,
            "guidance_scale": 1
        }

        headers = {
            'accept': 'application/json',
            'content-type': 'application/json',
            'Authorization': request.headers.get('Authorization')
        }

        response = requests.post(new_url, headers=headers, json=new_request_body)
        response_body = response.json()

        image_url = response_body['images'][0]['url']
        unique_id = int(time.time() * 1000)
        current_timestamp = unique_id // 1000

        if stream:
            response_payload = {
                "id": unique_id,
                "object": "chat.completion.chunk",
                "created": current_timestamp,
                "model": model,
                "choices": [
                    {
                        "index": 0,
                        "delta": {
                            "content": f"![]({image_url})"
                        },
                        "finish_reason": "stop"
                    }
                ]
            }
            data_string = json.dumps(response_payload)
            return Response(f"data: {data_string}\n\n", content_type='text/event-stream')
        else:
            response_payload = {
                "id": unique_id,
                "object": "chat.completion",
                "created": current_timestamp,
                "model": model,
                "choices": [
                    {
                        "index": 0,
                        "message": {
                            "role": "assistant",
                            "content": f"![]({image_url})"
                        },
                        "logprobs": None,
                        "finish_reason": "length"
                    }
                ],
                "usage": {
                    "prompt_tokens": len(prompt),
                    "completion_tokens": len(image_url),
                    "total_tokens": len(prompt) + len(image_url)
                }
            }
            data_string = json.dumps(response_payload)
            return Response(f"{data_string}\n\n", content_type='text/event-stream')

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

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