// Copyright (C) 2021 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 filesystem

import (
 "fmt"
 "strconv"

 "github.com/google/blueprint/proptools"

 "android/soong/android"
)

func init() {
 android.RegisterModuleType("logical_partition", logicalPartitionFactory)
}

type logicalPartition struct {
 android.ModuleBase

 properties logicalPartitionProperties

 output     android.Path
 installDir android.InstallPath
}

type logicalPartitionProperties struct {
 // Set the name of the output. Defaults to <module_name>.img.
 Stem *string

 // Total size of the logical partition. If set to "auto", total size is automatically
 // calcaulted as minimum.
 Size *string

 // List of partitions for default group. Default group has no size limit and automatically
 // minimized when creating an image.
 Default_group []partitionProperties

 // List of groups. A group defines a fixed sized region. It can host one or more logical
 // partitions and their total size is limited by the size of the group they are in.
 Groups []groupProperties

 // Whether the output is a sparse image or not. Default is false.
 Sparse *bool
}

type groupProperties struct {
 // Name of the partition group. Can't be "default"; use default_group instead.
 Name *string

 // Size of the partition group
 Size *string

 // List of logical partitions in this group
 Partitions []partitionProperties
}

type partitionProperties struct {
 // Name of the partition
 Name *string

 // Filesystem that is placed on the partition
 Filesystem *string `android:"path"`
}

// logical_partition is a partition image which has one or more logical partitions in it.
func logicalPartitionFactory() android.Module {
 module := &logicalPartition{}
 module.AddProperties(&module.properties)
 android.InitAndroidArchModule(module, android.DeviceSupported, android.MultilibFirst)
 return module
}

func (l *logicalPartition) DepsMutator(ctx android.BottomUpMutatorContext) {
 // do nothing
}

func (l *logicalPartition) installFileName() string {
 return proptools.StringDefault(l.properties.Stem, l.BaseModuleName()+".img")
}

func (l *logicalPartition) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 builder := android.NewRuleBuilder(pctx, ctx).SandboxDisabled()

 // Sparse the filesystem images and calculate their sizes
 sparseImages := make(map[string]android.Path)
 sparseImageSizes := make(map[string]android.Path)

 sparsePartitions := func(partitions []partitionProperties) {
  for _, part := range partitions {
   if part.Filesystem == nil {
    continue
   }
   sparseImg, sizeTxt := sparseFilesystem(ctx, part, builder)
   pName := proptools.String(part.Name)
   sparseImages[pName] = sparseImg
   sparseImageSizes[pName] = sizeTxt
  }
 }

 for _, group := range l.properties.Groups {
  sparsePartitions(group.Partitions)
 }

 sparsePartitions(l.properties.Default_group)

 cmd := builder.Command().BuiltTool("lpmake")

 size := proptools.String(l.properties.Size)
 if size == "" {
  ctx.PropertyErrorf("size", "must be set")
 } else if _, err := strconv.Atoi(size); err != nil && size != "auto" {
  ctx.PropertyErrorf("size", `must be a number or "auto"`)
 }
 cmd.FlagWithArg("--device-size=", size)

 // TODO(jiyong): consider supporting A/B devices. Then we need to adjust num of slots.
 cmd.FlagWithArg("--metadata-slots=", "2")
 cmd.FlagWithArg("--metadata-size=", "65536")

 if proptools.Bool(l.properties.Sparse) {
  cmd.Flag("--sparse")
 }

 groupNames := make(map[string]bool)
 partitionNames := make(map[string]bool)

 addPartitionsToGroup := func(partitions []partitionProperties, gName string) {
  for _, part := range partitions {
   pName := proptools.String(part.Name)
   if pName == "" {
    ctx.PropertyErrorf("groups.partitions.name", "must be set")
   }
   if _, ok := partitionNames[pName]; ok {
    ctx.PropertyErrorf("groups.partitions.name", "already exists")
   } else {
    partitionNames[pName] = true
   }
   // Get size of the partition by reading the -size.txt file
   var pSize string
   if size, hasSize := sparseImageSizes[pName]; hasSize {
    pSize = fmt.Sprintf("$(cat %s)", size)
   } else {
    pSize = "0"
   }
   cmd.FlagWithArg("--partition=", fmt.Sprintf("%s:readonly:%s:%s", pName, pSize, gName))
   if image, hasImage := sparseImages[pName]; hasImage {
    cmd.FlagWithInput("--image="+pName+"=", image)
   }
  }
 }

 addPartitionsToGroup(l.properties.Default_group, "default")

 for _, group := range l.properties.Groups {
  gName := proptools.String(group.Name)
  if gName == "" {
   ctx.PropertyErrorf("groups.name", "must be set")
  } else if gName == "default" {
   ctx.PropertyErrorf("groups.name", `can't use "default" as a group name. Use default_group instead`)
  }
  if _, ok := groupNames[gName]; ok {
   ctx.PropertyErrorf("group.name", "already exists")
  } else {
   groupNames[gName] = true
  }
  gSize := proptools.String(group.Size)
  if gSize == "" {
   ctx.PropertyErrorf("groups.size", "must be set")
  }
  if _, err := strconv.Atoi(gSize); err != nil {
   ctx.PropertyErrorf("groups.size", "must be a number")
  }
  cmd.FlagWithArg("--group=", gName+":"+gSize)

  addPartitionsToGroup(group.Partitions, gName)
 }

 output := android.PathForModuleOut(ctx, l.installFileName())
 cmd.FlagWithOutput("--output=", output)

 builder.Build("build_logical_partition", fmt.Sprintf("Creating %s", l.BaseModuleName()))

 l.installDir = android.PathForModuleInstall(ctx, "etc")
 ctx.InstallFile(l.installDir, l.installFileName(), output)

 ctx.SetOutputFiles([]android.Path{output}, "")
 l.output = output

 setCommonFilesystemInfo(ctx, l)
}

// Add a rule that converts the filesystem for the given partition to the given rule builder. The
// path to the sparse file and the text file having the size of the partition are returned.
func sparseFilesystem(ctx android.ModuleContext, p partitionProperties, builder *android.RuleBuilder) (android.Path, android.Path) {
 img := android.PathForModuleSrc(ctx, *p.Filesystem)
 name := proptools.String(p.Name)

 sparseImg := android.PathForModuleOut(ctx, name+".img")
 builder.Temporary(sparseImg)
 builder.Command().BuiltTool("img2simg").Input(img).Output(sparseImg)

 sizeTxt := android.PathForModuleOut(ctx, name+"-size.txt")
 builder.Temporary(sizeTxt)
 builder.Command().BuiltTool("sparse_img").Flag("--get_partition_size").Input(sparseImg).
  Text("| ").Text("tr").FlagWithArg("-d ", "'\n'").
  Text("> ").Output(sizeTxt)

 return sparseImg, sizeTxt
}

var _ android.AndroidMkEntriesProvider = (*logicalPartition)(nil)

// Implements android.AndroidMkEntriesProvider
func (l *logicalPartition) AndroidMkEntries() []android.AndroidMkEntries {
 return []android.AndroidMkEntries{android.AndroidMkEntries{
  Class:      "ETC",
  OutputFile: android.OptionalPathForPath(l.output),
  ExtraEntries: []android.AndroidMkExtraEntriesFunc{
   func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
    entries.SetString("LOCAL_MODULE_PATH", l.installDir.String())
    entries.SetString("LOCAL_INSTALLED_MODULE_STEM", l.installFileName())
   },
  },
 }}
}

var _ Filesystem = (*logicalPartition)(nil)

func (l *logicalPartition) OutputPath() android.Path {
 return l.output
}

func (l *logicalPartition) SignedOutputPath() android.Path {
 return nil // logical partition is not signed by itself
}