import gradio.networking
import gradio.queueing

from auth import attach_oauth


# 1. Patch => attach auth before starting app

old_create_app = gradio.networking.App.create_app


def patched_create_app(*args, **kwargs):
    app = old_create_app(*args, **kwargs)
    attach_oauth(app)
    return app


gradio.networking.App.create_app = patched_create_app

# 2. Patch => forward session info when using websockets (i.e. when queue is enabled which is the default on Spaces)

old_get_request_params = gradio.queueing.Queue.get_request_params


def new_get_request_params(self, websocket):
    params = old_get_request_params(self, websocket)
    params["session"] = websocket.session  # Forward session info to the internal request
    return params


gradio.queueing.Queue.get_request_params = new_get_request_params