// 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 filesystem

import (
 "github.com/google/blueprint"
 "github.com/google/blueprint/proptools"

 "android/soong/android"
)

//go:generate go run ../../blueprint/gobtools/codegen

// @auto-generate: gob
type bootloaderInfo struct {
 bootloaderImg android.Path
}

var bootloaderInfoProvider = blueprint.NewProvider[bootloaderInfo]()

type prebuiltBootloader struct {
 android.ModuleBase
 properties       PrebuiltBootloaderProperties
 vbmetaPartitions vbmetaPartitionInfos
}

type PrebuiltBootloaderProperties struct {
 Src *string `android:"path"`
 // List of OTA updatable partitions of bootloader.img.
 // These will be unpacked from bootloader.img and added to the list
 // of partitions to be updated.
 Ab_ota_partitions []string

 // Tool for unpacking bootloader.img
 Unpack_tool *string `android:"path"`

 // Additional files needed to run unpack_tool
 Unpack_tool_deps []string `android:"path"`
}

// TODO(soong-team): This module should be registered with the name
// `prebuilt_bootloader`. If not, update the property error in
// [androidDevice.checkRadioVersion]. Remove this TODO once the module type is
// registered.
func PrebuiltBootloaderFactory() android.Module {
 module := &prebuiltBootloader{}
 android.InitAndroidArchModule(module, android.DeviceSupported, android.MultilibFirst)
 module.AddProperties(&module.properties)
 return module
}

func (p *prebuiltBootloader) DepsMutator(ctx android.BottomUpMutatorContext) {
 ctx.AddHostToolDependencies("avbtool", "avb_openssl")
}

func (p *prebuiltBootloader) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 if p.properties.Src == nil {
  ctx.PropertyErrorf("src", "Source cannot be empty")
 }
 bootloader := android.PathForModuleSrc(ctx, proptools.String(p.properties.Src))
 ctx.SetOutputFiles([]android.Path{bootloader}, "")
 bootloaderPartitionFiles := p.partitionFilesBootloader(ctx)
 ctx.SetOutputFiles(bootloaderPartitionFiles, "bootloader_partitions")

 android.SetProvider(ctx, vbmetaPartitionsProvider, p.vbmetaPartitions)
 android.SetProvider(ctx, bootloaderInfoProvider, bootloaderInfo{
  bootloaderImg: bootloader,
 })
}

// Unpack a partition from a bootloader.img and add them to
// the list of partitions to be updated.
func (p *prebuiltBootloader) partitionFilesBootloader(ctx android.ModuleContext) android.Paths {
 bootloader := android.PathForModuleSrc(ctx, proptools.String(p.properties.Src))
 if len(p.properties.Ab_ota_partitions) == 0 {
  return nil
 }
 var files android.Paths
 unpackedDir := android.PathForModuleOut(ctx, "unpack_bootloader")
 builder := android.NewRuleBuilder(pctx, ctx).Sbox(
  unpackedDir,
  android.PathForModuleOut(ctx, "unpack_bootloader.textproto"),
 )
 for _, partition := range p.properties.Ab_ota_partitions {
  unpackedImg := unpackedDir.Join(ctx, partition+".img")
  cmd := builder.Command()
  cmd.Input(android.PathForModuleSrc(ctx, proptools.String(p.properties.Unpack_tool))).
   Implicits(android.PathsForModuleSrc(ctx, p.properties.Unpack_tool_deps)).
   Flag(" unpack ").
   Flag("-o ").
   Textf("%s ", cmd.PathForOutput(unpackedDir)).
   Input(bootloader).
   Text(" " + partition).
   ImplicitOutput(unpackedImg)

  files = append(files, unpackedImg)
  // avb signed
  signed := p.avbAddHash(ctx, builder, partition, unpackedImg)
  files = append(files, signed)
  p.vbmetaPartitions = append(p.vbmetaPartitions, vbmetaPartitionInfo{
   Name:                     partition,
   Output:                   signed,
   AbOtaBootloaderPartition: true,
  })
 }
 builder.Build("unpack_bootloader", "unpack_bootloader")
 return files
}

func (p *prebuiltBootloader) avbAddHash(ctx android.ModuleContext, builder *android.RuleBuilder, partitionName string, unpackedPartition android.OutputPath) android.Path {
 output := unpackedPartition.InSameDir(ctx, partitionName+"_vbfooted.img")
 builder.Command().BuiltTool("cp").Input(unpackedPartition).Output(output)

 cmd := builder.Command()
 cmd.BuiltTool("avbtool").
  ImplicitTool(ctx.Config().HostToolPath(ctx, "avb_openssl")).
  Text("add_hash_footer").
  FlagWithOutput("--image ", output).
  FlagWithArg("--partition_name ", partitionName).
  Textf(`--salt $(sha256sum "%s" "%s" | cut -d " " -f 1 | tr -d '\n')`, cmd.PathForInput(ctx.Config().BuildNumberFile(ctx)), cmd.PathForInput(ctx.Config().BuildDateFile(ctx))).
  OrderOnly(ctx.Config().BuildNumberFile(ctx)).OrderOnly(ctx.Config().BuildDateFile(ctx))

 vbmetaPaddingSize := 64*1024 + 4096
 cmd.Textf("--partition_size $(( %d + (( $(stat -c %%s %s) - 1) / 4096 + 1) * 4096 ))", vbmetaPaddingSize, cmd.PathForOutput(unpackedPartition))

 return output
}