from urllib.request import urlopen, Request
from flask import Flask, Blueprint, request, Response
from flask_login import current_user

frontend = Blueprint(__name__.rpartition('.')[2], __name__)

def app():
    app = Flask(__name__)
    app.register_blueprint(frontend)
    return app

@frontend.route("/videoplayback")
def videoplayback():
    fvip = request.args.get("fvip", "3")
    mn = request.args.get("mn", "").split(",")[-1]
    host = f"r{fvip}---{mn}.googlevideo.com"

    # cargo-culted from invidious
    REQUEST_HEADERS_WHITELIST  = {"accept", "accept-encoding", "cache-control", "content-length", "if-none-match", "range"}
    RESPONSE_HEADERS_BLACKLIST = {"access-control-allow-origin", "alt-svc", "server"}

    # NOTE: don't close 'u' (or use 'with'), or reading from it will fail!
    # TODO: what to do / how to close socket when the user stops reading?
    req_headers = {
        k:v for k,v in request.headers
        if k.lower() in REQUEST_HEADERS_WHITELIST
    }
    u = urlopen(Request(f"https://{host}/videoplayback?{request.query_string.decode()}", headers=req_headers))

    headers = {
        k:v for k,v in u.headers.items()
        if k.lower() not in RESPONSE_HEADERS_BLACKLIST
    }
    return Response(iter(lambda: u.read(8192), b''), status=u.status, headers=headers)

if __name__ == '__main__':
    app().run(debug=True)