"""
Generate user-facing views on top of stable tables.

If there are multiple versions of a given document type, the generated view
references only the most recent.

Note that no view will be generated if a corresponding view is already
present in the target directory, which allows manual overrides of views by
checking them into the sql/ tree of the default branch of the repository.
"""

import json
import logging
import tarfile
import urllib.request
from argparse import ArgumentParser
from dataclasses import dataclass
from functools import partial
from io import BytesIO
from itertools import groupby
from multiprocessing.pool import ThreadPool
from pathlib import Path
import re

from typing import List

from bigquery_etl.dryrun import DryRun
from bigquery_etl.format_sql.formatter import reformat
from bigquery_etl.metadata.parse_metadata import DatasetMetadata
from bigquery_etl.schema import Schema
from bigquery_etl.util import standard_args

MPS_URI = "https://github.com/mozilla-services/mozilla-pipeline-schemas"

SKIP_PREFIXES = (
    "pioneer",
    "rally",
)

VIEW_QUERY_TEMPLATE = """\
-- Generated by bigquery_etl.view.generate_stable_views
CREATE OR REPLACE VIEW
  `{full_view_id}`
AS
SELECT
  * REPLACE(
    {replacements})
FROM
  `{target}`
"""

VIEW_METADATA_TEMPLATE = """\
# Generated by bigquery_etl.view.generate_stable_views
---
friendly_name: Historical Pings for `{document_namespace}/{document_type}`
description: |-
  A historical view of pings sent for the
  `{document_namespace}/{document_type}`
  document type.

  This view is guaranteed to contain only complete days
  (per `submission_timestamp`)
  and to contain only one row per distinct `document_id` within a given date.

  Clustering fields: `normalized_channel`, `sample_id`
"""

VIEW_CREATE_REGEX = re.compile(
    r"CREATE OR REPLACE VIEW\n\s*[^\s]+\s*\nAS", re.IGNORECASE
)

parser = ArgumentParser(description=__doc__)
parser.add_argument(
    "--sql-dir", default="sql/", help="The path where generated SQL files are stored."
)
parser.add_argument(
    "--target-project",
    default="moz-fx-data-shared-prod",
    help="The project where the stable tables live.",
)
standard_args.add_log_level(parser)
standard_args.add_parallelism(parser)


@dataclass
class SchemaFile:
    """Container for metadata about a JSON schema and corresponding BQ table."""

    schema: dict
    schema_id: str
    bq_dataset_family: str
    bq_table: str
    document_namespace: str
    document_type: str
    document_version: int

    @property
    def bq_table_unversioned(self):
        """Return table_id with version suffix stripped."""
        return "_".join(self.bq_table.split("_")[:-1])

    @property
    def stable_table(self):
        """Return BQ stable table name in <dataset>.<table> form."""
        return f"{self.bq_dataset_family}_stable.{self.bq_table}"

    @property
    def user_facing_view(self):
        """Return user-facing view name in <dataset>.<view> form."""
        return f"{self.bq_dataset_family}.{self.bq_table_unversioned}"

    @property
    def sortkey(self):
        """Return variant of stable_table with zero-padded version for sorting."""
        return (
            "_".join(self.stable_table.split("_")[:-1]) + f"{self.document_version:04d}"
        )


def write_dataset_metadata_if_not_exists(
    target_project: str, sql_dir: Path, schema: SchemaFile
):
    """Write default dataset_metadata.yaml files where none exist.

    This function expects to be handed one representative `SchemaFile`
    object representing the dataset.
    """
    dataset_family = schema.bq_dataset_family
    project_dir = sql_dir / target_project

    # Derived dataset
    dataset_name = f"{dataset_family}_derived"
    target = project_dir / dataset_name / "dataset_metadata.yaml"
    target.parent.mkdir(parents=True, exist_ok=True)
    if not target.exists():
        print(f"Creating {target}")
        DatasetMetadata(
            friendly_name=f"{schema.document_namespace} Derived",
            description=(
                f"Derived tables related to document namespace"
                f" {schema.document_namespace},"
                f" usually populated via queries defined in"
                f" https://github.com/mozilla/bigquery-etl"
                f" and managed by Airflow"
            ),
            dataset_base_acl="derived",
            user_facing=False,
        ).write(target)

    # User-facing dataset
    dataset_name = dataset_family
    target = project_dir / dataset_name / "dataset_metadata.yaml"
    target.parent.mkdir(parents=True, exist_ok=True)
    if not target.exists():
        print(f"Creating {target}")
        DatasetMetadata(
            friendly_name=f"{schema.document_namespace}",
            description=(
                f"User-facing views related to document namespace"
                f" {schema.document_namespace}; see https://github.com/"
                f"mozilla-services/mozilla-pipeline-schemas/tree/"
                f"generated-schemas/schemas/{schema.document_namespace}"
            ),
            dataset_base_acl="view",
            user_facing=True,
        ).write(target)


def write_view_if_not_exists(target_project: str, sql_dir: Path, schema: SchemaFile):
    """If a view.sql does not already exist, write one to the target directory."""
    target_dir = (
        sql_dir
        / target_project
        / schema.bq_dataset_family
        / schema.bq_table_unversioned
    )

    target_file = target_dir / "view.sql"

    if target_file.exists():
        return

    full_source_id = f"{target_project}.{schema.stable_table}"
    full_view_id = f"{target_project}.{schema.user_facing_view}"
    replacements = ["mozfun.norm.metadata(metadata) AS metadata"]
    if schema.schema_id == "moz://mozilla.org/schemas/glean/ping/1":
        replacements += ["mozfun.norm.glean_ping_info(ping_info) AS ping_info"]
        if schema.bq_table == "baseline_v1":
            replacements += [
                "mozfun.norm.glean_baseline_client_info"
                "(client_info, metrics) AS client_info"
            ]
        if (
            schema.bq_dataset_family == "org_mozilla_fenix"
            and schema.bq_table == "metrics_v1"
        ):
            # todo: use mozfun udfs
            replacements += [
                "mozdata.udf.normalize_fenix_metrics"
                "(client_info.telemetry_sdk_build, metrics)"
                " AS metrics"
            ]
        if schema.bq_dataset_family == "firefox_desktop":
            # FOG does not provide an app_name, so we inject the one that
            # people already associate with desktop Firefox per bug 1672191.
            replacements += [
                "'Firefox' AS normalized_app_name",
            ]
    elif schema.schema_id.startswith("moz://mozilla.org/schemas/main/ping/"):
        replacements += ["mozdata.udf.normalize_main_payload(payload) AS payload"]
    replacements_str = ",\n    ".join(replacements)
    full_sql = reformat(
        VIEW_QUERY_TEMPLATE.format(
            target=full_source_id,
            replacements=replacements_str,
            full_view_id=full_view_id,
        )
    )
    print(f"Creating {target_file}")
    target_dir.mkdir(parents=True, exist_ok=True)
    with target_file.open("w") as f:
        f.write(full_sql)
    metadata_content = VIEW_METADATA_TEMPLATE.format(
        document_namespace=schema.document_namespace,
        document_type=schema.document_type,
    )
    metadata_file = target_dir / "metadata.yaml"
    if not metadata_file.exists():
        with metadata_file.open("w") as f:
            f.write(metadata_content)

    # get view schema with descriptions
    try:
        content = VIEW_CREATE_REGEX.sub("", target_file.read_text())
        content += " WHERE DATE(submission_timestamp) = '2020-01-01'"
        view_schema = Schema.from_query_file(target_file, content=content)

        stable_table_schema = Schema.from_json({"fields": schema.schema})
        view_schema.merge(stable_table_schema, add_missing_fields=False)
        view_schema.to_yaml_file(target_dir / "schema.yaml")
    except Exception as e:
        print(f"Cannot generate schema.yaml for {target_file}: {e}")


def get_stable_table_schemas() -> List[SchemaFile]:
    """Fetch last schema metadata per doctype by version."""
    schemas_uri = prod_schemas_uri()
    with urllib.request.urlopen(schemas_uri) as f:
        tarbytes = BytesIO(f.read())

    schemas = []
    with tarfile.open(fileobj=tarbytes, mode="r:gz") as tar:
        for tarinfo in tar:
            if tarinfo.name.endswith(".schema.json"):
                *_, document_namespace, document_type, basename = tarinfo.name.split(
                    "/"
                )
                version = int(basename.split(".")[1])
                schema = json.load(tar.extractfile(tarinfo.name))  # type: ignore
                bq_schema = {}

                try:
                    bq_schema_file = tar.extractfile(
                        tarinfo.name.replace(".schema.json", ".bq")
                    )
                    bq_schema = json.load(bq_schema_file)  # type: ignore
                except KeyError as e:
                    print(f"Cannot get Bigquery schema for {tarinfo.name}: {e}")

                pipeline_meta = schema.get("mozPipelineMetadata", None)
                if pipeline_meta is None:
                    continue
                schemas.append(
                    SchemaFile(
                        schema=bq_schema,
                        schema_id=schema.get("$id", ""),
                        bq_dataset_family=pipeline_meta["bq_dataset_family"],
                        bq_table=pipeline_meta["bq_table"],
                        document_namespace=document_namespace,
                        document_type=document_type,
                        document_version=version,
                    )
                )

    # Exclude doctypes maintained in separate projects.
    for prefix in SKIP_PREFIXES:
        schemas = [
            schema
            for schema in schemas
            if not schema.document_namespace.startswith(prefix)
        ]

    # Retain only the highest version per doctype.
    schemas = sorted(
        schemas,
        key=lambda t: f"{t.document_namespace}/{t.document_type}/{t.document_version:03d}",
    )
    schemas = [
        last
        for k, (*_, last) in groupby(
            schemas, lambda t: f"{t.document_namespace}/{t.document_type}"
        )
    ]

    return schemas


def prod_schemas_uri():
    """Return URI for the schemas tarball deployed to shared-prod.

    We construct a fake query and send it to the dry run service in order
    to read dataset labels, which contains the commit hash associated
    with the most recent production schemas deploy.
    """
    dryrun = DryRun("telemetry_derived/foo/query.sql", content="SELECT 1")
    build_id = dryrun.get_dataset_labels()["schemas_build_id"]
    commit_hash = build_id.split("_")[-1]
    return f"{MPS_URI}/archive/{commit_hash}.tar.gz"


def main():
    """
    Generate view definitions.

    Metadata about document types is from the generated-schemas branch
    of mozilla-pipeline-schemas. We write out generated views for each
    document type in parallel.

    There is a performance bottleneck here due to the need to dry-run each
    view to ensure the source table actually exists.
    """
    args = parser.parse_args()

    # set log level
    try:
        logging.basicConfig(level=args.log_level, format="%(levelname)s %(message)s")
    except ValueError as e:
        parser.error(f"argument --log-level: {e}")

    schemas = get_stable_table_schemas()
    one_schema_per_dataset = [
        last for k, (*_, last) in groupby(schemas, lambda t: t.bq_dataset_family)
    ]

    with ThreadPool(args.parallelism) as pool:
        pool.map(
            partial(
                write_view_if_not_exists,
                args.target_project,
                Path(args.sql_dir),
            ),
            schemas,
            chunksize=1,
        )
        pool.map(
            partial(
                write_dataset_metadata_if_not_exists,
                args.target_project,
                Path(args.sql_dir),
            ),
            one_schema_per_dataset,
            chunksize=1,
        )


if __name__ == "__main__":
    main()
