rebuild

#!/usr/bin/env python3

"""
Builds a fresh data repo from source data
"""
import argparse
import json
import time
from collections import defaultdict
from copy import deepcopy
from os import getcwd
from os.path import dirname, realpath, join, relpath, isfile

from lib.changelog import changelog_prepare, changelog_get_unreleased_section
from lib.container import dict_get, dict_get_required, find_index_by, first, format_list, \
  dict_remove_many, find_duplicates, dict_cleanup, find, unique
from lib.date import now_iso, iso_to_iso_safe
from lib.fasta import fasta_read_exactly_one_seq
from lib.fs import json_read, find_files, json_write, copy, make_zip, file_write, rmrf
from lib.git import git_get_modified_files, git_dir_is_clean, git_get_dirty_files, git_check_tag, \
  github_create_release, git_pull, git_commit_all, git_push
from lib.logger import l
from lib.minimizer import make_ref_search_index, serialize_ref_search_index


def get_dataset_capabilities(pathogen_json: dict, dataset_dir: str):
  ref_filename = dict_get(pathogen_json, ["files", "reference"])
  if not ref_filename:
    raise FileNotFoundError(f"Reference sequence file must be declared `.files.reference` field of pathogen.json")

  files = dict_get_required(pathogen_json, ["files"])
  for (name, filename) in files.items():
    filepath = join(dataset_dir, filename)
    if not isfile(filepath):
      raise FileNotFoundError(
        f"'Filename '{filename}' is declared in `.files.{name}` field of pathogen.json, but the actual file is not "
        f"found: '{filepath}'")

  other = []

  tree_filename = dict_get(pathogen_json, ["files", "treeJson"])
  tree_json_path = join(dataset_dir, tree_filename) if tree_filename else None
  clades = []
  custom_clades = {}
  if tree_json_path is not None and isfile(tree_json_path):
    tree_json = json_read(tree_json_path)
    clades = tree_find_clades(tree_json)
    custom_clades = tree_find_clade_like_attrs(tree_json)

  if dict_get(pathogen_json, ["mutLabels"]) is not None:
    other.append("mutLabels")

  if dict_get(pathogen_json, ["phenotypeData"]) is not None:
    other.append("phenotypeData")

  if dict_get(pathogen_json, ["aaMotifs"]) is not None:
    other.append("aaMotifs")

  qc = []
  for k, q in (dict_get(pathogen_json, ["qc"]) or {}).items():
    if dict_get(q, ["enabled"]):
      qc.append(k)

  custom_clades = dict_cleanup({attr: len(values) for attr, values in custom_clades.items() if len(values) > 0})

  return dict_cleanup({
    "clades": len(clades) if len(clades) > 0 else None,
    "customClades": custom_clades,
    "qc": qc,
    "primers": True if len(dict_get(pathogen_json, ["primers"]) or []) > 0 else None,
    "other": other
  })


def tree_find_clades(auspice_json):
  def tree_find_clades_recursive(node, clades=None):
    if clades is None:
      clades = []

    clade_membership = node.get('node_attrs', {}).get('clade_membership', {}).get('value')
    if clade_membership:
      clades.append(clade_membership)

    children = node.get('children', [])
    for child in children:
      tree_find_clades_recursive(child, clades)

    return clades

  clades = tree_find_clades_recursive(auspice_json["tree"])
  return list(sorted(unique(clades)))


def tree_find_clade_like_attrs(auspice_json):
  def tree_find_clade_like_attrs_recursive(node, attr_names, attributes=None):
    if attributes is None:
      attributes = {attr: [] for attr in attr_names}

    for attr in attr_names:
      attr_value = node.get('node_attrs', {}).get(attr, {}).get('value')
      if attr_value is not None:
        attributes[attr].append(attr_value)

    children = node.get('children', [])
    for child in children:
      tree_find_clade_like_attrs_recursive(child, attr_names, attributes)

    return attributes

  clade_node_attrs = dict_get(auspice_json, ["meta", "extensions", "nextclade", "clade_node_attrs"]) or []
  attr_names = [attr["name"] for attr in clade_node_attrs]
  attributes = tree_find_clade_like_attrs_recursive(auspice_json["tree"], attr_names)
  return {attr: list(sorted(unique(values))) for attr, values in attributes.items() if len(values) > 0}


def dataset_get_versions(dataset):
  versions = dict_get(dataset, ["versions"]) or []
  versions = list(filter(lambda version: version["tag"] != "unreleased", versions))
  last_version = first(sorted(versions, reverse=True, key=lambda v: v["tag"]))
  return versions, last_version


def index_one_dataset(args, pathogen_json_path: str, dataset: object, tag: str, updated_at: str):
  pathogen_json = json_read(pathogen_json_path)
  dataset_dir = dirname(pathogen_json_path)
  path = relpath(dataset_dir, args.input_dir)

  ref = get_ref_seq(pathogen_json, dataset_dir)
  check_ref_seq_mismatch(path, ref, pathogen_json, dataset_dir)

  versions, last_version = dataset_get_versions(dataset)

  if dataset_has_changes(dataset, dataset_dir):
    new_version = dict_cleanup({
      "updatedAt": None if tag == "unreleased" else updated_at,
      "tag": tag,
      "compatibility": dict_get(pathogen_json, ["compatibility"]),
    })
    versions.insert(0, new_version)
    last_version = new_version

  dataset_out = dict_cleanup({
    "path": path,
    "shortcuts": dict_get(pathogen_json, ["shortcuts"]),
    "enabled": False if dict_get(pathogen_json, ["enabled"]) == False else True,
    "attributes": dict_get(pathogen_json, ["attributes"]),
    "maintenance": dict_get(pathogen_json, ["maintenance"]),
    "files": dict_get_required(pathogen_json, ["files"]),
    "capabilities": get_dataset_capabilities(pathogen_json, dataset_dir),
    "versions": versions,
    "version": last_version,
  })

  return dataset_out, ref


def get_ref_seq(pathogen_json, dataset_dir):
  files = dict_get_required(pathogen_json, ["files"])
  ref_filename = dict_get_required(files, ["reference"])
  ref_filepath = join(dataset_dir, ref_filename)
  try:
    return fasta_read_exactly_one_seq(ref_filepath)
  except Exception as e:
    raise ValueError(f"When reading reference sequence") from e


def check_ref_seq_mismatch(path, standalone_ref, pathogen_json, dataset_dir):
  tree_filename = dict_get(pathogen_json, ["files", "treeJson"])
  tree_json_path = join(dataset_dir, tree_filename) if tree_filename else None
  if tree_json_path is not None and isfile(tree_json_path):
    tree_json = json_read(tree_json_path)
    tree_ref = dict_get(tree_json, ["root_sequence", "nuc"])
    if tree_ref is not None:
      if standalone_ref.seq != tree_ref:
        l.warning(
          f"{path}: Reference sequence provided does not exactly match reference (root) sequence in Auspice JSON. "
          f"This warning signals that there is a potential for failures if the mismatch is not intended."
        )


def get_new_dataset_order(datasets, dataset_order):
  paths = list(map(lambda d: d["path"], datasets))

  dupes = find_duplicates(dataset_order)
  if len(dupes):
    raise ValueError(
      f"The '.dataset_order' list in 'collection.json' contains duplicated entries: {format_list(dupes)}. "
      f"Please make sure the entries are unique, to avoid ambiguity."
    )

  # extra = set(dataset_order).difference(set(paths))
  # if len(extra) > 0:
  #   raise ValueError(
  #     f"The '.dataset_order' list in 'collection.json' contains entries for datasets that are not found: "
  #     f"{format_list(extra)}. "
  #     f"Please double check the existence of datasets and the spelling of their paths in the '.dataset_order' list. "
  #     f"The full list of datasets that are found:\n  {format_list(paths)}."
  #   )

  missing = set(paths).difference(set(dataset_order))
  if len(missing):
    l.info(
      f"Adding '.dataset_order' entries to 'collection.json' for the following datasets: {format_list(missing)}. "
      f"Please reorder them manually as needed. This order is used when displaying datasets of the collection in the "
      f"user interface."
    )
    dataset_order += list(missing)
  return dataset_order


def sort_collections(collections, collections_order):
  return [collection for x in collections_order for collection in collections if collection["meta"]["id"] == x]


def sort_datasets(datasets, dataset_order):
  return [dataset for x in dataset_order for dataset in datasets if dataset["path"] == x]


def sort_release_infos(release_infos, dataset_order):
  return [release_info for x in dataset_order for release_info in release_infos if release_info["dataset"]["path"] == x]


def parse_args():
  parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

  parser.add_argument('--input-dir', required=True, help='Directory of source files of a dataset collection')
  parser.add_argument('--output-dir', required=True, help='Where to output prepared files')
  parser.add_argument('--temp-dir', default=join(getcwd(), "data_temp"),
                      help='Where to temporarily store release files')

  parser.add_argument('--no-pull', action='store_true',
                      help="Skip pulling latest changes from remote. Note that if "
                           "there are incompatible changes this can create hard to resolve merge conflicts.")
  parser.add_argument('--allow-dirty', action='store_true',
                      help="Allow working directory to contain uncommitted files. Note that these files will also be "
                           "committed and can potentially be released along with the datasets",
                      )
  parser.add_argument('--commit', action='store_true', help="Commit updates into the repository")
  parser.add_argument('--push', action='store_true', help="Push updates into the remote repository. Implies --commit")
  parser.add_argument('--release', action='store_true', help="Release updates. Implies --commit and --push")
  parser.add_argument('--repo', required=False,
                      help="GitHub repo to push and to release to. You need to have write permission for that."
                      )

  args = parser.parse_args()

  if args.push and not args.repo:
    parser.error("--push requires --repo")
  if args.release and not args.repo:
    parser.error("--release requires --repo")

  if args.release:
    args.push = True

  if args.push:
    args.commit = True

  return args


def validate_shortcuts(collections):
  all_shortcuts = []
  for collection in collections:
    for dataset in collection['datasets']:
      shortcuts = dict_get(dataset, ['shortcuts'])
      if shortcuts:
        all_shortcuts.extend(shortcuts)
  duplicated_shortcuts = find_duplicates(all_shortcuts)

  duplicated_shortcuts_map = defaultdict(list)
  for collection in collections:
    for dataset in collection['datasets']:
      shortcuts = dict_get(dataset, ['shortcuts'])
      if shortcuts:
        for shortcut in shortcuts:
          if shortcut in duplicated_shortcuts:
            duplicated_shortcuts_map[shortcut].append(dataset["path"])

  if len(duplicated_shortcuts_map) > 0:
    dupes = [""]
    for shortcut, datasets in duplicated_shortcuts_map.items():
      datasets = ", ".join(datasets)
      dupes.append(f"{shortcut}: {datasets}")
    dupes = "\n - ".join(dupes)
    raise ValueError(
      f"Dataset name shortcuts must be unique, but found duplicates:\n{dupes}\n\nPlease ensure shortcuts are "
      f"unique by modifying or removing 'shortcuts' array from 'pathogen.json' files of these datasets.")


def main():
  args = parse_args()

  if not args.allow_dirty and not git_dir_is_clean():
    dirty_files = "\n  ".join(git_get_dirty_files())
    raise ValueError(
      f"Uncommitted changes detected. Refusing to proceed. Commit or stash changes first, or use --allow-dirty to "
      f"override (not recommended):\n  {dirty_files}"
    )

  if not args.no_pull:
    git_pull()

  updated_at = now_iso()
  tag = iso_to_iso_safe(updated_at) if args.release else "unreleased"
  git_check_tag(tag)

  collection_json_paths = list(find_files("collection.json", args.input_dir))

  index_json_path = join(args.output_dir, "index.json")
  try:
    index_json = json_read(index_json_path)
  except FileNotFoundError as _:
    index_json = {}

  collections = []
  release_infos = []
  all_refs = {}
  for collection_json_path in collection_json_paths:
    collection, release_infos_for_dataset, refs = process_one_collection(
      collection_json_path, index_json, args, tag, updated_at
    )
    collections.append(collection)
    release_infos.extend(release_infos_for_dataset)
    all_refs.update(refs)

  collections = sort_collections(collections, ["nextstrain", "community"])
  validate_shortcuts(collections)

  minimizer_index = make_ref_search_index(all_refs)

  minimizer_index_json_path = "minimizer_index.json"
  json_write(
    serialize_ref_search_index(minimizer_index),
    join(args.output_dir, minimizer_index_json_path),
    no_sort_keys=True
  )

  # if len(release_infos) == 0:
  #   l.info("No dataset modifications detected. Will not release anything.")
  #   return

  index_json = {
    "schemaVersion": "3.0.0",
    "collections": collections,
    "minimizerIndex": [
      {
        "version": minimizer_index["version"],
        "path": minimizer_index_json_path,
      }
    ]
  }

  json_write(index_json, index_json_path, no_sort_keys=True)

  if args.commit:
    commit_hash = commit_changes(args, tag, release_infos)

    if args.push:
      l.info("Pushing committed changes to GitHub")
      git_push()

      if args.release:
        l.info("Releasing to GitHub")
        release_notes = aggregate_release_notes(release_infos)
        l.info(f"Release notes:\n-------\n{release_notes}\n-------\nEnd of release notes\n")
        time.sleep(5)
        publish_to_github_releases(args, tag, commit_hash, release_notes)


def dataset_has_changes(dataset_from_index, dataset_dir):
  _, last_version = dataset_get_versions(dataset_from_index)
  modified_files = list(git_get_modified_files(from_revision=dict_get(last_version, ["tag"]), dirs=dataset_dir))

  modified_files = list(map(lambda f: realpath(f), modified_files))
  return len(modified_files) > 0


def process_one_collection(collection_json_path, index_json, args, tag, updated_at):
  collection_json = json_read(collection_json_path)
  collection_dir = dirname(collection_json_path)
  collection_id = dict_get_required(collection_json, ["meta", "id"])
  datasets_from_index_json = get_datasets_from_index_json(index_json, collection_id)

  datasets = []
  refs = {}
  for pathogen_json_path in find_files("pathogen.json", collection_dir):
    dataset_dir = dirname(pathogen_json_path)
    path = relpath(dataset_dir, args.input_dir)
    dataset_from_index = find(lambda dataset: dataset["path"] == path, datasets_from_index_json) or {}

    try:
      dataset, ref = index_one_dataset(args, pathogen_json_path, dataset_from_index, tag, updated_at)
      datasets.append(dataset)
      refs.update({dataset["path"]: ref})
    except Exception as e:
      raise ValueError(f"When processing '{pathogen_json_path}'") from e

  dataset_order = get_new_dataset_order(datasets, dict_get_required(collection_json, ["dataset_order"]) or [])
  datasets = sort_datasets(datasets, dataset_order)

  collection_json = {
    **collection_json,
    "dataset_order": dataset_order,
  }
  json_write(collection_json, collection_json_path, no_sort_keys=True)

  release_infos = prepare_dataset_release_infos(args, datasets, datasets_from_index_json, collection_dir, tag,
                                                updated_at)
  release_infos = sort_release_infos(release_infos, dataset_order)

  collection_info = deepcopy(collection_json)
  dict_remove_many(collection_info, ["dataset_order"])
  dict_remove_many(collection_info, ["schemaVersion"])

  collection_info = {
    **collection_info,
    "datasets": datasets,
  }

  return collection_info, release_infos, refs


def get_datasets_from_index_json(index_json, collection_id):
  collections = dict_get(index_json, ["collections"])
  collection = find(lambda c: dict_get_required(c, ["meta", "id"]) == collection_id, collections)
  return dict_get(collection, ["datasets"]) or []


def prepare_dataset_release_infos(args, datasets, datasets_from_index_json, collection_dir, tag, updated_at):
  release_infos = []
  for pathogen_json_path in find_files("pathogen.json", collection_dir):
    pathogen_json = json_read(pathogen_json_path)
    dataset_dir = dirname(pathogen_json_path)
    dataset_dir_rel = relpath(dataset_dir, args.input_dir)

    dataset_from_index = pathogen_json
    i_dataset = find_index_by(lambda dataset: dataset["path"] == dataset_dir_rel, datasets_from_index_json)
    if i_dataset is not None:
      dataset_from_index = datasets_from_index_json[i_dataset]

    if not dataset_has_changes(dataset_from_index, dataset_dir):
      continue

    dataset_new = dataset_from_index
    i_dataset = find_index_by(lambda dataset: dataset["path"] == dataset_dir_rel, datasets)
    if i_dataset is not None:
      dataset_new = datasets[i_dataset]

    path = dict_get(dataset_new, ['path']) or dataset_dir_rel

    changelog_path = join(dataset_dir, "CHANGELOG.md")
    release_notes = changelog_get_unreleased_section(changelog_path)
    if len(release_notes) == 0:
      raise ValueError(
        f"Changelog not found: dataset '{path}' has changes, but no valid changelog entry is found. Automated "
        f"releases will fail.\n\nTo fix this, please modify the file '{changelog_path}': add '## Unreleased' second "
        f"level heading (spelled exactly like this, without quotes) on the very top of the file and briefly "
        f"summarize the changes you've added. If the heading already exists, append your changes to the existing "
        f"changes.\n\nRefer to documentation for the instructions on how to prepare a dataset correctly."
      )

    if args.release:
      _, last_version = dataset_get_versions(dataset_from_index)
      release_info = prepare_dataset_release_info(
        dataset_dir, dataset_from_index, dict_get(last_version, ["tag"]), updated_at
      )
      if release_info is None:
        continue

      release_infos.append(release_info)

    create_dataset_package(args, dataset_new, path, tag, dataset_dir)

  return release_infos


def aggregate_release_notes(release_infos):
  dataset_list = format_dataset_list(release_infos)
  release_notes = f"This release contains changes for datasets:\n\n{dataset_list}\n\n\n"
  for release_info in release_infos:
    release_notes += f'\n{release_info["release_notes"]}\n\n'
  return release_notes


def format_dataset_list(release_infos):
  entries = [format_dataset_list_entry(release_info) for release_info in release_infos]
  return format_list(entries, sep="\n", marker="- ", quote=False)


def format_dataset_list_entry(release_info):
  path = dict_get_required(release_info, ['dataset', 'path'])
  name = get_dataset_name_friendly(release_info['dataset'])
  return f"{name} ({path})"


def commit_changes(args, tag, release_infos):
  l.info(f"Committing changes for '{tag}'")

  commit_message = "chore: rebuild [skip ci]"
  if args.release:
    dataset_list = format_dataset_list(release_infos)
    commit_message = f"chore: release '{tag}'\n\nUpdated datasets:\n\n{dataset_list}"

  l.info(f"Commit message:\n--------\n{commit_message}\n--------\nEnd of commit message\n\n")

  return git_commit_all(commit_message)


def publish_to_github_releases(args, tag, commit_hash, release_notes):
  l.info(f"Publishing to GitHub Releases: tag: '{tag}', commit: '{commit_hash}'")
  release_files = list(find_files("*", join(args.temp_dir)))
  github_create_release(
    repo=args.repo,
    version=tag,
    commit_hash=commit_hash,
    release_notes=release_notes,
    files=release_files
  )


def prepare_dataset_release_info(dataset_dir, dataset, last_version, updated_at):
  # modified_files = list(git_get_modified_files(from_revision=last_version, dirs=dataset_dir))
  # modified_files = list(map(lambda f: realpath(f), modified_files))

  # if len(modified_files) == 0:
  #   return None

  path = dataset["path"]
  l.info(f"Preparing release of '{path}'")

  changelog_path = join(dataset_dir, "CHANGELOG.md")

  # if changelog_path not in modified_files:
  #   raise ValueError(
  #     f"Cannot release dataset '{path}' without changelog. Please add or modify file '{changelog_path}': add '## "
  #     f"Unreleased' section and briefly summarize the changes being released"
  #   )

  release_notes = changelog_prepare(dataset, updated_at, changelog_path)

  return {"dataset": dataset, "release_notes": release_notes, "dataset_dir": dataset_dir}


def create_dataset_package(args, dataset, path, tag, dataset_dir):
  files = dict_get_required(dataset, ["files"])
  dict_get_required(files, ["reference"])

  out_dir = join(args.output_dir, path, tag)
  for _, file in files.items():
    inpath = join(dataset_dir, file)
    outpath = join(out_dir, file)
    if file == dict_get(files, ["tree"]):
      # Minify tree.json
      json.dump(json_read(inpath), open(outpath, "w"), separators=(",", ":"), indent=None)
    elif file == "pathogen.json":
      pathogen_json = json_read(inpath)
      pathogen_json["version"] = dataset["version"]
      json_write(pathogen_json, outpath, no_sort_keys=True)
    else:
      copy(inpath, outpath)

  path_safe = path.replace("/", "__")
  zip_filename = join(args.temp_dir, f"{path_safe}__{tag}.zip")
  make_zip(out_dir, zip_filename)

  copy(f"{zip_filename}", join(out_dir, f"dataset.zip"))

  not_found_json = {"status": 404, "message": "Not found"}
  json_write(not_found_json, join(args.output_dir, "404.json"))
  file_write("User-agent: *\nDisallow: /\n", join(args.output_dir, "robots.txt"))

  if tag != "unreleased":
    rmrf(join(args.output_dir, path, "unreleased"))


def get_dataset_name_friendly(dataset):
  return dict_get_required(dataset, ["attributes", "name"])


if __name__ == '__main__':
  main()