#!/usr/bin/env python3 -W ignore

"""Generate view to unnest parquet-format list and map fields."""

import json
import sys
from argparse import ArgumentParser
from textwrap import dedent

from google.cloud import bigquery


def qualify(table, dataset, project):
    """Fully qualify table in standard sql format without `s."""
    table = table.replace("`", "").replace(":", ".")
    dots = table.count(".")
    if dots == 0:
        return f"{project}.{dataset}.{table}"
    elif dots == 1:
        return f"{project}.{table}"
    else:
        return table


def replace(field, unnest_layer=0, *prefix):
    """Convert BigQuery field to a SQL expression for use in a REPLACE block."""
    if field.field_type == "RECORD":
        is_list = (
            len(field.fields) == 1
            and field.fields[0].name == "list"
            and field.fields[0].mode == "REPEATED"
            and len(field.fields[0].fields) == 1
            and field.fields[0].fields[0].name == "element"
        )
        is_map = (
            len(field.fields) == 1
            and field.fields[0].name == "key_value"
            and field.fields[0].mode == "REPEATED"
        )
        if is_list or is_map:
            # prevent naming collisions via `UNNEST(...) AS`
            unnest_as = f"_{unnest_layer}"
            if is_list:
                unnest = ".".join(prefix + (field.name, "list"))
                nested = replace(field.fields[0].fields[0], unnest_layer + 1, unnest_as)
                # handle simplest case without unnest_as
                if nested == f"{unnest_as}.element":
                    return f"ARRAY(SELECT * FROM UNNEST({unnest})) AS {field.name}"
            elif is_map:
                unnest = ".".join(prefix + (field.name, "key_value"))
                nested = replace(
                    bigquery.SchemaField(
                        field_type="RECORD",
                        fields=field.fields[0].fields,
                        name=unnest_as,
                    ),
                    unnest_layer + 1,
                )
                # handle simplest case without array unnest_layer
                if nested == unnest_as:
                    return f"{unnest} AS {field.name}"
            # shorten special case for nested struct
            special_prefix = f"(SELECT AS STRUCT {unnest_as}.* "
            if nested.startswith(special_prefix):
                index = len(special_prefix)
                suffix = nested[index:].rsplit(")", 1)[0]
                select = f"SELECT AS STRUCT * {suffix}"
            # shorten nested struct
            elif nested.startswith("(SELECT AS STRUCT"):
                select = nested[1:].rsplit(")", 1)[0]
            # handle nested array
            elif nested.startswith("ARRAY"):
                select = f"SELECT AS STRUCT {nested}"
            else:
                select = f"SELECT {nested}"
            from_ = f"FROM UNNEST({unnest}) AS {unnest_as}"
            return f"ARRAY({select} {from_}) AS {field.name}"
        # handle other structs with nested replacements
        replacements = []
        for subfield in field.fields:
            replacement = replace(subfield, unnest_layer, *prefix, field.name)
            if " AS " in replacement:
                replacements += [replacement]
        if replacements:
            select = f"SELECT AS STRUCT {'.'.join(prefix + (field.name, '*'))}"
            replacements = ", ".join(replacements)
            return f"({select} REPLACE ({replacements})) AS {field.name}"
    # no unnesting needed
    return ".".join(prefix + (field.name,))


def generate(table, view):
    """Generate a CREATE VIEW statement that unnests lists and maps."""
    schema = bigquery.Client().get_table(table).schema
    replacements = []
    for index, field in enumerate(schema):
        try:
            replacement = replace(field)
        except Exception:
            json_field = json.dumps(field.to_api_repr(), indent=2)
            print(f"problem with field {index}:\n{json_field}", file=sys.stderr)
            raise
        if " AS " in replacement:
            replacements += [replacement]

    replacements.sort(key=lambda x: x.rsplit(" AS ", 1)[-1])

    return dedent(
        f"""
        -- Query generated by: {' '.join(sys.argv)}
        CREATE OR REPLACE VIEW
          `{view}` AS
        SELECT
          * REPLACE (
            """
        + f""",
            """.join(
            replacements
        )
        + f"""
          )
        FROM
          `{table}`
        """
    ).strip()


def main():
    """Parse argv and print CREATE VIEW statement."""
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("view")
    parser.add_argument("table")
    parser.add_argument(
        "--project",
        default="moz-fx-data-shared-prod",
        help="project to use for partially qualified table and view, defaults to"
        " moz-fx-data-shared-prod",
    )
    parser.add_argument(
        "--dataset",
        default="telemetry",
        help="dataset to use for partially qualified table and view, defaults"
        " telemetry",
    )
    args = parser.parse_args()
    view = qualify(args.view, args.dataset, args.project)
    table = qualify(args.table, args.dataset, args.project)
    sql = generate(table, view)
    print(sql)


if __name__ == "__main__":
    main()
