# Dremio 25 only
def get_nessie_token(tokenendpoint, password):
    import requests

    # Exchange Dremio PAT to Nessie token
    token_request_body = {
        "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
        "scope": "dremio.all",
        "subject_token_type": "urn:ietf:params:oauth:token-type:dremio:personal-access-token",
        "subject_token": password,
    }
    x = requests.post(tokenendpoint, data=token_request_body)
    x.raise_for_status()
    return x.json()["access_token"]


def get_spark_session(host, nessie_token):
    import stackit_spark

    catalog_name_in_spark = "stackit"
    return stackit_spark.get_spark(
        additional_config={
            "spark.jars.packages": "org.apache.iceberg:iceberg-spark-runtime-3.5_2.12:1.6.1,org.apache.iceberg:iceberg-aws-bundle:1.6.1",
            f"spark.sql.catalog.{catalog_name_in_spark}": "org.apache.iceberg.spark.SparkCatalog",
            f"spark.sql.catalog.{catalog_name_in_spark}.type": "rest",
            f"spark.sql.catalog.{catalog_name_in_spark}.uri": host,
            f"spark.sql.catalog.{catalog_name_in_spark}.token": nessie_token,
        }
    )


if __name__ == "__main__":
    get_nessie_token(
        "https://dremio-internal.data-platform-dev.stackit.run/oauth/token",
        "xxxxx",
    )