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

_COOKIES = os.environ.get("COOKIES", "")
print(_COOKIES)

app = Flask(__name__)

@app.route('/v1/chat/completions', methods=['POST'])
def chat_completions():
    print("start")
    try:
        # 获取OpenAI格式的请求数据
        data = request.json
        
        # 生成唯一ID
        chat_id = str(uuid.uuid4()).replace('-', '')[:16]
        
        # 构建Akash格式的请求数据
        akash_data = {
            "id": chat_id,
            "messages": data.get('messages', []),
            "model": data.get('model', "DeepSeek-R1"),
            "system": data.get('system_message', "You are a helpful assistant."),
            "temperature": data.get('temperature', 0.6),
            "topP": data.get('top_p', 0.95)
        }
        # 构建请求头
        headers = {
            "Content-Type": "application/json",
            "Cookie": _COOKIES
        }
        print(headers)
        print(akash_data)
        _stream=data.get('stream', True)
        # 发送请求到Akash API
        response = requests.post(
            'https://chat.akash.network/api/chat',
            json=akash_data,
            headers=headers,
            stream=_stream
        )
        
        def generate():
            content_buffer = ""
            for line in response.iter_lines():
                if not line:
                    continue
                    
                try:
                    # 解析行数据，格式为 "type:json_data"
                    line_str = line.decode('utf-8')
                    msg_type, msg_data = line_str.split(':', 1)
                    
                    # 处理内容类型的消息
                    if msg_type == '0':
                        # 只去掉两边的双引号
                        if msg_data.startswith('"') and msg_data.endswith('"'):
                            msg_data = msg_data.replace('\\"', '"')
                            msg_data = msg_data[1:-1]
                        msg_data = msg_data.replace("\\n", "\n")
                        content_buffer += msg_data
                        
                        # 构建 OpenAI 格式的响应块
                        chunk = {
                            "id": f"chatcmpl-{chat_id}",
                            "object": "chat.completion.chunk",
                            "created": int(time.time()),
                            "model": data.get('model', "DeepSeek-R1"),
                            "choices": [{
                                "delta": {"content": msg_data},
                                "index": 0,
                                "finish_reason": None
                            }]
                        }
                        yield f"data: {json.dumps(chunk)}\n\n"
                    
                    # 处理结束消息
                    elif msg_type in ['e', 'd']:
                        chunk = {
                            "id": f"chatcmpl-{chat_id}",
                            "object": "chat.completion.chunk",
                            "created": int(time.time()),
                            "model": data.get('model', "DeepSeek-R1"),
                            "choices": [{
                                "delta": {},
                                "index": 0,
                                "finish_reason": "stop"
                            }]
                        }
                        yield f"data: {json.dumps(chunk)}\n\n"
                        yield "data: [DONE]\n\n"
                        break
                        
                except Exception as e:
                    print(f"Error processing line: {e}")
                    continue

        if _stream == True:
            return Response(
                generate(),
                mimetype='text/event-stream',
                headers={
                    'Cache-Control': 'no-cache',
                    'Connection': 'keep-alive',
                    'Content-Type': 'text/event-stream'
                }
            )
        else:
            result_json = response.json()
            return Response(
                response=json.dumps(result_json),
                status=response.status_code,
                headers={
                    'Cache-Control': 'no-cache',
                    'Connection': 'keep-alive',
                    'Content-Type': 'application/json'
                }
            )
    
    except Exception as e:
        return jsonify({"error": str(e)}), 500

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