import pandas as pd

from gcp_airflow_foundations.base_class.dlp_source_config import PolicyTagConfig


def get_dlp_results_sql(project_id, dataset_id, table_id, min_count=0) -> str:
    """ Generate sql to query the DLP results table: https://cloud.google.com/dlp/docs/querying-findings
        and counts the number of finding in each field,info-type, likelihood grouping
    :param project_id: Project Id
    :param dataset_id: Dataset
    :param table_id: Table
    :return:a sql query that generates a result set with the following columns [field_name, info_type_name, likelihood, count_total]
    """

    return """SELECT
        locations.record_location.field_id.name AS field_name,
        info_type.name as info_type_name,
        likelihood as likelihood,
        COUNT(*) AS count_total
    FROM {}.{}.{},
        UNNEST(location.content_locations) AS locations
    GROUP BY
        locations.record_location.field_id.name,
        info_type.name,
        likelihood
    HAVING count_total > {}
    """.format(
        project_id, dataset_id, table_id, min_count
    )


def results_to_bq_policy_tags(
    project_id, fields, poicy_tag_config: PolicyTagConfig
) -> list:
    """ Generates a tag to be applied for each column (if any)
    :param fields: A fields as generated by a query from get_dlp_results_sql
    :param tag: The tag to apply
    :return: A list of columns with am updated policyTags as per  https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#TableSchema
 """
    df = pd.DataFrame(
        fields, columns=["field_name", "info_type_name", "likelihood", "count_total"]
    )
    tags = []
    unique = df.drop_duplicates(subset=["field_name"])
    for field in unique["field_name"]:
        tags.append(
            {
                "name": field,
                "policy_tags": {
                    "names": [policy_tag_name(fields, project_id, poicy_tag_config)]
                },
            }
        )
    return tags


def policy_tag_name(field, project_id, poicy_tag_config: PolicyTagConfig):
    return f"projects/{project_id}/locations/{poicy_tag_config.location}/taxonomies/{poicy_tag_config.taxonomy}/policyTags/{poicy_tag_config.tag}"
