# Your url for ST will be something like this:
# https://username-spaceName.hf.space/proxy/anthropic/v1
# Streaming is not supported

import logging
import json
from flask import Flask, request, Response
import requests

app = Flask(__name__)

# Configure the logging level
app.logger.setLevel(logging.INFO)

@app.route('/', defaults={'path': ''})
@app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'PATCH', 'DELETE'])
def relay_request(path):
    # Extract the base URL to relay requests
    base_url = 'https://claude.ai'  # Replace with the URL you want to relay requests to

    # Prepare the relay URL by appending the path from the original request
    relay_url = base_url + '/' + path

    app.logger.info(request.headers)

    # If the request is an EventStream request, use the stream method
    if request.headers.get('Content-Type') == 'text/event-stream':
        response = requests.request(
            method=request.method,
            url=relay_url,
            headers=request.headers,
            allow_redirects=False,
            stream=True  # Enable streaming response
        )
        headers = [(name, value) for name, value in response.raw.headers.items()]

        def generate():
            for chunk in response.iter_content(chunk_size=8192):
                yield chunk

        return Response(generate(), response.status_code, headers)

    # Remove the original IP from the headers
    headers = {k: v for k, v in request.headers.items() if k != 'X-Forwarded-For'}

    # If the request includes JSON data, relay it as JSON
    if 'application/json' in request.headers.get('Content-Type', ''):
        data = request.get_json()
        response = requests.request(
            method=request.method,
            url=relay_url,
            headers=headers,
            json=data,
            allow_redirects=False
        )
    else:
        response = requests.request(
            method=request.method,
            url=relay_url,
            headers=headers,
            data=request.data,
            allow_redirects=False
        )

    # Prepare the response to return to the original API call
    headers = [(name, value) for name, value in response.headers.items() if name != 'Transfer-Encoding']
    return Response(response.content, response.status_code, headers)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)