// Copyright 2024 Google Inc. All rights reserved.
//
// 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 cc

import (
 "bytes"
 _ "embed"
 "fmt"
 "path/filepath"
 "slices"
 "sort"
 "strings"
 "text/template"

 "android/soong/android"

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

const veryVerbose bool = false

//go:embed cmake_main.txt
var templateCmakeMainRaw string
var templateCmakeMain *template.Template = parseTemplate(templateCmakeMainRaw)

//go:embed cmake_module_cc.txt
var templateCmakeModuleCcRaw string
var templateCmakeModuleCc *template.Template = parseTemplate(templateCmakeModuleCcRaw)

//go:embed cmake_module_aidl.txt
var templateCmakeModuleAidlRaw string
var templateCmakeModuleAidl *template.Template = parseTemplate(templateCmakeModuleAidlRaw)

//go:embed cmake_ext_add_aidl_library.txt
var cmakeExtAddAidlLibrary string

//go:embed cmake_ext_append_flags.txt
var cmakeExtAppendFlags string

var defaultUnportableFlags []string = []string{
 "-Wno-c99-designator",
 "-Wno-class-memaccess",
 "-Wno-exit-time-destructors",
 "-Winconsistent-missing-override",
 "-Wno-inconsistent-missing-override",
 "-Wreorder-init-list",
 "-Wno-reorder-init-list",
 "-Wno-restrict",
 "-Wno-stringop-overread",
 "-Wno-subobject-linkage",
}

var ignoredSystemLibs []string = []string{
 "crtbegin_dynamic",
 "crtend_android",
 "libc",
 "libc++",
 "libc++_static",
 "libc++demangle",
 "libc_musl",
 "libc_musl_crtbegin_so",
 "libc_musl_crtbegin_static",
 "libc_musl_crtend",
 "libc_musl_crtend_so",
 "libdl",
 "libm",
 "prebuilt_libclang_rt.builtins",
 "prebuilt_libclang_rt.ubsan_minimal",
}

// Mapping entry between Android's library name and the one used when building outside Android tree.
type LibraryMappingProperty struct {
 // Android library name.
 Android_name string

 // Library name used when building outside Android.
 Mapped_name string

 // If the make file is already present in Android source tree, specify its location.
 Package_pregenerated string

 // If the package is expected to be installed on the build host OS, specify its name.
 Package_system string
}

type CmakeSnapshotProperties struct {
 // Host modules to add to the snapshot package. Their dependencies are pulled in automatically.
 Modules_host []string

 // System modules to add to the snapshot package. Their dependencies are pulled in automatically.
 Modules_system []string

 // Vendor modules to add to the snapshot package. Their dependencies are pulled in automatically.
 Modules_vendor []string

 // Host prebuilts to bundle with the snapshot. These are tools needed to build outside Android.
 Prebuilts []string

 // Global cflags to add when building outside Android.
 Cflags []string

 // Flags to skip when building outside Android.
 Cflags_ignored []string

 // Mapping between library names used in Android tree and externally.
 Library_mapping []LibraryMappingProperty

 // List of cflags that are not portable between compilers that could potentially be used to
 // build a generated package. If left empty, it's initialized with a default list.
 Unportable_flags []string

 // Whether to include source code as part of the snapshot package.
 Include_sources bool
}

var cmakeSnapshotSourcesProvider = blueprint.NewProvider[android.Paths]()

type CmakeSnapshot struct {
 android.ModuleBase
 blueprint.ModuleUsesIncrementalWalkDeps

 Properties CmakeSnapshotProperties

 zipPath android.WritablePath
}

type cmakeProcessedProperties struct {
 LibraryMapping       map[string]LibraryMappingProperty
 PregeneratedPackages []string
 SystemPackages       []string
}

type cmakeSnapshotDependencyTag struct {
 blueprint.BaseDependencyTag
 name string
}

var (
 cmakeSnapshotModuleTag   = cmakeSnapshotDependencyTag{name: "cmake-snapshot-module"}
 cmakeSnapshotPrebuiltTag = cmakeSnapshotDependencyTag{name: "cmake-snapshot-prebuilt"}
)

func parseTemplate(templateContents string) *template.Template {
 funcMap := template.FuncMap{
  "setList": func(name string, nameSuffix string, itemPrefix string, items []string) string {
   var list strings.Builder
   list.WriteString("set(" + name + nameSuffix)
   templateListBuilder(&list, itemPrefix, items)
   return list.String()
  },
  "toStrings": func(files android.Paths) []string {
   return files.Strings()
  },
  "concat5": func(list1 []string, list2 []string, list3 []string, list4 []string, list5 []string) []string {
   return append(append(append(append(list1, list2...), list3...), list4...), list5...)
  },
  "cflagsList": func(name string, nameSuffix string, flags []string,
   unportableFlags []string, ignoredFlags []string) string {
   if len(unportableFlags) == 0 {
    unportableFlags = defaultUnportableFlags
   }

   var filteredPortable []string
   var filteredUnportable []string
   for _, flag := range flags {
    if slices.Contains(ignoredFlags, flag) {
     continue
    } else if slices.Contains(unportableFlags, flag) {
     filteredUnportable = append(filteredUnportable, flag)
    } else {
     filteredPortable = append(filteredPortable, flag)
    }
   }

   var list strings.Builder

   list.WriteString("set(" + name + nameSuffix)
   templateListBuilder(&list, "", filteredPortable)

   if len(filteredUnportable) > 0 {
    list.WriteString("\nappend_cxx_flags_if_supported(" + name + nameSuffix)
    templateListBuilder(&list, "", filteredUnportable)
   }

   return list.String()
  },
  "getSources": func(ctx android.ModuleContext, info *CcInfo) android.Paths {
   return info.CompilerInfo.Srcs
  },
  "getModuleType": getModuleType,
  "getAidlInterface": func(info *CcInfo) AidlInterfaceInfo {
   return info.CompilerInfo.AidlInterfaceInfo
  },
  "getCflagsProperty": func(ctx android.ModuleContext, info *CcInfo) []string {
   return info.CompilerInfo.Cflags
  },
  "getWholeStaticLibsProperty": func(ctx android.ModuleContext, info *CcInfo) []string {
   return info.LinkerInfo.WholeStaticLibs
  },
  "getStaticLibsProperty": func(ctx android.ModuleContext, info *CcInfo) []string {
   return info.LinkerInfo.StaticLibs
  },
  "getSharedLibsProperty": func(ctx android.ModuleContext, info *CcInfo) []string {
   return info.LinkerInfo.SharedLibs
  },
  "getHeaderLibsProperty": func(ctx android.ModuleContext, info *CcInfo) []string {
   return info.LinkerInfo.HeaderLibs
  },
  "getExtraLibs":   getExtraLibs,
  "getIncludeDirs": getIncludeDirs,
  "mapLibraries": func(ctx android.ModuleContext, m android.ModuleProxy, libs []string, mapping map[string]LibraryMappingProperty) []string {
   var mappedLibs []string
   for _, lib := range libs {
    mappedLib, exists := mapping[lib]
    if exists {
     lib = mappedLib.Mapped_name
    } else {
     if !ctx.OtherModuleExists(lib) {
      ctx.OtherModuleErrorf(m, "Dependency %s doesn't exist", lib)
     }
     lib = "android::" + lib
    }
    if lib == "" {
     continue
    }
    mappedLibs = append(mappedLibs, lib)
   }
   sort.Strings(mappedLibs)
   mappedLibs = slices.Compact(mappedLibs)
   return mappedLibs
  },
  "getAidlSources": func(info *CcInfo) []string {
   aidlInterface := info.CompilerInfo.AidlInterfaceInfo
   aidlRoot := aidlInterface.AidlRoot + string(filepath.Separator)
   if aidlInterface.AidlRoot == "" {
    aidlRoot = ""
   }
   var sources []string
   for _, src := range aidlInterface.Sources {
    if !strings.HasPrefix(src, aidlRoot) {
     panic(fmt.Sprintf("Aidl source '%v' doesn't start with '%v'", src, aidlRoot))
    }
    sources = append(sources, src[len(aidlRoot):])
   }
   return sources
  },
 }

 return template.Must(template.New("").Delims("<<", ">>").Funcs(funcMap).Parse(templateContents))
}

func sliceWithPrefix(prefix string, slice []string) []string {
 output := make([]string, len(slice))
 for i, elem := range slice {
  output[i] = prefix + elem
 }
 return output
}

func templateListBuilder(builder *strings.Builder, itemPrefix string, items []string) {
 if len(items) > 0 {
  builder.WriteString("\n")
  for _, item := range items {
   builder.WriteString("    " + itemPrefix + item + "\n")
  }
 }
 builder.WriteString(")")
}

func executeTemplate(templ *template.Template, buffer *bytes.Buffer, data any) string {
 buffer.Reset()
 if err := templ.Execute(buffer, data); err != nil {
  panic(err)
 }
 output := strings.TrimSpace(buffer.String())
 buffer.Reset()
 return output
}

func (m *CmakeSnapshot) DepsMutator(ctx android.BottomUpMutatorContext) {
 deviceVariations := ctx.Config().AndroidFirstDeviceTarget.Variations()
 deviceSystemVariations := append(deviceVariations, blueprint.Variation{"image", ""})
 deviceVendorVariations := append(deviceVariations, blueprint.Variation{"image", "vendor"})
 hostVariations := ctx.Config().BuildOSTarget.Variations()

 ctx.AddVariationDependencies(hostVariations, cmakeSnapshotModuleTag, m.Properties.Modules_host...)
 ctx.AddVariationDependencies(deviceSystemVariations, cmakeSnapshotModuleTag, m.Properties.Modules_system...)
 ctx.AddVariationDependencies(deviceVendorVariations, cmakeSnapshotModuleTag, m.Properties.Modules_vendor...)

 if len(m.Properties.Prebuilts) > 0 {
  prebuilts := append(m.Properties.Prebuilts, "libc++")
  ctx.AddVariationDependencies(hostVariations, cmakeSnapshotPrebuiltTag, prebuilts...)
 }
}

func (m *CmakeSnapshot) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 var templateBuffer bytes.Buffer
 var pprop cmakeProcessedProperties
 m.zipPath = android.PathForModuleOut(ctx, ctx.ModuleName()+".zip")

 // Process Library_mapping for more efficient lookups
 pprop.LibraryMapping = map[string]LibraryMappingProperty{}
 for _, elem := range m.Properties.Library_mapping {
  pprop.LibraryMapping[elem.Android_name] = elem

  if elem.Package_pregenerated != "" {
   pprop.PregeneratedPackages = append(pprop.PregeneratedPackages, elem.Package_pregenerated)
  }
  sort.Strings(pprop.PregeneratedPackages)
  pprop.PregeneratedPackages = slices.Compact(pprop.PregeneratedPackages)

  if elem.Package_system != "" {
   pprop.SystemPackages = append(pprop.SystemPackages, elem.Package_system)
  }
  sort.Strings(pprop.SystemPackages)
  pprop.SystemPackages = slices.Compact(pprop.SystemPackages)
 }

 // Generating CMakeLists.txt rules for all modules in dependency tree
 moduleDirs := map[string][]string{}
 sourceFiles := map[string]android.Path{}
 visitedModules := map[string]bool{}
 var pregeneratedModules []android.ModuleProxy
 ctx.WalkDepsProxy(func(dep, parent android.ModuleProxy) bool {
  moduleName := ctx.OtherModuleName(dep)
  if visited := visitedModules[moduleName]; visited {
   return false // visit only once
  }
  visitedModules[moduleName] = true
  ccInfo, ok := android.OtherModuleProvider(ctx, dep, CcInfoProvider)
  if !ok {
   return false // not a cc module
  }
  if mapping, ok := pprop.LibraryMapping[moduleName]; ok {
   if mapping.Package_pregenerated != "" {
    pregeneratedModules = append(pregeneratedModules, dep)
   }
   return false // mapped to system or pregenerated (we'll handle these later)
  }
  if ctx.OtherModuleDependencyTag(dep) == cmakeSnapshotPrebuiltTag {
   return false // we'll handle cmakeSnapshotPrebuiltTag later
  }
  if slices.Contains(ignoredSystemLibs, moduleName) {
   return false // system libs built in-tree for Android
  }
  if ccInfo.IsPrebuilt {
   return false // prebuilts are not supported
  }
  if ccInfo.CompilerInfo == nil {
   return false // unsupported module type
  }
  isAidlModule := ccInfo.CompilerInfo.AidlInterfaceInfo.Lang != ""

  if !ccInfo.CmakeSnapshotSupported {
   ctx.OtherModulePropertyErrorf(dep, "cmake_snapshot_supported",
    "CMake snapshots not supported, despite being a dependency for %s",
    ctx.OtherModuleName(parent))
   return false
  }

  if veryVerbose {
   fmt.Println("WalkDeps: " + ctx.OtherModuleName(parent) + " -> " + moduleName)
  }

  // Generate CMakeLists.txt fragment for this module
  templateToUse := templateCmakeModuleCc
  if isAidlModule {
   templateToUse = templateCmakeModuleAidl
  }
  moduleFragment := executeTemplate(templateToUse, &templateBuffer, struct {
   Ctx      *android.ModuleContext
   M        android.ModuleProxy
   CcInfo   *CcInfo
   Snapshot *CmakeSnapshot
   Pprop    *cmakeProcessedProperties
  }{
   &ctx,
   dep,
   ccInfo,
   m,
   &pprop,
  })
  moduleDir := ctx.OtherModuleDir(dep)
  moduleDirs[moduleDir] = append(moduleDirs[moduleDir], moduleFragment)

  if m.Properties.Include_sources {
   files, _ := android.OtherModuleProvider(ctx, dep, cmakeSnapshotSourcesProvider)
   for _, file := range files {
    sourceFiles[file.String()] = file
   }
  }

  // if it's AIDL module, no need to dive into their dependencies
  return !isAidlModule
 })

 // Enumerate sources for pregenerated modules
 if m.Properties.Include_sources {
  for _, dep := range pregeneratedModules {
   if !android.OtherModuleProviderOrDefault(ctx, dep, CcInfoProvider).CmakeSnapshotSupported {
    ctx.OtherModulePropertyErrorf(dep, "cmake_snapshot_supported",
     "Pregenerated CMake snapshots not supported, despite being requested for %s",
     ctx.ModuleName())
    continue
   }

   files, _ := android.OtherModuleProvider(ctx, dep, cmakeSnapshotSourcesProvider)
   for _, file := range files {
    sourceFiles[file.String()] = file
   }
  }
 }

 // Merging CMakeLists.txt contents for every module directory
 var makefilesList android.Paths
 for _, moduleDir := range android.SortedKeys(moduleDirs) {
  fragments := moduleDirs[moduleDir]
  moduleCmakePath := android.PathForModuleGen(ctx, moduleDir, "CMakeLists.txt")
  makefilesList = append(makefilesList, moduleCmakePath)
  sort.Strings(fragments)
  android.WriteFileRule(ctx, moduleCmakePath, strings.Join(fragments, "\n\n\n"))
 }

 // Generating top-level CMakeLists.txt
 mainCmakePath := android.PathForModuleGen(ctx, "CMakeLists.txt")
 makefilesList = append(makefilesList, mainCmakePath)
 mainContents := executeTemplate(templateCmakeMain, &templateBuffer, struct {
  Ctx        *android.ModuleContext
  M          *CmakeSnapshot
  ModuleDirs map[string][]string
  Pprop      *cmakeProcessedProperties
 }{
  &ctx,
  m,
  moduleDirs,
  &pprop,
 })
 android.WriteFileRule(ctx, mainCmakePath, mainContents)

 // Generating CMake extensions
 extPath := android.PathForModuleGen(ctx, "cmake", "AppendCxxFlagsIfSupported.cmake")
 makefilesList = append(makefilesList, extPath)
 android.WriteFileRuleVerbatim(ctx, extPath, cmakeExtAppendFlags)
 extPath = android.PathForModuleGen(ctx, "cmake", "AddAidlLibrary.cmake")
 makefilesList = append(makefilesList, extPath)
 android.WriteFileRuleVerbatim(ctx, extPath, cmakeExtAddAidlLibrary)

 // Generating the final zip file
 zipRule := android.NewRuleBuilder(pctx, ctx).SandboxDisabled()
 zipCmd := zipRule.Command().
  BuiltTool("soong_zip").
  FlagWithOutput("-o ", m.zipPath)

 // Packaging all sources into the zip file
 if m.Properties.Include_sources {
  var sourcesList android.Paths
  for _, file := range android.SortedKeys(sourceFiles) {
   path := sourceFiles[file]
   sourcesList = append(sourcesList, path)
  }

  sourcesRspFile := android.PathForModuleObj(ctx, ctx.ModuleName()+"_sources.rsp")
  zipCmd.FlagWithRspFileInputList("-r ", sourcesRspFile, sourcesList)
 }

 // Packaging all make files into the zip file
 makefilesRspFile := android.PathForModuleObj(ctx, ctx.ModuleName()+"_makefiles.rsp")
 zipCmd.
  FlagWithArg("-C ", android.PathForModuleGen(ctx).String()).
  FlagWithRspFileInputList("-r ", makefilesRspFile, makefilesList)

 // Packaging all prebuilts into the zip file
 if len(m.Properties.Prebuilts) > 0 {
  var prebuiltsList android.Paths

  ctx.VisitDirectDepsProxyWithTag(cmakeSnapshotPrebuiltTag, func(dep android.ModuleProxy) {
   for _, file := range android.GetInstallFiles(
    ctx, dep).InstallFiles {
    prebuiltsList = append(prebuiltsList, file)
   }
  })

  prebuiltsRspFile := android.PathForModuleObj(ctx, ctx.ModuleName()+"_prebuilts.rsp")
  zipCmd.
   FlagWithArg("-C ", android.PathForArbitraryOutput(ctx).String()).
   FlagWithArg("-P ", "prebuilts").
   FlagWithRspFileInputList("-r ", prebuiltsRspFile, prebuiltsList)
 }

 // Finish generating the final zip file
 zipRule.Build(m.zipPath.String(), "archiving "+ctx.ModuleName())

 ctx.SetOutputFiles(android.Paths{m.zipPath}, "")

 moduleInfoJSON := ctx.ModuleInfoJSON()
 moduleInfoJSON.Class = []string{"DATA"}
 moduleInfoJSON.SystemSharedLibs = []string{"none"}
}

func (m *CmakeSnapshot) AndroidMkEntries() []android.AndroidMkEntries {
 return []android.AndroidMkEntries{{
  Class:      "DATA",
  OutputFile: android.OptionalPathForPath(m.zipPath),
  ExtraEntries: []android.AndroidMkExtraEntriesFunc{
   func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
    entries.SetBool("LOCAL_UNINSTALLABLE_MODULE", true)
   },
  },
 }}
}

func getModuleType(info *CcInfo) string {
 if info.LinkerInfo.BinaryDecoratorInfo != nil {
  return "executable"
 } else if info.LinkerInfo.LibraryDecoratorInfo != nil {
  return "library"
 } else if info.LinkerInfo.TestBinaryInfo != nil || info.LinkerInfo.BenchmarkDecoratorInfo != nil {
  return "test"
 } else if info.LinkerInfo.ObjectLinkerInfo != nil {
  return "object"
 }
 panic(fmt.Sprintf("Unexpected module type for LinkerInfo"))
}

func getExtraLibs(info *CcInfo) []string {
 if info.LinkerInfo.TestBinaryInfo != nil {
  if info.LinkerInfo.TestBinaryInfo.Gtest {
   return []string{
    "libgtest",
    "libgtest_main",
   }
  }
 } else if info.LinkerInfo.BenchmarkDecoratorInfo != nil {
  return []string{"libgoogle-benchmark"}
 }
 return nil
}

func getIncludeDirs(ctx android.ModuleContext, m android.ModuleProxy, info *CcInfo) []string {
 moduleDir := ctx.OtherModuleDir(m) + string(filepath.Separator)
 if info.CompilerInfo.LibraryDecoratorInfo != nil {
  return sliceWithPrefix(moduleDir, info.CompilerInfo.LibraryDecoratorInfo.ExportIncludeDirs)
 }
 return nil
}

func cmakeSnapshotLoadHook(ctx android.LoadHookContext) {
 props := struct {
  Target struct {
   Darwin struct {
    Enabled *bool
   }
   Windows struct {
    Enabled *bool
   }
  }
 }{}
 props.Target.Darwin.Enabled = proptools.BoolPtr(false)
 props.Target.Windows.Enabled = proptools.BoolPtr(false)
 ctx.AppendProperties(&props)
}

// cmake_snapshot allows defining source packages for release outside of Android build tree.
// As a result of cmake_snapshot module build, a zip file is generated with CMake build definitions
// for selected source modules, their dependencies and optionally also the source code itself.
func CmakeSnapshotFactory() android.Module {
 module := &CmakeSnapshot{}
 module.AddProperties(&module.Properties)
 android.AddLoadHook(module, cmakeSnapshotLoadHook)
 android.InitAndroidArchModule(module, android.HostSupported, android.MultilibFirst)
 return module
}

func init() {
 android.InitRegistrationContext.RegisterModuleType("cc_cmake_snapshot", CmakeSnapshotFactory)
}