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

app = Flask(__name__)

@app.route('/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
        }

        # 从传入的 Authorization 头中随机选择一个 token
        authorization_header = request.headers.get('Authorization')
        if authorization_header:
            # 去掉 "Bearer " 前缀并分割 token
            tokens = authorization_header.replace("Bearer ", "").split(',')
            if len(tokens) > 1:
                selected_token = random.choice(tokens).strip()
            else:
                selected_token = tokens[0].strip()
            # 重新格式化为 "Bearer 随机选择的token"
            selected_token = f"Bearer {selected_token}"
        else:
            return jsonify({"error": "Unauthorized: Missing Authorization header"}), 401

        headers = {
            'accept': 'application/json',
            'content-type': 'application/json',
            'Authorization': selected_token
        }

        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)