from flask import Flask, request, Response
import requests
import os

app = Flask(__name__)

# 代理的目标域名
TARGET_DOMAIN = "https://generativelanguage.googleapis.com"

@app.route('/', defaults={'path': ''}, methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
@app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
def proxy(path):
    """
    代理所有请求到目标域名。
    """
    target_url = f"{TARGET_DOMAIN}/{path}?{request.query_string.decode()}"

    try:
        # 构造请求头
        headers = dict(request.headers)
        headers.pop('Host', None)
        headers.pop('Content-Length', None)

        # 发送请求
        response = requests.request(
            method=request.method,
            url=target_url,
            headers=headers,
            data=request.get_data(),
            stream=True,
            allow_redirects=False
        )

        # 获取响应内容的编码，默认使用 UTF-8
        encoding = response.encoding or 'utf-8'
        content = response.content.decode(encoding, errors='replace')

        # 构造响应
        excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
        response_headers = [(name, value) for (name, value) in response.headers.items()
                            if name.lower() not in excluded_headers]

        # 检查原始响应的 Content-Type，如果包含 json，则设置为 application/json
        original_content_type = response.headers.get('Content-Type', 'text/plain')
        if 'application/json' in original_content_type.lower():
            content_type = 'application/json; charset=utf-8'
        else:
            content_type = f'text/plain; charset={encoding}'

        proxy_response = Response(content, response.status_code, headers=response_headers)
        proxy_response.headers['Content-Type'] = content_type

        # 添加 CORS 头
        proxy_response.headers['Access-Control-Allow-Origin'] = '*'
        proxy_response.headers['Access-Control-Allow-Methods'] = 'GET, POST, OPTIONS, PUT, DELETE, PATCH'
        proxy_response.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization'
        proxy_response.headers['Access-Control-Allow-Credentials'] = 'true'

        return proxy_response

    except requests.exceptions.RequestException as e:
        print(f"代理请求出错: {e}")
        return Response(f"代理服务器出错: {e}", status=500)

    except Exception as e:
        print(f"其他错误: {e}")
        return Response(f"代理服务器出错: {e}", status=500)

@app.after_request
def after_request(response):
    """
    处理 OPTIONS 请求，添加 CORS 头。
    """
    if request.method == 'OPTIONS':
        response.headers['Access-Control-Allow-Origin'] = '*'
        response.headers['Access-Control-Allow-Methods'] = 'GET, POST, OPTIONS, PUT, DELETE, PATCH'
        response.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization'
        response.headers['Access-Control-Allow-Credentials'] = 'true'
    return response

if __name__ == '__main__':
    port = int(os.environ.get("PORT", 7860))
    print(f"Flask 服务器已启动，监听 {port} 端口")
    app.run(debug=True, host='0.0.0.0', port=port)