#!/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()
# Unique value used by all scripts in this folder.  Edit "run_id.txt" to change
run_id = tb.util.read_run_id()

##########################################################################
# GET AUTHENTICATION TOKENS AND ESTABLISH CONNECTION TO THE ROUTER
#
# Establish the connection details to reach the TripleBlind instance.
# Unless explicitly specified, all operations will occur via this default
# session as the user 'organization_one'.
tb.initialize(api_token=tb.config.example_user1["token"])

# Register the first training database with the Router and place on Access Point
try:
    # Start a connection as 'organization_two' user
    session2 = tb.Session(api_token=tb.config.example_user2["token"], from_default=True)

    print(
        f"Creating dataset 'train-ssd-{run_id}' on {tb.config.example_user2['login']}'s Access Point..."
    )
    dataset_train0 = tb.Asset.position(
        file_handle=tb.Package.load("train.zip"),
        name=f"train-ssd-{run_id}",
        desc="Single shot detector training dataset",
        is_discoverable=True,
        session=session2,
    )

    if tb.config.create_agreements_on_example_input_assets:
        # For this example we will attach an Agreement to the datasets.  This
        # agreement makes the dataset available to the other team,
        # meaning the training step will not require the dataset owner to
        # explicitly grant permission for use.
        dataset_train0.add_agreement(
            with_team=tb.config.example_user1["team_id"],
            operation=tb.Operation.REGION_OF_INTEREST_TRAIN,
            session=session2,
        )
        print("Created Agreement")
except tb.TripleblindAssetAlreadyExists:
    print(
        f"Asset 'train-ssd-{run_id}' already exists, change value in run_id.txt to make it unique.'"
    )


try:
    # Start a connection as 'organization_two' user
    session3 = tb.Session(api_token=tb.config.example_user3["token"], from_default=True)

    print(
        f"Creating dataset 'train1-ssd-{run_id}' on {tb.config.example_user3['login']}'s Access Point..."
    )
    dataset_train1 = tb.Asset.position(
        file_handle=tb.Package.load("train2.zip"),
        name=f"train1-ssd-{run_id}",
        desc="Single shot detector training dataset",
        is_discoverable=True,
        session=session3,
    )

    if tb.config.create_agreements_on_example_input_assets:
        dataset_train1.add_agreement(
            with_team=tb.config.example_user1["team_id"],
            operation=tb.Operation.REGION_OF_INTEREST_TRAIN,
            session=session3,
        )
        print("Created Agreement")
except tb.TripleblindAssetAlreadyExists:
    print(
        f"Asset 'train1-ssd-{run_id}' already exists, change value in run_id.txt to make it unique.'"
    )

print("Data is in position.")
