//    Copyright 2019 Google LLC
//
//    Licensed under the Apache License, Version 2.0 (the "License");
//    you may not use this file except in compliance with the License.
//    You may obtain a copy of the License at
//
//        https://www.apache.org/licenses/LICENSE-2.0
//
//    Unless required by applicable law or agreed to in writing, software
//    distributed under the License is distributed on an "AS IS" BASIS,
//    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//    See the License for the specific language governing permissions and
//    limitations under the License.

package com.google.fhir.protogen;

import static java.nio.charset.StandardCharsets.UTF_8;

import com.google.common.base.Ascii;
import com.google.common.io.ByteStreams;
import com.google.fhir.common.InvalidFhirException;
import com.google.fhir.common.JsonFormat;
import com.google.fhir.proto.Annotations.FhirVersion;
import com.google.fhir.proto.PackageInfo;
import com.google.fhir.r4.core.CodeSystem;
import com.google.fhir.r4.core.SearchParameter;
import com.google.fhir.r4.core.StructureDefinition;
import com.google.fhir.r4.core.ValueSet;
import com.google.gson.JsonArray;
import com.google.gson.JsonElement;
import com.google.gson.JsonParser;
import com.google.gson.stream.JsonReader;
import com.google.protobuf.TextFormat;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.StringReader;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.function.Predicate;
import org.apache.commons.compress.archivers.ArchiveEntry;
import org.apache.commons.compress.archivers.ArchiveInputStream;
import org.apache.commons.compress.archivers.tar.TarArchiveInputStream;
import org.apache.commons.compress.archivers.zip.ZipArchiveInputStream;
import org.apache.commons.compress.compressors.gzip.GzipCompressorInputStream;

/**
 * Represents a FHIR Proto package.
 *
 * <p>Includes resource definitions and package information.
 *
 * <p>Constructed from an archive file (ZIP or TAR) containing the JSON files of the resources and
 * optionally a PackageInfo file that describes the package. ZIP archives are generated by the
 * `fhir_package` rule in protogen.bzl.
 *
 * <p>TODO(b/235876918): Support NPM-based FHIR Packages for proto generation.
 */
public class FhirPackage {
  public final PackageInfo packageInfo;
  private final ResourceCollection<StructureDefinition> structureDefinitionCollection;
  private final ResourceCollection<SearchParameter> searchParameterCollection;
  private final ResourceCollection<CodeSystem> codeSystemCollection;
  private final ResourceCollection<ValueSet> valueSetCollection;

  public final Iterable<StructureDefinition> structureDefinitions() {
    return structureDefinitionCollection;
  }

  public final Iterable<SearchParameter> searchParameters() {
    return searchParameterCollection;
  }

  public final Iterable<CodeSystem> codeSystems() {
    return codeSystemCollection;
  }

  public final Iterable<ValueSet> valueSets() {
    return valueSetCollection;
  }

  private FhirPackage(
      PackageInfo packageInfo,
      ResourceCollection<StructureDefinition> structureDefinitions,
      ResourceCollection<SearchParameter> searchParameters,
      ResourceCollection<CodeSystem> codeSystems,
      ResourceCollection<ValueSet> valueSets) {
    this.packageInfo = packageInfo;
    this.structureDefinitionCollection = structureDefinitions;
    this.searchParameterCollection = searchParameters;
    this.codeSystemCollection = codeSystems;
    this.valueSetCollection = valueSets;
  }

  @Override
  public boolean equals(Object other) {
    return other instanceof FhirPackage
        && ((FhirPackage) other)
            .packageInfo
            .getProtoPackage()
            .equals(packageInfo.getProtoPackage());
  }

  @Override
  public int hashCode() {
    return packageInfo.getProtoPackage().hashCode();
  }

  public static boolean isCorePackage(PackageInfo packageInfo) {
    return packageInfo
        .getProtoPackage()
        .equals(
            com.google.fhir.common.FhirVersion.fromAnnotation(packageInfo.getFhirVersion())
                .coreProtoPackage);
  }

  public boolean isCorePackage() {
    return isCorePackage(packageInfo);
  }

  /**
   * Loads a package archive into a FhirPackage.
   *
   * <p>Parses all defining resources that are part of the spec.
   *
   * <p>A PackageInfo can be provided and must be named with the suffix "package_info.prototxt" or
   * "package_info.textproto".
   *
   * <p>To side load a PackageInfo proto, use the version of this that accepts a PackageInfo proto.
   *
   * <p>TODO(b/235876918): Support NPM-based FHIR Packages for proto generation.
   *
   * @param archiveFilePath The absolute path to the archive file (ZIP or TAR) that is loaded.
   *     Expected to end with ".zip" "tar.gz" or ".tgz".
   */
  public static FhirPackage load(String archiveFilePath) throws IOException, InvalidFhirException {
    return load(archiveFilePath, /* packageInfo= */ null);
  }

  /**
   * Loads a package ZIP into a FhirPackage.
   *
   * <p>Parses all defining resources that are part of the spec.
   *
   * <p>To read PackageInfo out of the ZIP itself, use the version of this that does not accept a
   * PackageProto argument.
   *
   * @param archiveFilePath The absolute path to the archive file (ZIP or TAR) that is loaded.
   *     Expected to end with ".zip" "tar.gz" or ".tgz".
   * @param packageInfo The package information to load. This package information is used,
   *     irrespective of whether the ZIP contains one.
   */
  public static FhirPackage load(String archiveFilePath, PackageInfo packageInfo)
      throws IOException, InvalidFhirException {
    PackageInfo extractedPackageInfo = null;
    try (ArchiveInputStream archiveEntries = getZipOrTarInputStream(archiveFilePath)) {
      List<JsonFile> jsonFiles = new ArrayList<>();
      ArchiveEntry entry = null;
      while ((entry = archiveEntries.getNextEntry()) != null) {
        String entryName = Ascii.toLowerCase(entry.getName());
        if (entryName.endsWith(".json")) {
          JsonReader reader =
              new JsonReader(
                  new StringReader(new String(ByteStreams.toByteArray(archiveEntries), UTF_8)));
          JsonElement json = JsonParser.parseReader(reader);
          jsonFiles.add(new JsonFile(entry.getName(), json));
        } else if (packageInfo == null
            && (entryName.endsWith("package_info.prototxt")
                || entryName.endsWith("package_info.textproto"))) {
          if (extractedPackageInfo != null) {
            throw new IllegalArgumentException(
                "Only one PackageInfo should be provided: " + archiveFilePath);
          }
          String protoTxt = new String(ByteStreams.toByteArray(archiveEntries), UTF_8);
          PackageInfo.Builder infoBuilder = PackageInfo.newBuilder();
          TextFormat.getParser().merge(protoTxt, infoBuilder);
          extractedPackageInfo = infoBuilder.build();
        }
      }

      return makeFromJsonAndPackageInfo(
          jsonFiles, packageInfo == null ? extractedPackageInfo : packageInfo);
    }
  }

  /**
   * Filters the package to only serve resources (`StructureDefinition`-s) matching the provided
   * `filter`.
   *
   * <p>Filters are not additive. In other words, if this method is called multiple times only the
   * latest `filter` is applied.
   */
  FhirPackage filterResources(Predicate<StructureDefinition> filter) {
    structureDefinitionCollection.setFilter(filter);
    return new FhirPackage(
        packageInfo,
        structureDefinitionCollection,
        searchParameterCollection,
        codeSystemCollection,
        valueSetCollection);
  }

  /**
   * Returns the `StructureDefinition` from the package with the specified `uri` or an empty
   * `Optional` if this package does not contain a resource for the `uri`.
   *
   * @param uri of the `StructureDefinition` to retrieve as defined in the package.
   */
  public Optional<StructureDefinition> getStructureDefinition(String uri)
      throws InvalidFhirException {
    return structureDefinitionCollection.get(uri);
  }

  /**
   * Returns the `SearchParameter` from the package with the specified `uri` or an empty `Optional`
   * if this package does not contain a resource for the `uri`.
   *
   * @param uri of the `SearchParameter` to retrieve as defined in the package.
   */
  public Optional<SearchParameter> getSearchParameter(String uri) throws InvalidFhirException {
    return searchParameterCollection.get(uri);
  }

  /**
   * Returns the `CodeSystem` from the package with the specified `uri` or an empty `Optional` if
   * this package does not contain a resource for the `uri`.
   *
   * @param uri of the `CodeSystem` to retrieve as defined in the package.
   */
  public Optional<CodeSystem> getCodeSystem(String uri) throws InvalidFhirException {
    return codeSystemCollection.get(uri);
  }

  /**
   * Returns the `ValueSet` from the package with the specified `uri` or an empty `Optional` if this
   * package does not contain a resource for the `uri`.
   *
   * @param uri of the `ValueSet` to retrieve as defined in the package.
   */
  public Optional<ValueSet> getValueSet(String uri) throws InvalidFhirException {
    return valueSetCollection.get(uri);
  }

  private static class JsonFile {
    JsonFile(String name, JsonElement json) {
      this.name = name;
      this.json = json;
    }

    final String name;
    final JsonElement json;
  }

  private static class ResourceCollections {
    ResourceCollection<StructureDefinition> structureDefinitions;
    ResourceCollection<SearchParameter> searchParameters;
    ResourceCollection<CodeSystem> codeSystems;
    ResourceCollection<ValueSet> valueSets;
  }

  /**
   * Gets the archive input stream from a ZIP or TAR.GZ file.
   *
   * @param archiveFilePath The absolute path to the archive file (ZIP or TAR) that is loaded.
   *     Expected to end with ".zip" "tar.gz" or ".tgz".
   */
  private static ArchiveInputStream getZipOrTarInputStream(String archiveFilePath)
      throws IOException {
    FileInputStream fileStream = new FileInputStream(archiveFilePath);
    if (archiveFilePath.endsWith(".zip")) {
      return new ZipArchiveInputStream(fileStream);
    } else if (archiveFilePath.endsWith(".tar.gz") || archiveFilePath.endsWith(".tgz")) {
      return new TarArchiveInputStream(new GzipCompressorInputStream(fileStream));
    } else {
      throw new IllegalArgumentException(
          "`archiveFilePath` must end with '.zip', 'tar.gz' or '.tgz': " + archiveFilePath);
    }
  }

  private static Optional<String> getResourceType(JsonElement json) {
    if (!json.isJsonObject() || !json.getAsJsonObject().has("resourceType")) {
      return Optional.empty();
    }
    return Optional.of(json.getAsJsonObject().get("resourceType").getAsString());
  }

  private static void buildResourceCollections(
      List<JsonFile> jsonFiles, JsonFormat.Parser parser, ResourceCollections resourceCollections) {

    for (JsonFile jsonFile : jsonFiles) {
      JsonElement json = jsonFile.json;
      Optional<String> expectedType = getResourceType(json);
      if (expectedType.isEmpty()) {
        System.out.println("Unhandled JSON entry: " + jsonFile.name);
        continue;
      }
      switch (expectedType.get()) {
        case "StructureDefinition":
          resourceCollections.structureDefinitions.put(json);
          break;
        case "SearchParameter":
          resourceCollections.searchParameters.put(json);
          break;
        case "CodeSystem":
          resourceCollections.codeSystems.put(json);
          break;
        case "ValueSet":
          resourceCollections.valueSets.put(json);
          break;
        case "Bundle":
          JsonElement bundleEntry;
          if (!json.isJsonObject()
              || !json.getAsJsonObject().has("entry")
              || !(bundleEntry = json.getAsJsonObject().get("entry")).isJsonArray()) {
            break;
          }
          JsonArray entries = bundleEntry.getAsJsonArray();
          List<JsonFile> bundleEntries = new ArrayList<>();
          for (JsonElement entry : entries) {
            bundleEntries.add(new JsonFile(jsonFile.name, entry.getAsJsonObject().get("resource")));
          }
          buildResourceCollections(bundleEntries, parser, resourceCollections);
          break;
        default:
          System.out.println("Unhandled JSON entry: " + jsonFile.name);
      }
    }
  }

  private static FhirPackage makeFromJsonAndPackageInfo(
      List<JsonFile> jsonFiles, PackageInfo packageInfo) {
    if (packageInfo != null) {
      if (packageInfo.getProtoPackage().isEmpty()) {
        throw new IllegalArgumentException(
            "When PackageInfo is provided, must specify `proto_package`.");
      }
      if (packageInfo.getFhirVersion() == FhirVersion.FHIR_VERSION_UNKNOWN) {
        throw new IllegalArgumentException(
            "When PackageInfo is provided, must specify `fhir_version`.");
      }
    }

    JsonFormat.Parser parser = JsonFormat.getParser();

    ResourceCollections resourceCollections = new ResourceCollections();
    resourceCollections.structureDefinitions =
        new ResourceCollection<>(parser, StructureDefinition.class);
    resourceCollections.searchParameters = new ResourceCollection<>(parser, SearchParameter.class);
    resourceCollections.codeSystems = new ResourceCollection<>(parser, CodeSystem.class);
    resourceCollections.valueSets = new ResourceCollection<>(parser, ValueSet.class);

    buildResourceCollections(jsonFiles, parser, resourceCollections);

    return new FhirPackage(
        packageInfo,
        resourceCollections.structureDefinitions,
        resourceCollections.searchParameters,
        resourceCollections.codeSystems,
        resourceCollections.valueSets);
  }
}
