#!env python3
""" Flask with OAuth:

This application provides end-to-end demonstration of Databricks SDK for Python
capabilities of OAuth Authorization Code flow with PKCE security enabled. This
can help you build a hosted app with every user using their own identity to
access Databricks resources.

If you have already Custom App:

./flask_app_with_oauth.py --host <databricks workspace url> \
    --client_id <app-client-id> \
    --client_secret <app-secret> \
    --port 5001

If you want this script to register Custom App and redirect URL for you:

./flask_app_with_oauth.py --port 5001 --profile <databricks account profile>

You'll get prompted for Databricks Account username and password for
script to enroll your account into OAuth and create a custom app with
http://localhost:5001/callback as the redirect callback. Client and
secret credentials for this OAuth app will be printed to the console,
so that you could resume testing this app at a later stage.

Once started, please open http://localhost:5001 in your browser and
go through SSO flow to get a list of clusters on <databricks workspace url>.
"""

import argparse
import logging
import sys

from databricks.sdk.oauth import OAuthClient, get_workspace_endpoints
from databricks.sdk.service.compute import ListClustersFilterBy, State

APP_NAME = "flask-demo"
all_clusters_template = """<ul>
{% for cluster in clusters -%}
    <li><a 
        target="_blank" 
        href="{{ workspace_host }}/#setting/clusters/{{ cluster.cluster_id }}/configuration">
        {{ cluster.cluster_name }}</a> is {{ cluster.state }}</li>
{% endfor %}
</ul>"""


def create_flask_app(workspace_host: str, client_id: str, client_secret: str):
    """The create_flask_app function creates a Flask app that is enabled with OAuth.

    It initializes the app and web session secret keys with a randomly generated token. It defines two routes for
    handling the callback and index pages.
    """
    import secrets

    from flask import (Flask, redirect, render_template_string, request,
                       session, url_for)

    app = Flask(APP_NAME)
    app.secret_key = secrets.token_urlsafe(32)

    @app.route("/callback")
    def callback():
        """The callback route initiates consent using the OAuth client, exchanges
        the callback parameters, and redirects the user to the index page."""
        from databricks.sdk.oauth import Consent

        consent = Consent.from_dict(session["consent"], client_secret=client_secret)
        session["creds"] = consent.exchange_callback_parameters(request.args).as_dict()
        return redirect(url_for("index"))

    @app.route("/")
    def index():
        """The index page checks if the user has already authenticated and retrieves the user's credentials using
        the Databricks SDK WorkspaceClient. It then renders the template with the clusters' list.
        """
        oidc_endpoints = get_workspace_endpoints(workspace_host)
        port = request.environ.get("SERVER_PORT")
        redirect_url = f"http://localhost:{port}/callback"
        if "creds" not in session:
            oauth_client = OAuthClient(
                oidc_endpoints=oidc_endpoints,
                client_id=client_id,
                client_secret=client_secret,
                redirect_url=redirect_url,
            )
            consent = oauth_client.initiate_consent()
            session["consent"] = consent.as_dict()
            return redirect(consent.authorization_url)

        from databricks.sdk import WorkspaceClient
        from databricks.sdk.oauth import SessionCredentials

        credentials_strategy = SessionCredentials.from_dict(
            session["creds"],
            token_endpoint=oidc_endpoints.token_endpoint,
            client_id=client_id,
            client_secret=client_secret,
            redirect_url=redirect_url,
        )
        workspace_client = WorkspaceClient(
            host=workspace_host,
            product=APP_NAME,
            credentials_strategy=credentials_strategy,
        )
        clusters = workspace_client.clusters.list(
            filter_by=ListClustersFilterBy(cluster_states=[State.RUNNING, State.PENDING])
        )
        return render_template_string(
            all_clusters_template,
            workspace_host=workspace_host,
            clusters=clusters,
        )

    return app


def register_custom_app(args: argparse.Namespace) -> tuple[str, str]:
    """Creates new Custom OAuth App in Databricks Account"""
    logging.info("No OAuth custom app client/secret provided, creating new app")

    from databricks.sdk import AccountClient

    account_client = AccountClient(profile=args.profile)

    custom_app = account_client.custom_app_integration.create(
        name=APP_NAME,
        redirect_urls=[
            f"http://localhost:{args.port}/callback",
        ],
        confidential=True,
        scopes=["all-apis"],
    )
    logging.info(
        f"Created new custom app: " f"--client_id {custom_app.client_id} " f"--client_secret {custom_app.client_secret}"
    )

    return custom_app.client_id, custom_app.client_secret


def parse_arguments() -> argparse.Namespace:
    """Parses arguments for this demo"""
    parser = argparse.ArgumentParser(prog=APP_NAME, description=__doc__.strip())
    parser.add_argument("--host")
    for flag in ["client_id", "client_secret"]:
        parser.add_argument(f"--{flag}")
    parser.add_argument("--port", default=5001, type=int)
    parser.add_argument(
        "--profile",
        default="DEFAULT",
        help="Databricks account profile to use for authentication.",
    )
    return parser.parse_args()


if __name__ == "__main__":
    logging.basicConfig(
        stream=sys.stdout,
        level=logging.INFO,
        format="%(asctime)s [%(name)s][%(levelname)s] %(message)s",
    )
    logging.getLogger("databricks.sdk").setLevel(logging.DEBUG)

    args = parse_arguments()
    client_id, client_secret = args.client_id, args.client_secret
    if not client_id:
        client_id, client_secret = register_custom_app(args)
    app = create_flask_app(args.host, client_id, client_secret)

    app.run(
        host="localhost",
        port=args.port,
        debug=True,
        # to simplify this demo experience, we create OAuth Custom App for you,
        # but it intervenes with the werkzeug reloader. So we disable it
        use_reloader=args.client_id is not None,
    )