// Copyright (C) 2025 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 ci_tests

import (
 "cmp"
 "fmt"
 "path/filepath"
 "slices"
 "strings"

 "android/soong/android"
 "android/soong/cc"
 "android/soong/java"

 "github.com/google/blueprint"
 "github.com/google/blueprint/proptools"
)

func init() {
 pctx.Import("android/soong/android")
 registerTestPackageZipBuildComponents(android.InitRegistrationContext)
}

func registerTestPackageZipBuildComponents(ctx android.RegistrationContext) {
 ctx.RegisterModuleType("test_package", TestPackageZipFactory)
}

type testPackageZip struct {
 android.ModuleBase
 android.DefaultableModuleBase
 blueprint.ModuleUsesIncrementalWalkDeps

 properties CITestPackageProperties

 output android.Path
}

type CITestPackageProperties struct {
 // test modules will be added as dependencies using the device os and the common architecture's variant.
 Tests proptools.Configurable[[]string] `android:"arch_variant"`
 // test modules that will be added as dependencies based on the first supported arch variant and the device os variant
 Device_first_tests proptools.Configurable[[]string] `android:"arch_variant"`
 // test modules that will be added as dependencies based on both 32bit and 64bit arch variant and the device os variant
 Device_both_tests proptools.Configurable[[]string] `android:"arch_variant"`
 // test modules that will be added as dependencies based on host
 Host_tests proptools.Configurable[[]string] `android:"arch_variant"`
 // git-main only test modules. Will only be added as dependencies using the device os and the common architecture's variant if exists.
 Tests_if_exist_common proptools.Configurable[[]string] `android:"arch_variant"`
 // git-main only test modules. Will only be added as dependencies based on both 32bit and 64bit arch variant and the device os variant if exists.
 Tests_if_exist_device_both proptools.Configurable[[]string] `android:"arch_variant"`
 // git-main only test modules. Will only be added as dependencies based on the first supported arch variant and the device os variant if exists.
 Tests_if_exist_device_first proptools.Configurable[[]string] `android:"arch_variant"`
 // git-main only test modules. Will only be added as dependencies based on host if exists.
 Tests_if_exist_host proptools.Configurable[[]string] `android:"arch_variant"`
}

type testPackageZipDepTagType struct {
 blueprint.BaseDependencyTag
}

var testPackageZipDepTag testPackageZipDepTagType

var (
 pctx = android.NewPackageContext("android/soong/ci_tests")
)

func (p *testPackageZip) DepsMutator(ctx android.BottomUpMutatorContext) {
 // adding tests property deps
 for _, t := range p.properties.Tests.GetOrDefault(ctx, nil) {
  ctx.AddVariationDependencies(ctx.Config().AndroidCommonTarget.Variations(), testPackageZipDepTag, t)
 }

 // adding device_first_tests property deps
 for _, t := range p.properties.Device_first_tests.GetOrDefault(ctx, nil) {
  ctx.AddVariationDependencies(ctx.Config().AndroidFirstDeviceTarget.Variations(), testPackageZipDepTag, t)
 }

 // adding device_both_tests property deps
 p.addDeviceBothDeps(ctx, false)

 // adding host_tests property deps
 for _, t := range p.properties.Host_tests.GetOrDefault(ctx, nil) {
  ctx.AddVariationDependencies(ctx.Config().BuildOSTarget.Variations(), testPackageZipDepTag, t)
 }
 // adding Tests_if_exist_host property deps
 for _, t := range p.properties.Tests_if_exist_host.GetOrDefault(ctx, nil) {
  if ctx.OtherModuleExists(t) {
   ctx.AddVariationDependencies(ctx.Config().BuildOSTarget.Variations(), testPackageZipDepTag, t)
  }
 }

 // adding Tests_if_exist_* property deps
 for _, t := range p.properties.Tests_if_exist_common.GetOrDefault(ctx, nil) {
  if ctx.OtherModuleExists(t) {
   ctx.AddVariationDependencies(ctx.Config().AndroidCommonTarget.Variations(), testPackageZipDepTag, t)
  }
 }
 for _, t := range p.properties.Tests_if_exist_device_first.GetOrDefault(ctx, nil) {
  if ctx.OtherModuleExists(t) {
   ctx.AddVariationDependencies(ctx.Config().AndroidFirstDeviceTarget.Variations(), testPackageZipDepTag, t)
  }
 }
 p.addDeviceBothDeps(ctx, true)
}

func (p *testPackageZip) addDeviceBothDeps(ctx android.BottomUpMutatorContext, checkIfExist bool) {
 android32TargetList := android.FirstTarget(ctx.Config().Targets[android.Android], "lib32")
 android64TargetList := android.FirstTarget(ctx.Config().Targets[android.Android], "lib64")
 if len(android32TargetList) > 0 {
  maybeAndroid32Target := &android32TargetList[0]
  if checkIfExist {
   for _, t := range p.properties.Tests_if_exist_device_both.GetOrDefault(ctx, nil) {
    if ctx.OtherModuleExists(t) {
     ctx.AddFarVariationDependencies(maybeAndroid32Target.Variations(), testPackageZipDepTag, t)
    }
   }
  } else {
   ctx.AddFarVariationDependencies(maybeAndroid32Target.Variations(), testPackageZipDepTag, p.properties.Device_both_tests.GetOrDefault(ctx, nil)...)
  }
 }
 if len(android64TargetList) > 0 {
  maybeAndroid64Target := &android64TargetList[0]
  if checkIfExist {
   for _, t := range p.properties.Tests_if_exist_device_both.GetOrDefault(ctx, nil) {
    if ctx.OtherModuleExists(t) {
     ctx.AddFarVariationDependencies(maybeAndroid64Target.Variations(), testPackageZipDepTag, t)
    }
   }
  } else {
   ctx.AddFarVariationDependencies(maybeAndroid64Target.Variations(), testPackageZipDepTag, p.properties.Device_both_tests.GetOrDefault(ctx, nil)...)
  }
 }
}

func TestPackageZipFactory() android.Module {
 module := &testPackageZip{}

 module.AddProperties(&module.properties)

 android.InitAndroidArchModule(module, android.DeviceSupported, android.MultilibCommon)
 android.InitDefaultableModule(module)

 return module
}

// We need to implement IsNativeCoverageNeeded so that in coverage builds we depend on the coverage
// variants of the tests. The non-coverage variants will have PreventInstall called on them,
// so they won't be able to be packaged into this test zip.
func (p *testPackageZip) IsNativeCoverageNeeded(ctx cc.IsNativeCoverageNeededContext) bool {
 return ctx.DeviceConfig().NativeCoverageEnabled()
}

var _ cc.UseCoverage = (*testPackageZip)(nil)

func (p *testPackageZip) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 // Never install this test package, it's for disting only
 p.SkipInstall()

 p.output = createOutput(ctx, pctx)

 ctx.SetOutputFiles(android.Paths{p.output}, "")
}

func getAllTestModules(ctx android.ModuleContext) []android.ModuleProxy {
 var ret []android.ModuleProxy
 ctx.WalkDepsProxy(func(child, parent android.ModuleProxy) bool {
  if info, ok := android.OtherModuleProvider(ctx, child, android.CommonModuleInfoProvider); !ok || !info.Enabled {
   return false
  }
  if android.EqualModules(parent, ctx.Module()) && ctx.OtherModuleDependencyTag(child) == testPackageZipDepTag {
   // handle direct deps
   ret = append(ret, child)
   return true
  } else if !android.EqualModules(parent, ctx.Module()) && ctx.OtherModuleDependencyTag(child) == android.RequiredDepTag {
   // handle the "required" from deps
   ret = append(ret, child)
   return true
  } else {
   return false
  }
 })
 ret = android.FirstUniqueInPlace(ret)
 slices.SortFunc(ret, func(a, b android.ModuleProxy) int {
  return cmp.Compare(a.String(), b.String())
 })
 return ret
}

func createOutput(ctx android.ModuleContext, pctx android.PackageContext) android.ModuleOutPath {
 productOut := android.PathForModuleInPartitionInstall(ctx, "").String()
 stagingDir := android.PathForModuleOut(ctx, "STAGING")
 productVariables := ctx.Config().ProductVariables()
 arch := proptools.String(productVariables.DeviceArch)
 secondArch := proptools.String(productVariables.DeviceSecondaryArch)

 builder := android.NewRuleBuilder(pctx, ctx).SandboxDisabled()
 builder.Command().Text("rm").Flag("-rf").Text(stagingDir.String())
 builder.Command().Text("mkdir").Flag("-p").Output(stagingDir)
 builder.Temporary(stagingDir)

 allTestModules := getAllTestModules(ctx)
 for _, dep := range allTestModules {
  extendBuilderCommand(ctx, dep, builder, stagingDir, productOut, arch, secondArch)
 }

 createSymbolsZip(ctx, allTestModules)
 createJacocoJar(ctx, allTestModules)

 output := android.PathForModuleOut(ctx, ctx.ModuleName()+".zip")
 builder.Command().
  BuiltTool("soong_zip").
  Flag("-o").Output(output).
  Flag("-C").Text(stagingDir.String()).
  Flag("-D").Text(stagingDir.String())
 builder.Command().Text("rm").Flag("-rf").Text(stagingDir.String())
 builder.Build("test_package", fmt.Sprintf("build test_package for %s", ctx.ModuleName()))
 return output
}

func createSymbolsZip(ctx android.ModuleContext, allModules []android.ModuleProxy) {
 symbolsZipFile := android.PathForModuleOut(ctx, "symbols.zip")
 symbolsMappingFile := android.PathForModuleOut(ctx, "symbols-mapping.textproto")
 android.BuildSymbolsZip(ctx, allModules, nil, symbolsZipFile, symbolsMappingFile)

 ctx.SetOutputFiles(android.Paths{symbolsZipFile}, ".symbols")
 ctx.SetOutputFiles(android.Paths{symbolsMappingFile}, ".elf_mapping")
}

func createJacocoJar(ctx android.ModuleContext, allModules []android.ModuleProxy) {
 if ctx.Config().JavaCoverageEnabled() {
  jacocoJar := android.PathForModuleOut(ctx, ctx.ModuleName()+"_jacoco_report_classes.jar")
  java.BuildJacocoZip(ctx, allModules, jacocoJar)
  ctx.SetOutputFiles(android.Paths{jacocoJar}, ".jacoco")
 }
}

func extendBuilderCommand(ctx android.ModuleContext, m android.ModuleProxy, builder *android.RuleBuilder, stagingDir android.ModuleOutPath, productOut, arch, secondArch string) {
 info := android.GetModuleInfoJSONInfo(ctx, m)
 if info == nil {
  ctx.OtherModuleErrorf(m, "doesn't set ModuleInfoJSON provider")
 } else if len(info.Data) != 1 {
  ctx.OtherModuleErrorf(m, "doesn't provide exactly one ModuleInfoJSON")
 }

 classes := info.Data[0].GetClass()
 if len(info.Data[0].Class) != 1 {
  ctx.OtherModuleErrorf(m, "doesn't have exactly one class in its ModuleInfoJSON")
 }
 class := strings.ToLower(classes[0])
 if class == "apps" {
  class = "app"
 } else if class == "java_libraries" {
  class = "framework"
 }

 commonInfo, ok := android.OtherModuleProvider(ctx, m, android.CommonModuleInfoProvider)
 if !ok || commonInfo.InstallFiles == nil {
  ctx.ModuleErrorf("Module %s doesn't set InstallFilesProvider", m.Name())
 }
 installedFilesInfo := commonInfo.InstallFiles
 for _, spec := range installedFilesInfo.PackagingSpecs {
  if spec.SrcPath() == nil {
   // Probably a symlink
   continue
  }
  installedFile := spec.FullInstallPath()

  ext := installedFile.Ext()
  // there are additional installed files for some app-class modules, we only need the .apk, .odex and .vdex files in the test package
  excludeInstalledFile := ext != ".apk" && ext != ".odex" && ext != ".vdex"
  if class == "app" && excludeInstalledFile {
   continue
  }
  // only .jar files should be included for a framework dep
  if class == "framework" && ext != ".jar" {
   continue
  }
  name := removeFileExtension(installedFile.Base())
  // some apks have other apk as installed files, these additional files shouldn't be included
  // But due to for override_android_test or override_android_app it will have OtherModuleName deps for the module
  // it wants to override, to prevent it being ignored only skip this deps if it not the direct dependency of the
  // test_packages.
  isAppOrFramework := class == "app" || class == "framework"
  if ctx.OtherModuleDependencyTag(m) != testPackageZipDepTag && isAppOrFramework && name != ctx.OtherModuleName(m) {
   continue
  }

  f := strings.TrimPrefix(installedFile.String(), productOut+"/")
  if strings.HasPrefix(f, "out") {
   continue
  }
  if strings.HasPrefix(f, "system/") {
   f = strings.Replace(f, "system/", "DATA/", 1)
  }
  f = strings.ReplaceAll(f, filepath.Join("testcases", name, arch), filepath.Join("DATA", class, name))
  f = strings.ReplaceAll(f, filepath.Join("testcases", name, secondArch), filepath.Join("DATA", class, name))
  if strings.HasPrefix(f, "testcases") {
   f = strings.Replace(f, "testcases", filepath.Join("DATA", class), 1)
  }
  if strings.HasPrefix(f, "data/") {
   f = strings.Replace(f, "data/", "DATA/", 1)
  }
  f = strings.ReplaceAll(f, "DATA_other", "system_other")
  f = strings.ReplaceAll(f, "system_other/DATA", "system_other/system")
  dir := filepath.Dir(f)

  tempOut := android.PathForModuleOut(ctx, "STAGING", f)
  builder.Command().Text("mkdir").Flag("-p").Text(filepath.Join(stagingDir.String(), dir))
  // Copy srcPath instead of installedFile because some rules like target-files.zip
  // are non-hermetic and would be affected if we built the installed files.
  builder.Command().Text("cp").Flag("-Rf").Input(spec.SrcPath()).Output(tempOut)
  builder.Temporary(tempOut)
 }
}

func removeFileExtension(filename string) string {
 return strings.TrimSuffix(filename, filepath.Ext(filename))
}

// The only purpose of this method is to make sure we can build the module directly
// without adding suffix "-soong"
func (p *testPackageZip) AndroidMkEntries() []android.AndroidMkEntries {
 return []android.AndroidMkEntries{
  android.AndroidMkEntries{
   Class:      "ETC",
   OutputFile: android.OptionalPathForPath(p.output),
   ExtraEntries: []android.AndroidMkExtraEntriesFunc{
    func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
     entries.SetBool("LOCAL_UNINSTALLABLE_MODULE", true)
    }},
  },
 }
}