/*
 * Copyright (C) 2022 The Android Open Source Project
 *
 * 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
 *
 *      http://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.android.build.shrinker

import com.android.SdkConstants.DOT_9PNG
import com.android.SdkConstants.DOT_PNG
import com.android.SdkConstants.DOT_XML
import com.android.aapt.Resources
import com.android.build.shrinker.DummyContent.TINY_9PNG
import com.android.build.shrinker.DummyContent.TINY_9PNG_CRC
import com.android.build.shrinker.DummyContent.TINY_BINARY_XML
import com.android.build.shrinker.DummyContent.TINY_BINARY_XML_CRC
import com.android.build.shrinker.DummyContent.TINY_PNG
import com.android.build.shrinker.DummyContent.TINY_PNG_CRC
import com.android.build.shrinker.DummyContent.TINY_PROTO_XML
import com.android.build.shrinker.DummyContent.TINY_PROTO_XML_CRC
import com.android.build.shrinker.gatherer.ResourcesGatherer
import com.android.build.shrinker.graph.ResourcesGraphBuilder
import com.android.build.shrinker.obfuscation.ObfuscationMappingsRecorder
import com.android.build.shrinker.usages.ResourceUsageRecorder
import com.android.ide.common.resources.findUnusedResources
import com.android.ide.common.resources.usage.ResourceStore
import com.android.ide.common.resources.usage.ResourceUsageModel.Resource
import com.android.resources.FolderTypeRelationship
import com.android.resources.ResourceFolderType
import com.android.resources.ResourceType
import com.google.common.io.Files
import java.io.BufferedOutputStream
import java.io.File
import java.io.FileOutputStream
import java.io.IOException
import java.io.InputStream
import java.util.jar.JarEntry
import java.util.jar.JarOutputStream
import java.util.zip.CRC32
import java.util.zip.ZipEntry
import java.util.zip.ZipFile

/**
 * Unit that analyzes all resources (after resource merging, compilation and code shrinking has been
 * completed) and figures out which resources are unused, and replaces them with dummy content
 * inside zip archive file.
 *
 * Resource shrinker implementation that allows to customize:
 * <ul>
 * <li>application resource gatherer (from R files, resource tables, etc);
 * <li>recorder for mappings from obfuscated class/methods to original class/methods;
 * <li>sources from which resource usages are recorded (Dex files, compiled JVM classes,
 *   AndroidManifests, etc);
 * <li>resources graph builder that connects resources dependent on each other (analyzing raw
 *   resources content in XML, HTML, CSS, JS, analyzing resource content in proto compiled format);
 * </ul>
 */
private class ResourceShrinkerImpl(
  private val resourcesGatherers: List<ResourcesGatherer>,
  private val obfuscationMappingsRecorder: ObfuscationMappingsRecorder?,
  private val usageRecorders: List<ResourceUsageRecorder>,
  private val graphBuilders: List<ResourcesGraphBuilder>,
  private val debugReporter: ShrinkerDebugReporter,
  private val supportMultipackages: Boolean,
  private val usePreciseShrinking: Boolean,
) : ResourceShrinker {
  private val model = ResourceShrinkerModel(debugReporter, supportMultipackages)
  private lateinit var unused: List<Resource>

  override fun analyze() {
    resourcesGatherers.forEach { it.gatherResourceValues(model) }
    obfuscationMappingsRecorder?.recordObfuscationMappings(model)
    usageRecorders.forEach { it.recordUsages(model) }
    graphBuilders.forEach { it.buildGraph(model) }

    model.resourceStore.processToolsAttributes()
    model.keepPossiblyReferencedResources()

    debugReporter.debug { model.resourceStore.dumpResourceModel() }

    unused =
      findUnusedResources(model.resourceStore.resources) { roots ->
        debugReporter.debug { "The root reachable resources are:" }
        debugReporter.debug { roots.joinToString("\n", transform = { " $it" }) }
      }
    debugReporter.debug { "Unused resources are: " }
    debugReporter.debug { unused.joinToString("\n", transform = { " $it" }) }
  }

  override fun close() {
    debugReporter.close()
  }

  override fun getUnusedResourceCount(): Int {
    return unused.size
  }

  override fun rewriteResourcesInApkFormat(
    source: File,
    dest: File,
    format: LinkedResourcesFormat,
  ) {
    rewriteResourceZip(source, dest, ApkArchiveFormat(model.resourceStore, format))
  }

  override fun rewriteResourcesInBundleFormat(
    source: File,
    dest: File,
    moduleNameToPackageNameMap: Map<String, String>,
  ) {
    rewriteResourceZip(
      source,
      dest,
      BundleArchiveFormat(model.resourceStore, moduleNameToPackageNameMap),
    )
  }

  private fun rewriteResourceZip(source: File, dest: File, format: ArchiveFormat) {
    if (dest.exists() && !dest.delete()) {
      throw IOException("Could not delete $dest")
    }
    JarOutputStream(BufferedOutputStream(FileOutputStream(dest))).use { zos ->
      ZipFile(source).use { zip ->
        // Rather than using Deflater.DEFAULT_COMPRESSION we use 9 here,  since that seems
        // to match the compressed sizes we observe in source .ap_ files encountered by the
        // resource shrinker:
        zos.setLevel(9)
        zip.entries().asSequence().forEach {
          if (format.fileIsNotReachable(it)) {
            // If we don't use precise shrinking we don't remove the files, see:
            // https://b.corp.google.com/issues/37010152
            if (!usePreciseShrinking) {
              replaceWithDummyEntry(zos, it, format.resourcesFormat)
            }
          } else if (it.name.endsWith("resources.pb") && usePreciseShrinking) {
            removeResourceUnusedTableEntries(zip.getInputStream(it), zos, it)
          } else {
            copyToOutput(zip.getInputStream(it), zos, it)
          }
        }
      }
    }
    // If net negative, copy original back. This is unusual, but can happen
    // in some circumstances, such as the one described in
    // https://plus.google.com/+SaidTahsinDane/posts/X9sTSwoVUhB
    // "Removed unused resources: Binary resource data reduced from 588KB to 595KB: Removed -1%"
    // Guard against that, and worst case, just use the original.
    val before = source.length()
    val after = dest.length()
    if (after > before) {
      debugReporter.info {
        "Resource shrinking did not work (grew from $before to $after); using original " + "instead"
      }
      Files.copy(source, dest)
    }
  }

  private fun removeResourceUnusedTableEntries(
    zis: InputStream,
    zos: JarOutputStream,
    srcEntry: ZipEntry,
  ) {
    val resourceIdsToRemove =
      unused.filterNot { it.type == ResourceType.ID }.map { resource -> resource.value }
    val shrunkenResourceTable =
      Resources.ResourceTable.parseFrom(zis).nullOutEntriesWithIds(resourceIdsToRemove)
    val bytes = shrunkenResourceTable.toByteArray()
    val outEntry = JarEntry(srcEntry.name)
    if (srcEntry.time != -1L) {
      outEntry.time = srcEntry.time
    }
    if (srcEntry.method == JarEntry.STORED) {
      outEntry.method = JarEntry.STORED
      outEntry.size = bytes.size.toLong()
      val crc = CRC32()
      crc.update(bytes, 0, bytes.size)
      outEntry.crc = crc.getValue()
    }
    zos.putNextEntry(outEntry)
    zos.write(bytes)
    zos.closeEntry()
  }

  /** Replaces the given entry with a minimal valid file of that type. */
  private fun replaceWithDummyEntry(
    zos: JarOutputStream,
    entry: ZipEntry,
    format: LinkedResourcesFormat,
  ) {
    // Create a new entry so that the compressed len is recomputed.
    val name = entry.name
    val (bytes, crc) =
      when {
        // DOT_9PNG (.9.png) must be always before DOT_PNG (.png)
        name.endsWith(DOT_9PNG) -> TINY_9PNG to TINY_9PNG_CRC
        name.endsWith(DOT_PNG) -> TINY_PNG to TINY_PNG_CRC
        name.endsWith(DOT_XML) && format == LinkedResourcesFormat.BINARY ->
          TINY_BINARY_XML to TINY_BINARY_XML_CRC
        name.endsWith(DOT_XML) && format == LinkedResourcesFormat.PROTO ->
          TINY_PROTO_XML to TINY_PROTO_XML_CRC
        else -> ByteArray(0) to 0L
      }

    val outEntry = JarEntry(name)
    if (entry.time != -1L) {
      outEntry.time = entry.time
    }
    if (entry.method == JarEntry.STORED) {
      outEntry.method = JarEntry.STORED
      outEntry.size = bytes.size.toLong()
      outEntry.crc = crc
    }
    zos.putNextEntry(outEntry)
    zos.write(bytes)
    zos.closeEntry()
    debugReporter.info {
      "Skipped unused resource $name: ${entry.size} bytes (replaced with small dummy file " +
        "of size ${bytes.size} bytes)"
    }
  }

  private fun copyToOutput(zis: InputStream, zos: JarOutputStream, entry: ZipEntry) {
    // We can't just compress all files; files that are not compressed in the source .ap_ file
    // must be left uncompressed here, since for example RAW files need to remain uncompressed
    // in the APK such that they can be mmap'ed at runtime.
    // Preserve the STORED method of the input entry.
    val outEntry =
      when (entry.method) {
        JarEntry.STORED -> JarEntry(entry)
        else -> JarEntry(entry.name)
      }
    if (entry.time != -1L) {
      outEntry.time = entry.time
    }
    zos.putNextEntry(outEntry)
    if (!entry.isDirectory) {
      zis.transferTo(zos)
    }
    zos.closeEntry()
  }
}

private interface ArchiveFormat {
  val resourcesFormat: LinkedResourcesFormat

  fun fileIsNotReachable(entry: ZipEntry): Boolean
}

private class ApkArchiveFormat(
  private val store: ResourceStore,
  override val resourcesFormat: LinkedResourcesFormat,
) : ArchiveFormat {

  override fun fileIsNotReachable(entry: ZipEntry): Boolean {
    if (entry.isDirectory || !entry.name.startsWith("res/")) {
      return false
    }
    val (_, folder, name) = entry.name.split('/', limit = 3)
    return !store.isJarPathReachable(folder, name)
  }
}

private class BundleArchiveFormat(
  private val store: ResourceStore,
  private val moduleNameToPackageName: Map<String, String>,
) : ArchiveFormat {

  override val resourcesFormat = LinkedResourcesFormat.PROTO

  override fun fileIsNotReachable(entry: ZipEntry): Boolean {
    val module = entry.name.substringBefore('/')
    val packageName = moduleNameToPackageName[module]
    if (entry.isDirectory || packageName == null || !entry.name.startsWith("$module/res/")) {
      return false
    }
    val (_, _, folder, name) = entry.name.split('/', limit = 4)
    return !store.isJarPathReachable(folder, name)
  }
}

private fun ResourceStore.isJarPathReachable(folder: String, name: String): Boolean {
  val folderType = ResourceFolderType.getFolderType(folder) ?: return true
  val resourceName = name.substringBefore('.')
  // Bundle format has a restriction: in case the same resource is duplicated in multiple modules
  // its content should be the same in all of them. This restriction means that we can't replace
  // resource with dummy content if its duplicate is used in some module.
  return FolderTypeRelationship.getRelatedResourceTypes(folderType)
    .filterNot { it == ResourceType.ID }
    .flatMap { getResources(it, resourceName) }
    .any { it.isReachable }
}

public fun ResourceStore.isJarPathReachable(path: String): Boolean {
  val (_, folder, name) = path.split('/', limit = 3)
  return isJarPathReachable(folder, name)
}

public fun ResourceStore.getResourcesFor(path: String): List<Resource> {
  val (_, folder, name) = path.split('/', limit = 3)
  val folderType = ResourceFolderType.getFolderType(folder) ?: return emptyList()
  val resourceName = name.substringBefore('.')
  return FolderTypeRelationship.getRelatedResourceTypes(folderType)
    .filterNot { it == ResourceType.ID }
    .flatMap { getResources(it, resourceName) }
    .toList()
}