import sys
import time
import os
import urllib

import requests

def setup_mlflow_environment(tracking_server_url: str) -> None:
    """
    Call this function to setup environment variables to allow you to interact with an
    MLFlow tracking server. It will print out a URL for the user to visit to authorize
    access.
    :param tracking_server_url:
    """
    os.environ['MLFLOW_TRACKING_URI'] = tracking_server_url
    os.environ['MLFLOW_TRACKING_TOKEN'] = get_mlflow_token(tracking_server_url)

def _get_keycloak_realm_and_client(mlflow_url: str):
    """
    NCSA Deployments of MLFlow are protected by the OAuth2-proxy sidecar. The OAuth2
    start endpoint redirects to our keycloak instance. Find this redirect and morph the
    redirect URL into a device flow start URL.
    :param mlflow_url:
    :return: A tuple with the device flow url and the mlflow instance's client_id
    """
    keycloak_query = requests.get(f"{mlflow_url}/oauth2/start", allow_redirects=False)
    keycloak_url = urllib.parse.urlparse(keycloak_query.next.url)
    if not keycloak_url.query:
        raise Exception("MLFLOW_TRACKING_URI malformed. Make sure there is no trailing slash")
    client_id = urllib.parse.parse_qs(keycloak_url.query)['client_id'][0]

    # Blank out the query args
    parts = list(keycloak_url[0:6])
    parts[4] = ''
    device_flow_start_url = urllib.parse.urlunparse(parts)
    return device_flow_start_url, client_id


def get_mlflow_token(tracking_server_url) -> str:
    """
    Kicks off a device flow by inspecting the MLFlow Tracking URI, requesting the user
    visit a constructed device flow url and waiting for them to approve the request.

    :return: The JWT generated by the device flow
    """
    realm_url, client_id = _get_keycloak_realm_and_client(tracking_server_url)

    r = requests.post(realm_url+"/device", data=[("client_id", client_id)])
    device_code = r.json()['device_code']
    print(f"Visit this link to authorize access to MLFlow instance {client_id}", file=sys.stderr)
    print(r.json()['verification_uri_complete'], file=sys.stderr)

    token_query_url = realm_url.replace("protocol/openid-connect/auth",
                                        "protocol/openid-connect/token")

    token = None
    while not token:
        time.sleep(5)
        r2 = requests.post(
            token_query_url,
            data=[
                ("client_id", client_id),
                ("grant_type", "urn:ietf:params:oauth:grant-type:device_code"),
                ("device_code", device_code)
            ])
        if r2.status_code==200:
            return r2.json()['access_token']
        elif r2.json()['error'] == 'authorization_pending':
            pass # Still waiting for user to complete login in web page
        else:
            raise Exception(r2.text)


def run():
    print(f"MLFLOW_TRACKING_TOKEN={get_mlflow_token(os.environ['MLFLOW_TRACKING_URI'])}")


if __name__ == '__main__':
    run()