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

from pathlib import Path

import yaml

import tripleblind as tb


DATASET_NAME1 = "DEMO - Snowflake Pharmacy Prescriptions"
DATASET_NAME2 = "DEMO - Hope Valley Hospital EHR"

OWNER1 = tb.config.example_user2
OWNER2 = tb.config.example_user3

############################################################################
# Connect Snowflake table of pharmacy prescriptions to dataset on org 2

tablename = "pharmacy_prescriptions"

# Load the Snowflake configuration file with authentication information
tb.util.set_script_dir_current()
data_dir = Path("example_data")
with open("snowflake.yaml", "r") as f:
    config = yaml.safe_load(f.read())

# Create the asset under the second example team
tb.initialize(api_token=OWNER1["token"])

try:
    # Create an asset that points to a Snowflake table
    asset = tb.asset.SnowflakeDatabase.create(
        config["USERNAME"],
        config["PASSWORD"],
        config["ACCOUNT"],
        config["DATABASE"],
        query=f"SELECT * FROM {tablename};",
        name=DATASET_NAME1,
        desc="Fictional report of pharmacy transactions, housed in a Snowflake data warehouse.",
        is_discoverable=True,
    )
    print(f"'{DATASET_NAME1}' asset created, owned by {OWNER1['name']}")
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{DATASET_NAME1}' already exists")
    asset = tb.TableAsset.find(DATASET_NAME1, owned_by=OWNER2["team_id"])


############################################################################
# Position CSV table of hospital EHR data to dataset on org 3

try:
    print(f"Uploading '{DATASET_NAME2}' to {OWNER2['login']}'s Access Point...")
    session2 = tb.Session(api_token=OWNER2["token"], from_default=True)
    asset = tb.Asset.position(
        file_handle=data_dir / "Hope_Valley_Hospital.csv",
        name=DATASET_NAME2,
        desc="Fictional hospital electronic health record (EHR) patient data.",
        is_discoverable=True,
        session=session2,
    )
    print(f"Asset created")
except tb.TripleblindAssetAlreadyExists:
    print(f"Asset '{DATASET_NAME2}' already exists, skipping.'")

# Give org3 permission to use in a Blind Join
asset.add_agreement(
    with_team=OWNER1["team_id"], operation=tb.Operation.BLIND_JOIN, session=session2
)
print("Created Agreement")


print("Data is in position.")
