// Copyright 2020 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 rust

import (
 "encoding/json"
 "fmt"
 "strings"

 "github.com/google/blueprint/proptools"

 "android/soong/android"
 "android/soong/rust/config"
)

// This singleton collects Rust crate definitions and generates a JSON file
// (${OUT_DIR}/soong/rust-project.json) which can be use by external tools,
// such as rust-analyzer. It does so when either make, mm, mma, mmm or mmma is
// called.  This singleton is enabled only if SOONG_GEN_RUST_PROJECT is set.
// For example,
//
//   $ SOONG_GEN_RUST_PROJECT=1 m nothing

const (
 // Environment variables used to control the behavior of this singleton.
 envVariableCollectRustDeps = "SOONG_GEN_RUST_PROJECT"
 envVariableUseKythe        = "XREF_CORPUS"
 rustProjectJsonFileName    = "rust-project.json"
 rustTargetMappingFileName  = "rust-target-mapping.json"
)

// The format of rust-project.json is not yet finalized. A current description is available at:
// https://github.com/rust-analyzer/rust-analyzer/blob/master/docs/user/manual.adoc#non-cargo-based-projects
type rustProjectDep struct {
 // The Crate attribute is the index of the dependency in the Crates array in rustProjectJson.
 Crate int    `json:"crate"`
 Name  string `json:"name"`
}

type rustProjectCrate struct {
 DisplayName    string                 `json:"display_name"`
 RootModule     string                 `json:"root_module"`
 Edition        string                 `json:"edition,omitempty"`
 Deps           []rustProjectDep       `json:"deps"`
 Cfg            []string               `json:"cfg"`
 Target         string                 `json:"target"`
 Env            map[string]string      `json:"env"`
 ProcMacro      bool                   `json:"is_proc_macro"`
 ProcMacroDylib *string                `json:"proc_macro_dylib_path"`
 Source         rustProjectIncludeDirs `json:"source"`
}

type rustProjectJson struct {
 Sysroot string             `json:"sysroot"`
 Crates  []rustProjectCrate `json:"crates"`
}

type rustTargetMappingsJson []rustTargetMappingJson

type rustTargetMappingJson struct {
 Name               string `json:"name"`
 BuildTarget        string `json:"build_target"`
 CheckTarget        string `json:"check_target"`
 SourceDir          string `json:"source_dir"`
 TargetProduct      string `json:"TARGET_PRODUCT"`
 TargetBuildVariant string `json:"TARGET_BUILD_VARIANT"`
}

type rustProjectIncludeDirs struct {
 Include_dirs []string `json:"include_dirs"`
 Exclude_dirs []string `json:"exclude_dirs"`
}

// crateInfo is used during the processing to keep track of the known crates.
type crateInfo struct {
 Idx    int            // Index of the crate in rustProjectJson.Crates slice.
 Deps   map[string]int // The keys are the module names and not the crate names.
 Device bool           // True if the crate at idx was a device crate
}

type projectGeneratorSingleton struct {
 project        rustProjectJson
 targetMappings rustTargetMappingsJson
 knownCrates    map[string]crateInfo // Keys are module names.
}

func rustProjectGeneratorSingleton() android.Singleton {
 return &projectGeneratorSingleton{}
}

func init() {
 android.RegisterParallelSingletonType("rust_project_generator", rustProjectGeneratorSingleton)
}

// mergeDependencies visits all the dependencies for module and updates crate and deps
// with any new dependency.
func (singleton *projectGeneratorSingleton) mergeDependencies(ctx android.SingletonContext,
 module android.ModuleProxy, crate *rustProjectCrate, deps map[string]int) {

 ctx.VisitDirectDepsProxies(module, func(child android.ModuleProxy) {
  // Skip intra-module dependencies (i.e., generated-source library depending on the source variant).
  if module.Name() == child.Name() {
   return
  }
  // Skip unsupported modules.
  rustInfo, commonInfo, ok := isModuleSupported(ctx, child)
  if !ok {
   return
  }
  // For unknown dependency, add it first.
  var childId int
  cInfo, known := singleton.knownCrates[child.Name()]
  if !known {
   childId, ok = singleton.addCrate(ctx, child, rustInfo, commonInfo)
   if !ok {
    return
   }
  } else {
   childId = cInfo.Idx
  }
  // Is this dependency known already?
  if _, ok = deps[child.Name()]; ok {
   return
  }
  crate.Deps = append(crate.Deps, rustProjectDep{Crate: childId, Name: rustInfo.CompilerInfo.CrateName})
  deps[child.Name()] = childId
 })
}

// isModuleSupported returns the RustModule if the module
// should be considered for inclusion in rust-project.json.
func isModuleSupported(ctx android.SingletonContext, module android.ModuleProxy) (*RustInfo, *android.CommonModuleInfo, bool) {
 info, ok := android.OtherModuleProvider(ctx, module, RustInfoProvider)
 if !ok {
  return nil, nil, false
 }
 commonInfo := android.OtherModuleProviderOrDefault(ctx, module, android.CommonModuleInfoProvider)
 if !commonInfo.Enabled {
  return nil, nil, false
 }
 return info, commonInfo, true
}

// addCrate adds a crate to singleton.project.Crates ensuring that required
// dependencies are also added. It returns the index of the new crate in
// singleton.project.Crates
func (singleton *projectGeneratorSingleton) addCrate(ctx android.SingletonContext, module android.ModuleProxy, rustInfo *RustInfo,
 commonInfo *android.CommonModuleInfo) (int, bool) {

 deps := make(map[string]int)

 var procMacroDylib *string = nil

 if procMacro := rustInfo.ProcMacroInfo; procMacro != nil {
  procMacroDylib = proptools.StringPtr(procMacro.Dylib.String())
 }

 var toolchain = config.FindToolchain(commonInfo.Target.Os, commonInfo.Target.Arch)

 var rootmodule = rustInfo.CompilerInfo.CrateRootPath.String()
 include_dir := []string{ctx.ModuleDir(module)}
 // Aidl generates rust files so we include those files in the include_dir parameter.
 if strings.Contains(rootmodule, "aidl") {
  include_dir = append(include_dir, rootmodule)
 }
 crate := rustProjectCrate{
  DisplayName:    module.Name(),
  RootModule:     rootmodule,
  Edition:        rustInfo.CompilerInfo.Edition,
  Deps:           make([]rustProjectDep, 0),
  Cfg:            make([]string, 0),
  Target:         toolchain.RustTriple(),
  Env:            make(map[string]string),
  ProcMacro:      procMacroDylib != nil,
  ProcMacroDylib: procMacroDylib,
  Source: rustProjectIncludeDirs{
   Include_dirs: include_dir,
   Exclude_dirs: []string{},
  }, // TODO: What should this value be?
 }

 if cargoOutDir := rustInfo.CompilerInfo.CargoOutDir; cargoOutDir.Valid() {
  crate.Env["OUT_DIR"] = cargoOutDir.String()
 }
 for _, feature := range rustInfo.CompilerInfo.Features {
  crate.Cfg = append(crate.Cfg, "feature=\""+feature+"\"")
 }

 singleton.mergeDependencies(ctx, module, &crate, deps)

 var idx int
 if cInfo, ok := singleton.knownCrates[module.Name()]; ok {
  idx = cInfo.Idx
  singleton.project.Crates[idx] = crate
 } else {
  idx = len(singleton.project.Crates)
  singleton.project.Crates = append(singleton.project.Crates, crate)
 }
 singleton.knownCrates[module.Name()] = crateInfo{Idx: idx, Deps: deps, Device: commonInfo.Target.Os.Class == android.Device}
 if rustInfo.CompilerInfo.BuildTarget != nil {
  buildVariant := "user"
  if ctx.Config().Eng() {
   buildVariant = "eng"
  } else if ctx.Config().Debuggable() {
   buildVariant = "userdebug"
  }
  mapping := rustTargetMappingJson{
   Name:               module.Name(),
   BuildTarget:        rustInfo.CompilerInfo.BuildTarget.String(),
   CheckTarget:        rustInfo.CompilerInfo.CheckTarget.String(),
   SourceDir:          ctx.ModuleDir(module),
   TargetProduct:      ctx.Config().DeviceProduct(),
   TargetBuildVariant: buildVariant,
  }
  singleton.targetMappings = append(singleton.targetMappings, mapping)
 }
 return idx, true
}

// appendCrateAndDependencies creates a rustProjectCrate for the module argument and appends it to singleton.project.
// It visits the dependencies of the module depth-first so the dependency ID can be added to the current module. If the
// current module is already in singleton.knownCrates, its dependencies are merged.
func (singleton *projectGeneratorSingleton) appendCrateAndDependencies(ctx android.SingletonContext, module android.ModuleProxy) {
 rustInfo, commonInfo, ok := isModuleSupported(ctx, module)
 if !ok {
  return
 }
 // If we have seen this crate already; merge any new dependencies.
 if cInfo, ok := singleton.knownCrates[module.Name()]; ok {
  // If we have a new device variant, override the old one
  if !cInfo.Device && commonInfo.Target.Os.Class == android.Device {
   singleton.addCrate(ctx, module, rustInfo, commonInfo)
   return
  }
  crate := singleton.project.Crates[cInfo.Idx]
  singleton.mergeDependencies(ctx, module, &crate, cInfo.Deps)
  singleton.project.Crates[cInfo.Idx] = crate
  return
 }
 singleton.addCrate(ctx, module, rustInfo, commonInfo)
}

func (singleton *projectGeneratorSingleton) GenerateBuildActions(ctx android.SingletonContext) {
 if !(ctx.Config().IsEnvTrue(envVariableCollectRustDeps) || ctx.Config().Getenv(envVariableUseKythe) != "") {
  return
 }

 singleton.project.Sysroot = config.RustPath(ctx)

 singleton.knownCrates = make(map[string]crateInfo)
 ctx.VisitAllModuleProxies(func(module android.ModuleProxy) {
  singleton.appendCrateAndDependencies(ctx, module)
 })

 path := android.PathForOutput(ctx, rustProjectJsonFileName)
 err := createJsonFile(singleton.project, path)
 if err != nil {
  ctx.Errorf(err.Error())

 }
 if ctx.Config().XrefCorpusName() != "" {
  rule := android.NewRuleBuilder(pctx, ctx).SandboxDisabled()
  jsonPath := android.PathForOutput(ctx, "rust-project.json")
  kzipPath := android.PathForOutput(ctx, "rust-project.kzip")
  vnames := android.PathForSource(ctx, "build/soong/vnames.json")
  rule.Command().PrebuiltBuildTool(ctx, "rust_project_to_kzip").
   FlagWithInput("-project_json ", jsonPath).
   FlagWithOutput("-output ", kzipPath).
   FlagWithArg("-corpus ", ctx.Config().XrefCorpusName()).
   FlagWithArg("-root ", ".").
   FlagWithInput("-vnames_json_path ", vnames)
  ruleName := "rust3-extraction"
  rule.Build(ruleName, "Turning Rust project into kzips")
  ctx.Phony("xref_rust", kzipPath)
 }

 rustTargetMappingPath := android.PathForOutput(ctx, rustTargetMappingFileName)
 if err := createJsonFile(singleton.targetMappings, rustTargetMappingPath); err != nil {
  ctx.Errorf(err.Error())
 }

}

func createJsonFile(project any, rustProjectPath android.WritablePath) error {
 buf, err := json.MarshalIndent(project, "", "  ")
 if err != nil {
  return fmt.Errorf("JSON marshal failed: %s", err)
 }

 if err := android.WriteFileToOutputDir(rustProjectPath, buf, 0666); err != nil {
  return fmt.Errorf("Writing rust-project to %s failed: %s", rustProjectPath.String(), err)
 }
 return nil
}