// Copyright 2026 Google LLC
//
// 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 main

import (
	"flag"
	"fmt"
	"log"
	"strings"

	"cos.googlesource.com/cos/tools.git/src/pkg/oempreloader"
)

// stringList implements a list flag variable. Each list item needs to be
// specified in its own instance on the command line.
type stringList []string

func (s *stringList) String() string {
	return strings.Join(*s, ", ")
}

func (s *stringList) Set(value string) error {
	*s = append(*s, value)
	return nil
}

func main() {
	srcImage := flag.String("src-image", "", "Path to the source COS image")
	outImage := flag.String("out-image", "", "Path to the output image destination")
	oemFSSize := flag.String("oem-fs-size", "", "Size of the OEM filesystem (e.g. 2G, 500M)")
	diskSizeGB := flag.Int("disk-size-gb", 0, "Total disk size in GB")
	installDir := flag.String("install-dir", "", "Path to the directory containing files to install")
	installExt4 := flag.String("install-ext4", "", "Path to a pre-built ext4 image")
	vmlinux := flag.String("vmlinux", "", "Path to vmlinux with symbol information for this image. Optional for released images.")
	var cmdlineSedCommands stringList
	flag.Var(&cmdlineSedCommands, "cmdline-sed-command", "Sed command to run on the kernel cmdline. Can be provided multiple times to run multiple commands.")

	flag.Parse()

	// Basic Validation
	if *srcImage == "" || *outImage == "" {
		log.Fatalf("--src-image and --out-image must be set.")
	}

	if (*installDir != "") == (*installExt4 != "") {
		log.Fatalf("Exactly one of --install-dir or --install-ext4 must be set")
	}

	fmt.Printf("Starting oem-preloader...\n")
	fmt.Printf("Source: %s\nTarget: %s\nOEM FS Size: %s\nDisk Size: %dGB\n",
		*srcImage, *outImage, *oemFSSize, *diskSizeGB)

	if err := oempreloader.Extend(*srcImage, *outImage, *oemFSSize, *diskSizeGB); err != nil {
		log.Fatalf("Failed to extend image disk: %v", err)
	}
	oemSize, err := oempreloader.ReadOEMPartitionSize(*outImage)
	if err != nil {
		log.Fatalf("Failed to read OEM partition size: %v", err)
	}

	oemFile := ""
	if *installDir != "" {
		oem, err := oempreloader.PreloadDir(*installDir, *oemFSSize, oemSize)
		if err != nil {
			log.Fatalf("Failed to preload dir: %v", err)
		}
		oemFile = oem
	} else {
		if err := oempreloader.ExtendExt4(*installExt4, oemSize); err != nil {
			log.Fatalf("Failed to extend ext4 file: %v", err)
		}
		oemFile = *installExt4
	}
	hash, salt, err := oempreloader.RunVeritysetup(oemFile, *oemFSSize)
	if err != nil {
		log.Fatalf("Failed to run veritysetup on OEM file: %v", err)
	}
	if err := oempreloader.WriteOEMFileToImage(oemFile, *outImage); err != nil {
		log.Fatalf("Failed to write OEM partition back to image: %v", err)
	}
	if err := oempreloader.UpdateKernelCmdLine(*outImage, hash, salt, *oemFSSize, *vmlinux, cmdlineSedCommands); err != nil {
		log.Fatalf("Failed to update kernel command line: %v", err)
	}
}