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

import tripleblind as tb


# The creation of the summary report is performed by the data owner.  In this
# example the first organization is the report owner with a database isolated
# behind their Access Point.
tb.initialize(api_token=tb.config.example_user1["token"])
tb.util.set_script_dir_current()

# This query template defines the information they are willing to share with
# others.  It contains two parameters -- "demographic" and "pay" -- both of
# which are used twice in the query & an unused "trt_code" parameter:
TEMPLATE = """
SELECT Dept_Name, {{demographic}}, avg({{pay}}) AS average_{{pay}} FROM tripleblind_datasets.city_of_somerville_payroll
-- {{#trt_code}}{{#_FIRST}}WHERE{{/_FIRST}} TRTCode = {{_VALUE}} {{^_LAST}}OR{{/_LAST}} {{/trt_code}}
GROUP BY Dept_Name, {{demographic}};
"""

# The first parameter has three possible values
demographic_param = tb.report_asset.ReportParameter.create_string(
    name="demographic",
    description="Specific demographic to filter the report on",
    display="Select demographic filter",
    options=["Gender", "Title", "Ethnicity"],
)


# The second parameter has a default value & required=False,
# making it an optional parameter.
pay_param = tb.report_asset.ReportParameter.create_string(
    name="pay",
    description="specific pay field to display",
    display="pay",
    default_value="Total_Gross",
    required=False,
    options=[
        tb.report_asset.ParameterOption(
            "Total_Gross", "Total Gross", "Pay amount before taxes"
        ),
        tb.report_asset.ParameterOption(
            "Reg", "Regular", "Paycheck amount after taxes"
        ),
    ],
)

# The third parameter is a code parameter which limits the report to only
# those in requested CPT treatment code classes. When used in the web interface,
# the user would be assisted in typing codes via the "cpt" codes stored on
# the Router.
# Because the user can select any number of codes, the mustache template
# for this parameter is slightly more complicated. Consult the documentation
# before using it in a Blind Report.
# NOTE: This parameter isn't actually used in the SQL query, but it is included
# in the template as a SQL comment to show how it would be used.
trt_code_param = tb.report_asset.ReportParameter.create_code(
    name="trt_code",
    description="CPT/ICD Treatment Code",
    display="CPT/ICD Treatment Code",
    systems=["cpt", "icd9_pcs", "icd10_pcs"],
    required=False,
)


# Optional: Validate your query_template and parameters.
# This is just a helper function to help confirm that your setup
# of the template and parameters is valid. Since the mustache templating
# language can be unintuitive, this can help catch errors quickly.
tb.report_asset.validate_report_template(
    TEMPLATE, [demographic_param, pay_param, trt_code_param]
)

# This example is built around a BigQuery database, but you can use any other database.
# For a comprehensive list of options, check the Data_Connectors directory.
# You will need to point to the credentials that Google provides,
# your GCP project associated with the database, and the database name itself.

# Path to a GCP Credentials file.  This example includes a read-only credential
# for a dataset hosted by TripleBlind using BigQuery.
# For help on obtaining one a JSON file for your own BigQuery database see:
# https://developers.google.com/workspace/guides/create-credentials#create_credentials_for_a_service_account
CREDENTIAL_FILE = "credentials.json"

# A BigQuery dataset exists within a project, so both the project ID and the
# dataset name must be specified.
GCP_PROJECT_ID = "tripleblind-datasets"
DATASET_NAME = "tripleblind_datasets"

# An Asset can be published for any organization to see by making it
# "discoverable".  It is also possible to create an Asset that is only
# visible to specific organizations.
PUBLISH_ON_ROUTER = True

# The name of the report asset
NAME = "EXAMPLE - Payroll Report"

# This post-processing script will be run after the query is executed. It can be
# used to polish the output before it is returned to the user.  The parameters
# are a pandas dataframe and a context dictionary.  The context contains all of
# the parameters selected by the user.
#
#        "name": package.meta.record.name,
#        "description": package.meta.record.description,     # str
#        "initiator_details": job_params.initiator_details,  # Dict[str, str]
#        "attributes": {
#            "report_values": display_params,                # Dict[str, str]
#            "raw_values": raw_params,                       # Dict[str, str]
#            "federation_members": fed_members,              # List[str] (only for federated reports)
#        },

script = """
def postprocess(df: "pd.Dataframe", ctx: dict):

    # Change "ISD" in the Dept_Name column to "Information Services"
    df["Dept_Name"] = df["Dept_Name"].str.replace("ISD", "INFORMATION SERVICES")

    # Rename a column based on the user-selected parameter
    df.rename(columns={
        'demographic': ctx["attributes"]["report_values"]["Select demographic filter"]
    }, inplace=True)

    return df
"""

try:
    print("Creating BigQuery Report...")
    report_asset = tb.report_asset.BigQueryReport.create(
        GCP_PROJECT_ID,
        DATASET_NAME,
        CREDENTIAL_FILE,
        query_template=TEMPLATE,
        params=[demographic_param, pay_param, trt_code_param],
        name=NAME,
        desc="Sample report from a BigQuery database of salaries with configurable demographics.",
        is_discoverable=PUBLISH_ON_ROUTER,
        post_processing=script,
    )
except tb.TripleblindAssetAlreadyExists:
    print("Summary report already exists.")
    report_asset = tb.ReportAsset.find(NAME)

try:
    report_asset.add_agreement(with_team="ANY", operation=report_asset.uuid)
    print("Added Agreement allowing any team to run the report.")
except Exception as e:
    print(e)
    print("Error: Could not create the agreement.")

print("Report ID:", report_asset.uuid)
