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

request_count = {}
ongoing_requests = {}
bypass = os.getenv("BYPASS")
MAX_REQUESTS_PER_MINUTE = 7
app = Flask("DiscordRocks Instant Chat Proxy")

DISCORD_ROCKS_HOST = 'chat.discord.rocks'
CHAT_UI_HOST = 'cutycat2000x-instantchat.static.hf.space'

@app.route('/', methods=['GET', 'POST'])
@app.route("/<path:path>", methods=['GET', 'POST'])
def handle_request(path=""):
    chat_ui_url = f'https://{DISCORD_ROCKS_HOST}/{path}'
    ip =  request.remote_addr
    method = request.method
    data = request.data
    params = request.args
    headers = None
    didReply = False
    if method == 'POST':
        headers = {'Content-Type': 'application/json'}
        if 'reply' in path:
            if request_count.get(ip, 0) >= MAX_REQUESTS_PER_MINUTE:
                return jsonify({"error": "Too Many Requests. Please try again later."}), 429
            if ip in ongoing_requests:
                return jsonify({"error": "Concurrent requests not allowed."}), 429
            ongoing_requests[ip] = True
            request_count[ip] = request_count.get(ip, 0) + 1
            didReply = True
            json_data = request.get_json()
            json_data['bypass'] = bypass
            data = json.dumps(json_data)
    
    response = requests.request(method, chat_ui_url, params=params, data=data, headers=headers)
    
    proxied_response = Response(response.content)
    proxied_response.status_code = response.status_code
    proxied_response.headers["Content-Type"] = "text/html"
    if didReply:
        decrement_request_count_task(ip)
    
    return proxied_response

def decrement_request_count_task(ip):
    @copy_current_request_context
    def decrement_count():
        try:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            loop.run_until_complete(decrement_request_count(ip))
        except Exception as e:
            print(e)

    threading.Thread(target=decrement_count).start()

async def decrement_request_count(ip):
    try:
        await asyncio.sleep(60)  # Ensure this runs asynchronously
        request_count[ip] = request_count.get(ip, 0) - 1
        print(request_count[ip])
    except Exception as es:
        print(es)

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