#!/usr/bin/env python
# Copyright (c) TripleBlind Holdings, Inc. Confidential and Proprietary. All rights reserved.

import tripleblind as tb


tb.util.set_script_dir_current()
tb.initialize(api_token=tb.config.example_user1["token"])

# Create an asset on the Router that exposes a projected view of data hosted
# in a MongoDB database.
try:
    print("Creating view into college acceptance dataset Access Point...")
    USER = ""
    PASSWORD = ""
    MONGO_HOST = ""
    MONGO_PORT = 27017
    CONNECTION_STR = f"mongodb://{USER}:{PASSWORD}@{MONGO_HOST}:{MONGO_PORT}/"
    asset0 = tb.asset.MongoDatabase.create(
        connection_str=CONNECTION_STR,
        query={},  # Empty query returns all records
        name=f"college-acceptance-mongo",
        desc="College acceptance dataset stored in a MongoDB database.",
        projection={
            "gmat": 1,
            "gpa": 1,
            "work_experience": 1,
            "age": 1,
            "admitted": 1,
            "_id": 0,
        },
        is_discoverable=True,
        database="datasets",
        collection="college_acceptance",
    )

    # The store creates an Agreement with corporate to allow access to this
    # specific view of data only for training the recommendation model.
    asset0.add_agreement(
        with_team=tb.config.example_user2["team_id"], operation=tb.Operation.BLIND_QUERY
    )
    asset0.add_agreement(
        with_team=tb.config.example_user2["team_id"],
        operation=tb.Operation.XGBOOST_TRAIN,
    )
    print("Created Agreement")
except tb.TripleblindAssetAlreadyExists:
    print("Asset already exists.")
