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

# 1. 初始化Flask应用
app = Flask(__name__)

# 2. gpt-oss.com API的固定配置
GPT_OSS_API_URL = "https://api.gpt-oss.com/chatkit"
GPT_OSS_HEADERS = {
    'authority': 'api.gpt-oss.com',
    'accept': 'text/event-stream',
    'content-type': 'application/json',
    'origin': 'https://gpt-oss.com',
    'referer': 'https://gpt-oss.com/',
    'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36',
    'x-selected-model': 'gpt-oss-120b',
}

# 3. 新增：根路由，用于健康检查和显示欢迎信息
@app.route('/', methods=['GET'])
def root():
    return jsonify({
        "message": "欢迎使用 GPT-OSS to OpenAI 格式代理API",
        "status": "ok",
        "api_endpoint": "/v1/chat/completions"
    })

# 4. 核心：创建OpenAI兼容的API端点
@app.route('/v1/chat/completions', methods=['POST'])
def chat_completions_proxy():
    try:
        openai_request_data = request.json
        messages = openai_request_data.get("messages", [])
        user_prompt = next((m['content'] for m in reversed(messages) if m.get('role') == 'user'), None)
        
        if not user_prompt:
            return jsonify({"error": "在请求中未找到用户消息。"}), 400
            
        stream_requested = openai_request_data.get("stream", False)
    except Exception as e:
        return jsonify({"error": f"请求格式无效: {e}"}), 400

    random_user_id = str(uuid.uuid4())
    gpt_oss_cookies = {'user_id': random_user_id}
    gpt_oss_payload = {
        "op": "threads.create",
        "params": {"input": {"text": user_prompt, "content": [{"type": "input_text", "text": user_prompt}]}}
    }
    
    def generate_stream():
        try:
            with requests.post(
                GPT_OSS_API_URL,
                headers=GPT_OSS_HEADERS,
                cookies=gpt_oss_cookies,
                json=gpt_oss_payload,
                stream=True,
                timeout=120
            ) as response:
                response.raise_for_status()

                for line in response.iter_lines():
                    if line:
                        line_str = line.decode('utf-8')
                        if line_str.startswith('data: '):
                            json_data_str = line_str[6:]
                            try:
                                gpt_oss_data = json.loads(json_data_str)
                                
                                event_type = gpt_oss_data.get('type')
                                if (event_type == 'thread.item_updated' and 
                                    gpt_oss_data.get('update', {}).get('type') == 'assistant_message.content_part.text_delta'):
                                    
                                    delta_content = gpt_oss_data['update'].get('delta', '')
                                    openai_chunk = {
                                        "id": f"chatcmpl-{str(uuid.uuid4())}",
                                        "object": "chat.completion.chunk",
                                        "created": int(time.time()),
                                        "model": "gpt-oss-120b",
                                        "choices": [{"index": 0, "delta": {"content": delta_content}, "finish_reason": None}]
                                    }
                                    yield f"data: {json.dumps(openai_chunk)}\n\n"

                            except json.JSONDecodeError:
                                continue

                yield "data: [DONE]\n\n"
        except requests.exceptions.RequestException as e:
            error_chunk = {"error": f"与后端服务通信失败: {e}"}
            yield f"data: {json.dumps(error_chunk)}\n\n"

    if stream_requested:
        return Response(generate_stream(), mimetype='text/event-stream')
    else:
        return jsonify({"error": "非流式响应目前不受支持，请在请求中设置 'stream': true"}), 501

# 5. 启动应用
if __name__ == '__main__':
    # 使用 gunicorn 部署时不会执行这里，但在本地测试时会用到
    app.run(host='0.0.0.0', port=7860)