// Copyright 2021 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 provisioner

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"golang.org/x/sys/unix"

	"cos.googlesource.com/cos/tools.git/src/pkg/tools/partutil"
	"cos.googlesource.com/cos/tools.git/src/pkg/utils"
)

func switchRoot(deps Deps, runState *state) (err error) {
	if !runState.data.Config.BootDisk.ReclaimSDA3 {
		log.Println("ReclaimSDA3 is not set, not switching root device")
		return nil
	}
	bootDev, err := partutil.BootDev(deps.RootdevCmd)
	if err != nil {
		return err
	}
	sda3Device := bootDev + "3"
	sda5Device := bootDev + "5"
	rootDev, err := exec.Command(deps.RootdevCmd, "-s").Output()
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			return fmt.Errorf("error running rootdev: %v: stderr = %q", exitErr, string(exitErr.Stderr))
		}
		return fmt.Errorf("error running rootdev: %v", err)
	}
	rootDevStr := strings.TrimSpace(string(rootDev))
	if rootDevStr == sda5Device {
		log.Printf("Current root device is %q, not switching root device", rootDevStr)
		return nil
	}
	log.Println("Need to switch root device")
	log.Println("Copying partition 3 to partition 5...")
	in, err := os.Open(sda3Device)
	if err != nil {
		return err
	}
	defer utils.CheckClose(in, fmt.Sprintf("error closing %q", sda3Device), &err)
	out, err := os.Create(sda5Device)
	if err != nil {
		return err
	}
	defer utils.CheckClose(out, fmt.Sprintf("error closing %q", sda5Device), &err)
	if _, err := io.Copy(out, in); err != nil {
		return fmt.Errorf("error copying from %q to %q: %v", sda3Device, sda5Device, err)
	}
	log.Println("Setting GPT priority...")
	if err := utils.RunCommand([]string{deps.CgptCmd, "prioritize", "-P", "5", "-i", "4", bootDev}, "", nil); err != nil {
		return err
	}
	log.Println("Reboot required to switch root device")
	return ErrRebootRequired
}

func shrinkSDA3(deps Deps, runState *state) error {
	if !runState.data.Config.BootDisk.ReclaimSDA3 {
		log.Println("ReclaimSDA3 is not set, not shrinking partition 3")
		return nil
	}
	bootDev, err := partutil.BootDev(deps.RootdevCmd)
	if err != nil {
		return err
	}
	minimal, err := partutil.IsPartitionMinimal(bootDev, 3)
	if err != nil {
		return fmt.Errorf("error checking partition 3 size: %v", err)
	}
	if minimal {
		log.Println("partition 3 is minimally sized, not shrinking partition 3")
		return nil
	}
	log.Println("ReclaimSDA3 is set, and partition 3 is not minimal; now shrinking partition 3")
	if _, err := partutil.MinimizePartition(bootDev, 3); err != nil {
		return fmt.Errorf("error minimizing partition 3: %v", err)
	}
	log.Println("Reboot required to reload partition table changes")
	return ErrRebootRequired
}

func setupOnShutdownUnit(deps Deps, runState *state) (err error) {
	if err := mountFunc("", filepath.Join(deps.RootDir, "tmp"), "", unix.MS_REMOUNT|unix.MS_NOSUID|unix.MS_NODEV, ""); err != nil {
		return fmt.Errorf("error remounting /tmp as exec: %v", err)
	}
	out := filepath.Join(deps.RootDir, "tmp", "handle_disk_layout.bin")
	if err := utils.CopyFile(deps.HandleDiskLayoutBin, out); err != nil {
		return err
	}
	if err := os.Chmod(out, 755); err != nil {
		return err
	}
	bootDev, err := partutil.BootDev(deps.RootdevCmd)
	if err != nil {
		return err
	}
	data := fmt.Sprintf(`[Unit]
Description=Run after everything unmounted
DefaultDependencies=false
Conflicts=shutdown.target
Before=mnt-stateful_partition.mount usr-share-oem.mount
After=tmp.mount

[Service]
Type=oneshot
RemainAfterExit=true
ExecStart=/bin/true
ExecStop=/bin/bash -c '/tmp/handle_disk_layout.bin %s 1 8 "%s" "%t" 2>&1 | sed "s/^/BuildStatus: /"'
TimeoutStopSec=600
StandardOutput=tty
StandardError=tty
TTYPath=/dev/ttyS2
`, bootDev, runState.data.Config.BootDisk.OEMSize, runState.data.Config.BootDisk.ReclaimSDA3)
	if err := ioutil.WriteFile(filepath.Join(deps.RootDir, "etc/systemd/system/last-run.service"), []byte(data), 0664); err != nil {
		return err
	}
	systemd := systemdClient{systemctl: deps.SystemctlCmd}
	if err := systemd.start("last-run.service", []string{"--no-block"}); err != nil {
		return err
	}
	// journald needs to be stopped in order for the stateful partition to be
	// unmounted at shutdown. We need the stateful partition to be unmounted so
	// that disk repartitioning can occur.
	if err := systemd.stopJournald(deps.RootDir); err != nil {
		return err
	}
	return nil
}

func calcSDA3End(device string) (uint64, error) {
	sda3Start, err := partutil.ReadPartitionStart(device, 3)
	if err != nil {
		return 0, err
	}
	sda3Size, err := partutil.ReadPartitionSize(device, 3)
	if err != nil {
		return 0, err
	}
	sda3End := sda3Start + sda3Size - 1
	return sda3End, nil
}

func waitForDiskResize(deps Deps, runState *state) error {
	if !runState.data.Config.BootDisk.WaitForDiskResize {
		log.Println("WaitForDiskResize is not set, not waiting for a boot disk resize")
		return nil
	}
	if runState.data.DiskResizeComplete {
		log.Println("Already finished waiting for disk resize, not waiting again")
		return nil
	}
	bootDev, err := partutil.BootDev(deps.RootdevCmd)
	if err != nil {
		return err
	}
	startSize, err := ioutil.ReadFile(filepath.Join(deps.RootDir, fmt.Sprintf("sys/class/block/%s/size", filepath.Base(bootDev))))
	if err != nil {
		return err
	}
	log.Println("WaitForDiskResize is set; waiting for the boot disk size to change. Timeout is 3 minutes")
	start := time.Now()
	end := start.Add(3 * time.Minute)
	for time.Now().Before(end) {
		curSize, err := ioutil.ReadFile(filepath.Join(deps.RootDir, fmt.Sprintf("sys/class/block/%s/size", filepath.Base(bootDev))))
		if err != nil {
			return err
		}
		if string(curSize) != string(startSize) {
			log.Printf("Boot disk size has changed: start %q, end %q", strings.TrimSpace(string(startSize)), strings.TrimSpace(string(curSize)))
			runState.data.DiskResizeComplete = true
			return runState.write()
		}
		time.Sleep(time.Second)
	}
	return errors.New("timed out waiting for disk resize")
}

func relocatePartitions(deps Deps, runState *state) error {
	if !runState.data.Config.BootDisk.ReclaimSDA3 && runState.data.Config.BootDisk.OEMSize == "" {
		log.Println("ReclaimSDA3 is not set, OEM resize not requested, not relocating partitions")
		return nil
	}
	bootDev, err := partutil.BootDev(deps.RootdevCmd)
	if err != nil {
		return err
	}
	if runState.data.Config.BootDisk.OEMSize != "" {
		// Check if OEM partition is after sda3; if so, then we're done
		oemStart, err := partutil.ReadPartitionStart(bootDev, 8)
		if err != nil {
			return err
		}
		sda3End, err := calcSDA3End(bootDev)
		if err != nil {
			return err
		}
		if oemStart > sda3End {
			log.Println("OEM resize requested, OEM appears to be relocated after sda3. Partition relocation is complete")
			return nil
		}
	} else {
		// Check two things:
		// 1. sda3 is minimal
		// 2. Stateful partition is located immediately after sda3
		//
		// If both are true, we are done.
		minimal, err := partutil.IsPartitionMinimal(bootDev, 3)
		if err != nil {
			return err
		}
		statefulStart, err := partutil.ReadPartitionStart(bootDev, 1)
		if err != nil {
			return err
		}
		sda3Start, err := partutil.ReadPartitionStart(bootDev, 3)
		// The stateful partition is relocated 4096 sectors after the start of sda3.
		// See src/pkg/tools/handle_disk_layout.go for details.
		if minimal && statefulStart == sda3Start+4096 {
			log.Println("ReclaimSDA3 is set, sda3 appears to have been reclaimed. Partition relocation is complete")
			return nil
		}
	}
	// Partition relocation must be done. Prepare for disk relocation to happen on
	// the next reboot
	log.Println("Partition relocation is required. Preparing for partition relocation to occur on the next reboot")
	if err := setupOnShutdownUnit(deps, runState); err != nil {
		return err
	}
	log.Println("Reboot required to relocate partitions")
	return ErrRebootRequired
}

func resizeOEMFileSystem(deps Deps, runState *state) error {
	if !runState.data.Config.BootDisk.ReclaimSDA3 && runState.data.Config.BootDisk.OEMSize == "" {
		log.Println("ReclaimSDA3 is not set, OEM resize not requested, partition relocation did not occur, FS resize unnecessary")
		return nil
	}
	// Check if OEM partition is after sda3; if so, then relocation occurred and
	// we need to resize the file system.
	bootDev, err := partutil.BootDev(deps.RootdevCmd)
	if err != nil {
		return err
	}
	sda3End, err := calcSDA3End(bootDev)
	if err != nil {
		return err
	}
	oemStart, err := partutil.ReadPartitionStart(bootDev, 8)
	if err != nil {
		return err
	}
	if oemStart < sda3End {
		log.Println("OEM partition is before sda3; relocation did not occur, FS resize unnecessary")
		return nil
	}
	log.Println("Partition relocation appears to have occurred, resizing the OEM file system")
	systemd := systemdClient{systemctl: deps.SystemctlCmd}
	if err := systemd.stop("usr-share-oem.mount"); err != nil {
		return err
	}
	sda8Device := bootDev + "8"
	if err := utils.RunCommand([]string{deps.E2fsckCmd, "-fp", sda8Device}, "", nil); err != nil {
		return err
	}
	resizeArgs := []string{deps.Resize2fsCmd, sda8Device}
	if runState.data.Config.BootDisk.OEMFSSize4K != 0 {
		resizeArgs = append(resizeArgs, strconv.FormatUint(runState.data.Config.BootDisk.OEMFSSize4K, 10))
	}
	if err := utils.RunCommand(resizeArgs, "", nil); err != nil {
		return err
	}
	if err := systemd.start("usr-share-oem.mount", nil); err != nil {
		return err
	}
	log.Println("OEM file system resized to account for available space")
	return nil
}

// repartitionBootDisk executes all behaviors related to repartitioning the boot
// disk. Most of these behaviors require a reboot. To keep reboots simple (e.g.
// we don't want to initiate a reboot when deferred statements are unresolved),
// we handle reboots by returning ErrRebootRequired and asking the caller to
// initiate the reboot.
func repartitionBootDisk(deps Deps, runState *state) error {
	if err := switchRoot(deps, runState); err != nil {
		return err
	}
	if err := shrinkSDA3(deps, runState); err != nil {
		return err
	}
	if err := waitForDiskResize(deps, runState); err != nil {
		return err
	}
	if err := relocatePartitions(deps, runState); err != nil {
		return err
	}
	if err := resizeOEMFileSystem(deps, runState); err != nil {
		return err
	}
	return nil
}