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

app = Flask(__name__)

# 代理的目标域名
TARGET_DOMAIN = os.environ.get("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)
       # 删除 Host 头，防止目标服务器出错
       headers.pop('Host', None)
       # 删除 Content-Length，让 requests 自动处理
       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 # 禁止重定向，让客户端处理
       )


       # 构造响应
       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]

       proxy_response = Response(response.raw.read(),  # 使用 raw.read() 读取流式响应
                                 response.status_code,
                                 headers=response_headers)

       # 添加 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)