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

app = Flask(__name__)

@app.route('/v1/chat/completions', methods=['POST'])
def chat_completions():
    auth_header = request.headers.get('Authorization', '')
    if auth_header.startswith('Bearer '):
        token = auth_header[7:]
    else:
        return jsonify({"error": "Authorization header must be a Bearer token"}), 401
    
    try:
        openai_request = request.json
        messages = openai_request.get('messages', [])
        
        user_message = ""
        for msg in reversed(messages):
            if msg.get('role') == 'user':
                user_message = msg.get('content', '')
                break
        
        if not user_message:
            return jsonify({"error": "No user message found"}), 400
        
        yun_api_url = 'https://ai.yun.139.com/api/outer/assistant/chat/add'
        
        yun_headers = {
            'User-Agent': 'Mozilla/5.0 (Linux; Android 14; 23049RAD8C Build/UKQ1.230804.001; wv) AppleWebKit/537.36 (KHTML, like Gecko) Version/4.0 Chrome/118.0.0.0 Mobile Safari/537.36 MCloudApp/11.4.4',
            'Accept': 'text/event-stream',
            'Content-Type': 'application/json',
            'sec-ch-ua': '"Chromium";v="118", "Android WebView";v="118", "Not=A?Brand";v="99"',
            'x-yun-app-channel': '101',
            'sec-ch-ua-mobile': '?1',
            'authorization': f'Basic {token}',
            'x-yun-api-version': 'v4',
            'sec-ch-ua-platform': '"Android"',
            'origin': 'https://yun.139.com',
            'x-requested-with': 'com.chinamobile.mcloud',
            'sec-fetch-site': 'same-site',
            'sec-fetch-mode': 'cors',
            'sec-fetch-dest': 'empty',
            'referer': 'https://yun.139.com/',
            'accept-language': 'zh,zh-CN;q=0.9,en-US;q=0.8,en;q=0.7'
        }
        
        yun_data = {
            "userId": "",
            "sessionId": "",
            "content": {
                "dialogue": user_message,
                "prompt": "",
                "commands": "",
                "resourceType": "0",
                "resourceId": "",
                "dialogueType": "0",
                "sourceChannel": "101",
                "extInfo": "{\"h5Version\":\"1.7.4\"}"
            },
            "applicationType": "chat",
            "applicationId": ""
        }
        
        stream_mode = openai_request.get('stream', False)
        
        if stream_mode:
            def generate():
                response_id = f"chatcmpl-{str(uuid.uuid4())}"
                start_time = time.time()
                
                response = requests.post(yun_api_url, headers=yun_headers, json=yun_data, stream=True)
                
                collected_text = ""
                for line in response.iter_lines():
                    if line:
                        line_text = line.decode('utf-8')
                        if line_text.startswith('data:'):
                            data_json = json.loads(line_text[5:].strip())
                            if 'content' in data_json:
                                delta_text = data_json['content']
                                collected_text += delta_text
                                
                                chunk = {
                                    "id": response_id,
                                    "object": "chat.completion.chunk",
                                    "created": int(time.time()),
                                    "model": "DeepSeek-R1",
                                    "choices": [
                                        {
                                            "index": 0,
                                            "delta": {
                                                "content": delta_text
                                            },
                                            "finish_reason": None
                                        }
                                    ]
                                }
                                
                                yield f"data: {json.dumps(chunk)}\n\n"
                
                final_chunk = {
                    "id": response_id,
                    "object": "chat.completion.chunk",
                    "created": int(time.time()),
                    "model": "DeepSeek-R1",
                    "choices": [
                        {
                            "index": 0,
                            "delta": {},
                            "finish_reason": "stop"
                        }
                    ]
                }
                yield f"data: {json.dumps(final_chunk)}\n\n"
                yield "data: [DONE]\n\n"
                
            return Response(stream_with_context(generate()), content_type='text/event-stream')
        
        else:
            response = requests.post(yun_api_url, headers=yun_headers, json=yun_data)
            response_data = response.json()
            
            assistant_message = ""
            if 'data' in response_data and 'content' in response_data['data']:
                assistant_message = response_data['data']['content']
            
            openai_response = {
                "id": f"chatcmpl-{str(uuid.uuid4())}",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": "DeepSeek-R1",
                "choices": [
                    {
                        "index": 0,
                        "message": {
                            "role": "assistant",
                            "content": assistant_message
                        },
                        "finish_reason": "stop"
                    }
                ],
                "usage": {
                    "prompt_tokens": len(user_message),
                    "completion_tokens": len(assistant_message),
                    "total_tokens": len(user_message) + len(assistant_message)
                }
            }
            
            return jsonify(openai_response)
    
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/v1/models', methods=['GET'])
def list_models():
    models = {
        "object": "list",
        "data": [
            {
                "id": "DeepSeek-R1",
                "object": "model",
                "created": int(time.time()),
                "owned_by": "139.com"
            }
        ]
    }
    return jsonify(models)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8080, debug=True)