from flask import Flask, request, jsonify
from flask_cors import CORS
import jwt
import time
import uuid
import requests
import os

app = Flask(__name__)
CORS(app, origins=["https://catptain-coze-api-01.hf.space"] + [
    "https://x-raremeta.com",
    "https://cybercity.top",
    "https://play-1.x-raremeta.com",
    "https://play.cybercity.top",
    "https://play.x-raremeta.com",
    "https://www.x-raremeta.com",
    "https://www.cybercity.top"
])

# 你的配置信息
CLIENT_ID = "1243934778935"
PRIVATE_KEY_FILE_PATH = "private_key.pem"
KID = "tlrohMMZyKMrrpP3GtxF_3_cerDhVIMINs0LOW91m7w"
VALIDATION_TOKEN = "cybercity2025"

def generate_jwt(client_id, private_key, kid):
    header = {
        "alg": "RS256",
        "typ": "JWT",
        "kid": kid
    }
    payload = {
        "iss": client_id,
        "aud": "api.coze.cn",
        "iat": int(time.time()),
        "exp": int(time.time()) + 3600,  # JWT 有效期为 1 小时
        "jti": uuid.uuid4().hex  # 防止重放攻击
        "connector_id": "4213660371992"  # 根据实际要求设置
        "user_id": "4213660371992"  # 根据实际要求设置
        "user": "4213660371992"  # 根据实际要求设置
    }
    return jwt.encode(payload, private_key, algorithm="RS256", headers=header)

def get_access_token(jwt_token):
    url = "https://api.coze.cn/api/permission/oauth2/token"
    data = {
        "duration_seconds": 86399,
        "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer"
    }
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {jwt_token}"
    }
    response = requests.post(url, json=data, headers=headers)
    return response.json()

# 添加根路由，帮助Hugging Face识别应用已经就绪
@app.route('/', methods=['GET'])
def index():
    return jsonify({"status": "Service is running", "endpoints": ["/get_token"]}), 200

# 正确的token获取路由
@app.route('/get_token', methods=['GET'])
def get_token_from_flask():
    auth_header = request.headers.get('Authorization')
    if auth_header != VALIDATION_TOKEN:
        return jsonify({"error": "Invalid authorization token"}), 401
    
    try:
        with open(PRIVATE_KEY_FILE_PATH, "r") as f:
            private_key = f.read()
        
        jwt_token = generate_jwt(CLIENT_ID, private_key, KID)
        response = get_access_token(jwt_token)
        
        if "access_token" in response:
            return jsonify({
                "access_token": response["access_token"],
                "expires_in": response["expires_in"]
            })
        else:
            return jsonify({"error": "Failed to get access token", "details": response}), 500
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# 使用环境变量设置端口
port = int(os.environ.get("PORT", 7860))
if __name__ == '__main__':
    app.run(host="0.0.0.0", port=port)