// Copyright (c) Tailscale Inc & contributors
// SPDX-License-Identifier: BSD-3-Clause

// This program builds the Tailscale Appliance Gokrazy image.
//
// As of 2024-06-02 this is a exploratory work in progress and is
// not intended for serious use.
//
// Tracking issue is https://github.com/tailscale/tailscale/issues/1866
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"tailscale.com/gokrazy/mkfs"
)

var (
	app    = flag.String("app", "tsapp", "appliance name; one of the subdirectories of gokrazy/")
	bucket = flag.String("bucket", "tskrazy-import", "S3 bucket to upload disk image to while making AMI")
	build  = flag.Bool("build", false, "if true, just build locally and stop, without uploading")
	gaf    = flag.Bool("gaf", false, "if true, build a gokrazy archive format file instead of a full disk image")
)

// baseImageSizeBytes is the size of the disk image we ask monogok to
// produce (and that the AWS AMI import expects). It has to be large
// enough to fit gokrazy's standard partition layout (see
// github.com/bradfitz/monogok/disklayout):
//
//	  4 MiB gap before the first partition
//	100 MiB boot      (FAT)
//	500 MiB root A    (squashfs; the partition OTA updates write into)
//	500 MiB root B    (squashfs)
//	 ~96 MiB /perm    (ext4; rest of the disk minus the secondary GPT)
//
// Bump this to give /perm more room (and to make the produced .img
// file larger). The same value is passed to monogok via
// --target_storage_bytes and to mkfs.Perm so the GPT and the ext4
// inside it agree on the disk's size.
//
// imageSizeBytesFor may round this up; callers should use that helper
// instead of this constant.
const baseImageSizeBytes = 1258299392

// imageSizeBytesFor returns the disk image size to use for app. For Raspberry
// Pi appliances the size is rounded up to the next power of two because
// qemu-system-aarch64's raspi3b machine rejects SD card images whose size
// isn't a power of two.
func imageSizeBytesFor(app string) int64 {
	if !strings.HasPrefix(app, "tsapp-pi.") {
		return baseImageSizeBytes
	}
	n := int64(1)
	for n < baseImageSizeBytes {
		n <<= 1
	}
	return n
}

var conf gokrazyConfig

// gokrazyConfig is the subset of gokrazy/internal/config.Struct
// that we care about.
type gokrazyConfig struct {
	// Environment is os.Environment pairs to use when
	// building userspace.
	// See https://gokrazy.org/userguide/instance-config/#environment
	Environment []string
}

func (c *gokrazyConfig) GOARCH() string {
	for _, e := range c.Environment {
		if v, ok := strings.CutPrefix(e, "GOARCH="); ok {
			return v
		}
	}
	return ""
}

func main() {
	flag.Parse()

	if *app == "" || strings.Contains(*app, "/") {
		log.Fatalf("--app must be non-empty name such as 'tsapp' or 'natlabapp'")
	}

	confJSON, err := os.ReadFile(filepath.Join(*app, "config.json"))
	if err != nil {
		log.Fatalf("reading config.json: %v", err)
	}
	if err := json.Unmarshal(confJSON, &conf); err != nil {
		log.Fatalf("unmarshaling config.json: %v", err)
	}
	switch conf.GOARCH() {
	case "amd64", "arm64":
	default:
		log.Fatalf("config.json GOARCH %q must be amd64 or arm64", conf.GOARCH())
	}

	if err := buildImage(); err != nil {
		log.Fatalf("build image: %v", err)
	}
	if *build || *gaf {
		log.Printf("built. stopping.")
		return
	}

	if err := copyToS3(); err != nil {
		log.Fatalf("copy to S3: %v", err)
	}

	importTask, err := startImportSnapshot()
	if err != nil {
		log.Fatalf("start import snapshot: %v", err)
	}
	snapID, err := waitForImportSnapshot(importTask)
	if err != nil {
		log.Fatalf("waitForImportSnapshot(%v): %v", importTask, err)
	}
	log.Printf("snap ID: %v", snapID)

	ami, err := makeAMI(fmt.Sprintf(*app+"-%d", time.Now().Unix()), snapID)
	if err != nil {
		log.Fatalf("makeAMI: %v", err)
	}
	log.Printf("made AMI: %v", ami)
}

func buildImage() error {
	dir, err := os.Getwd()
	if err != nil {
		return err
	}
	if fi, err := os.Stat(filepath.Join(dir, *app)); err != nil || !fi.IsDir() {
		return fmt.Errorf("in wrong directory %v; no %q subdirectory found", dir, *app)
	}

	args := []string{"run", "github.com/bradfitz/monogok/cmd/monogok"}
	if *gaf {
		args = append(args,
			"overwrite",
			"--gaf", filepath.Join(dir, *app+".gaf"),
		)
	} else {
		args = append(args,
			"overwrite",
			"--full", filepath.Join(dir, *app+".img"),
			fmt.Sprintf("--target_storage_bytes=%d", imageSizeBytesFor(*app)),
		)
	}

	cmd := exec.Command("go", args...)
	cmd.Dir = filepath.Join(dir, *app)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return err
	}
	if *gaf {
		return nil
	}

	imgPath := filepath.Join(dir, *app+".img")
	f, err := os.OpenFile(imgPath, os.O_RDWR, 0)
	if err != nil {
		return fmt.Errorf("open %s: %w", imgPath, err)
	}
	defer f.Close()
	if err := mkfs.Perm(f, imageSizeBytesFor(*app)); err != nil {
		return fmt.Errorf("formatting /perm in %s: %v", imgPath, err)
	}
	log.Printf("Wrote ext4 /perm filesystem to %s.", imgPath)
	return nil
}

func copyToS3() error {
	cmd := exec.Command("aws", "s3", "cp", *app+".img", "s3://"+*bucket+"/")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func startImportSnapshot() (importTaskID string, err error) {
	out, err := exec.Command("aws", "ec2", "import-snapshot", "--disk-container", "Url=s3://"+*bucket+"/"+*app+".img").CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("import snapshot: %v: %s", err, out)
	}
	var resp struct {
		ImportTaskID string `json:"ImportTaskId"`
	}
	/*
		{
			"ImportTaskId": "import-snap-0d2d72622b4359567",
			"SnapshotTaskDetail": {
				"DiskImageSize": 0.0,
				"Progress": "0",
				"Status": "active",
				"StatusMessage": "pending",
				"Url": "s3://tskrazy-import/tskrazy.img"
			},
			"Tags": []
		}
	*/
	if err := json.Unmarshal(out, &resp); err != nil {
		return "", fmt.Errorf("unmarshal response: %v: %s", err, out)
	}
	return resp.ImportTaskID, nil
}

/*
% aws ec2 describe-import-snapshot-tasks --import-task-ids import-snap-0d2d72622b4359567
{
    "ImportSnapshotTasks": [
        {
            "ImportTaskId": "import-snap-0d2d72622b4359567",
            "SnapshotTaskDetail": {
                "DiskImageSize": 1258299392.0,
                "Format": "RAW",
                "SnapshotId": "snap-053efd3539d787927",
                "Status": "completed",
                "Url": "s3://tskrazy-import/tskrazy.img",
                "UserBucket": {
                    "S3Bucket": "tskrazy-import",
                    "S3Key": "tskrazy.img"
                }
            },
            "Tags": []
        }
    ]
}
*/

func waitForImportSnapshot(importTaskID string) (snapID string, err error) {
	for {
		out, err := exec.Command("aws", "ec2", "describe-import-snapshot-tasks", "--import-task-ids", importTaskID).CombinedOutput()
		if err != nil {
			return "", fmt.Errorf("describe import snapshot tasks: %v: %s", err, out)
		}

		var resp struct {
			ImportSnapshotTasks []struct {
				SnapshotTaskDetail struct {
					SnapshotID string `json:"SnapshotId"`
					Status     string `json:"Status"`
				} `json:"SnapshotTaskDetail"`
			} `json:"ImportSnapshotTasks"`
		}
		if err := json.Unmarshal(out, &resp); err != nil {
			return "", fmt.Errorf("unmarshal response: %v: %s", err, out)
		}
		if len(resp.ImportSnapshotTasks) > 0 {
			first := &resp.ImportSnapshotTasks[0]
			if first.SnapshotTaskDetail.Status == "completed" {
				return first.SnapshotTaskDetail.SnapshotID, nil
			}
		}
		log.Printf("Still waiting; got: %s", out)
		time.Sleep(5 * time.Second)

		// TODO(bradfitz): percentage bar?
		// Looks like:
		/* 2024/05/14 13:03:21 Still waiting; got: {
		    "ImportSnapshotTasks": [
		        {
		            "ImportTaskId": "import-snap-0232251d0fbcb33fd",
		            "SnapshotTaskDetail": {
		                "DiskImageSize": 1258299392.0,
		                "Format": "RAW",
		                "Progress": "32",
		                "Status": "active",
		                "StatusMessage": "validated",
		                "Url": "s3://tskrazy-import/tskrazy.img",
		                "UserBucket": {
		                    "S3Bucket": "tskrazy-import",
		                    "S3Key": "tskrazy.img"
		                }
		            },
		            "Tags": []
		        }
		    ]
		}*/
	}
}

func makeAMI(name, ebsSnapID string) (ami string, err error) {
	var arch string
	switch conf.GOARCH() {
	case "arm64":
		arch = "arm64"
	case "amd64":
		arch = "x86_64"
	default:
		return "", fmt.Errorf("unknown arch %q", conf.GOARCH())
	}
	out, err := exec.Command("aws", "ec2", "register-image",
		"--name", name,
		"--architecture", arch,
		"--root-device-name", "/dev/sda",
		"--ena-support",
		"--imds-support", "v2.0",
		"--boot-mode", "uefi-preferred",
		"--block-device-mappings", "DeviceName=/dev/sda,Ebs={SnapshotId="+ebsSnapID+"}").CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("register image: %v: %s", err, out)
	}
	/*
		On success:
		{
		    "ImageId": "ami-052e1538166886ad2"
		}
	*/
	var resp struct {
		ImageID string `json:"ImageId"`
	}
	if err := json.Unmarshal(out, &resp); err != nil {
		return "", fmt.Errorf("unmarshal response: %v: %s", err, out)
	}
	if resp.ImageID == "" {
		return "", fmt.Errorf("empty image ID in response: %s", out)
	}
	return resp.ImageID, nil
}