import json
from datetime import datetime
import aiohttp
from aiohttp import web

@app.route('/')
def index():
    return "flux2api with siliconflow", 200

async def handle_chat_completions(request):
    try:
        body = await 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 web.json_response({'error': 'Bad Request: Missing required fields'}, status=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
        }

        async with aiohttp.ClientSession() as session:
            async with session.post(new_url, json=new_request_body, headers={
                'accept': 'application/json',
                'content-type': 'application/json',
                'Authorization': request.headers.get('Authorization')
            }) as response:
                response_body = await response.json()

        image_url = response_body['images'][0]['url']
        unique_id = int(datetime.now().timestamp() * 1000)
        current_timestamp = int(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 web.Response(text=f"data: {data_string}\n\n", 
                                status=200, 
                                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)
                }
            }
            return web.json_response(response_payload)

    except Exception as error:
        return web.json_response({'error': f"Internal Server Error: {str(error)}"}, status=500)

app = web.Application()
app.router.add_post('/ai/v1/chat/completions', handle_chat_completions)


if __name__ == '__main__':
    web.run_app(app, port=8000)