package main

import (
	"flag"
	"fmt"
	"github.com/diskfs/go-diskfs/disk"
	"github.com/diskfs/go-diskfs/filesystem"
	"github.com/diskfs/go-diskfs/filesystem/fat32"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	ack "git.swzry.com/zry/go-cli-ack"
	diskfs "github.com/diskfs/go-diskfs"
	"github.com/diskfs/go-diskfs/partition/mbr"
	humanize "github.com/dustin/go-humanize"
	"github.com/liushuochen/gotable"
	yaml "gopkg.in/yaml.v3"
)

const (
	BLOCK_SIZE        = 512
	UBOOT_OFFSET      = 8192
	FIRST_PART_OFFSET = 1024 * 1024
	LAST_RESERVED     = 4 * 1024 * 1024
)

type ConfigDef struct {
	UBootImg   string `yaml:"u_boot_img"`
	BootDir    string `yaml:"boot_dir"`
	BootFsSize int64  `yaml:"boot_fs_size"`
	RootImg    string `yaml:"root_img"`
	UpperImg   string `yaml:"upper_img"`
	SwapSize   int64  `yaml:"swap_size"`
}

var OutputImg string
var PrintHelp bool
var ConfigFile string
var Cfg ConfigDef
var QuietMode bool

func main() {
	flag.BoolVar(&PrintHelp, "h", false, "print this help")
	flag.StringVar(&OutputImg, "o", "sdcard.img", "output image file")
	flag.StringVar(&ConfigFile, "i", "config.yaml", "input config file")
	flag.BoolVar(&QuietMode, "q", false, "quiet mode, which does not ask the user whether to continue")
	flag.Parse()
	if PrintHelp {
		flag.PrintDefaults()
		return
	}
	fcfg, err := os.ReadFile(ConfigFile)
	if err != nil {
		panic(fmt.Sprint("failed read config file: ", err))
	}
	err = yaml.Unmarshal(fcfg, &Cfg)
	if err != nil {
		panic(fmt.Sprint("failed parse yaml config: ", err))
	}
	if Cfg.SwapSize < 4 {
		panic("swap size should not less than 4MB.")
	}
	if Cfg.BootFsSize < 4 {
		panic("swap size should not less than 4MB.")
	}
	ubootLen, err := CheckSize("u-boot", Cfg.UBootImg)
	if err != nil {
		panic(err)
	}
	ubootMaxSize := int64(FIRST_PART_OFFSET - UBOOT_OFFSET)
	if ubootLen > ubootMaxSize {
		panic(fmt.Sprintf(
			"The maximum size of u-boot is %d bytes, but yourimage has %d bytes, which exceeds %d bytes.",
			ubootMaxSize, ubootLen, ubootMaxSize-ubootLen,
		))
	}
	CheckDir(Cfg.BootDir)
	rootLen, err := CheckSize("root filesystem", Cfg.RootImg)
	if err != nil {
		panic(err)
	}
	upperLen, err := CheckSize("upper filesystem", Cfg.UpperImg)
	if err != nil {
		panic(err)
	}
	bootActualSize := Cfg.BootFsSize * 1024 * 1024
	swapActualSize := Cfg.SwapSize * 1024 * 1024
	p1start := AlignToBlockSize(FIRST_PART_OFFSET)
	p1size := AlignToBlockSize(bootActualSize)
	p2start := p1start + p1size
	p2size := AlignToBlockSize(rootLen)
	p3start := p2start + p2size
	p3size := AlignToBlockSize(swapActualSize)
	p4start := p3start + p3size
	p4size := AlignToBlockSize(upperLen)
	diskSize := int64(p4start+p4size)*BLOCK_SIZE + LAST_RESERVED
	fmt.Println("==== Partition Table Brief ====")
	PrintPartitionTableBrief([]PTBrief{
		{
			StartSector: int(p1start),
			SizeSector:  int(p1size),
			ShowName:    "boot",
		},
		{
			StartSector: int(p2start),
			SizeSector:  int(p2size),
			ShowName:    "root",
		},
		{
			StartSector: int(p3start),
			SizeSector:  int(p3size),
			ShowName:    "swap",
		},
		{
			StartSector: int(p4start),
			SizeSector:  int(p4size),
			ShowName:    "data",
		},
	})
	fmt.Println("Total Size: ", humanize.IBytes(uint64(diskSize)))
	fmt.Println("================================")
	if !QuietMode {
		if !ack.Acknowledge("Continue?", true) {
			fmt.Println("User canceled.")
			return
		}
	}
	if !CheckExists() {
		return
	}
	idisk, err := diskfs.Create(OutputImg, diskSize, diskfs.Raw, diskfs.SectorSizeDefault)
	table := &mbr.Table{
		Partitions: []*mbr.Partition{
			{
				Bootable: false,
				Type:     mbr.Fat32LBA,
				Start:    p1start,
				Size:     p1size,
			},
			{
				Bootable: false,
				Type:     mbr.Linux,
				Start:    p2start,
				Size:     p2size,
			},
			{
				Bootable: false,
				Type:     mbr.LinuxSwap,
				Start:    p3start,
				Size:     p3size,
			},
			{
				Bootable: false,
				Type:     mbr.Linux,
				Start:    p4start,
				Size:     p4size,
			},
		},
		LogicalSectorSize:  BLOCK_SIZE,
		PhysicalSectorSize: BLOCK_SIZE,
	}
	err = idisk.Partition(table)
	if err != nil {
		panic(fmt.Sprint("failed create disk partition: ", err))
	}
	fmt.Println("processing Partition 1: boot...")
	CreateFat32Fs(idisk, 1, Cfg.BootDir)
	fmt.Println("processing Partition 2: root...")
	CopyImgToPartition(idisk, 2, Cfg.RootImg)
	fmt.Println("processing Partition 3: swap...")
	FillPartition(idisk, 3, int(swapActualSize))
	fmt.Println("processing Partition 4: data...")
	CopyImgToPartition(idisk, 4, Cfg.UpperImg)
	fmt.Println("writing bootloader...")
	WriteUboot()
	fmt.Println("==== Done ====")
}

func CheckExists() bool {
	stat, err := os.Stat(OutputImg)
	if err != nil {
		return true
	}
	if !QuietMode {
		if !ack.Acknowledge(fmt.Sprintf("'%s' already exists, override?", OutputImg), true) {
			fmt.Println("User canceled.")
			return false
		}
	}
	if stat.IsDir() {
		panic(fmt.Sprintf("'%s' is a existed dir", OutputImg))
	}
	err = os.Remove(OutputImg)
	if err != nil {
		panic(fmt.Sprintf("'%s' is a existed file, but can not remove: %v", OutputImg, err))
	}
	return true
}

func AlignToBlockSize(size int64) uint32 {
	remain := size % BLOCK_SIZE
	m := size - remain
	mr := m / BLOCK_SIZE
	if remain > 0 {
		mr = mr + 1
	}
	return uint32(mr)
}

func CheckSize(dname, fname string) (int64, error) {
	fstat, err := os.Stat(fname)
	if err != nil {
		return 0, fmt.Errorf("can not stat %s image '%s': %v", dname, fname, err)
	}
	l := fstat.Size()
	if l <= 0 {
		return 0, fmt.Errorf("%s image is empty", dname)
	}
	return l, nil
}

func CheckDir(dname string) {
	fstat, err := os.Stat(dname)
	if err != nil {
		panic(fmt.Sprintf("can not stat dir '%s': %v", dname, err))
	}
	if !fstat.IsDir() {
		panic(fmt.Sprintf("'%s' is not a dir", dname))
	}
}

type PTBrief struct {
	ShowName    string
	StartSector int
	SizeSector  int
}

func PrintPartitionTableBrief(brief []PTBrief) {
	tb, err := gotable.Create("#", "Used For", "Offset(Sectors)", "Offset(Bytes)", "Size(Sectors)", "Size(Bytes)")
	if err != nil {
		panic(err)
	}
	for i, v := range brief {
		err = tb.AddRow(map[string]string{
			"#":               strconv.Itoa(i + 1),
			"Used For":        v.ShowName,
			"Offset(Sectors)": strconv.Itoa(v.StartSector),
			"Offset(Bytes)":   humanize.IBytes(uint64(v.StartSector) * BLOCK_SIZE),
			"Size(Sectors)":   strconv.Itoa(v.SizeSector),
			"Size(Bytes)":     humanize.IBytes(uint64(v.SizeSector) * BLOCK_SIZE),
		})
		if err != nil {
			panic(err)
		}
	}
	fmt.Println(tb.String())
}

func CreateFat32Fs(d *disk.Disk, ptn int, from string) {
	spec := disk.FilesystemSpec{Partition: ptn, FSType: filesystem.TypeFat32, VolumeLabel: "boot"}
	ifs, err := d.CreateFilesystem(spec)
	if err != nil {
		panic(fmt.Sprintf("can not create filesystem for partition %d: %v", ptn, err))
	}
	err = filepath.Walk(from, func(path string, info fs.FileInfo, err error) error {
		if err != nil {
			fmt.Println("[Warn] error in '", path, "': ", err)
			return nil
		}
		relPath, err := filepath.Rel(from, path)
		if err != nil {
			fmt.Println("[Warn] error in '", path, "': can not get rel path: ", err)
			return nil
		}
		relPath = "/" + strings.TrimPrefix(relPath, ".")
		if info.IsDir() {
			err = ifs.Mkdir(relPath)
			if err != nil {
				fmt.Printf("[Warn] error in 'PART%d://%s': can not mkdir: %v\n", ptn, relPath, err)
				return nil
			}
			return nil
		} else {
			rw, xerr := ifs.OpenFile(relPath, os.O_CREATE|os.O_RDWR)
			if xerr != nil {
				fmt.Printf("[Warn] error in 'PART%d://%s': can not open file: %v\n", ptn, relPath, xerr)
				return nil
			}
			defer rw.Close()
			in, xerr := os.Open(path)
			if xerr != nil {
				fmt.Println("[Warn] error in '", path, "': can not open file: ", xerr)
				return nil
			}
			defer in.Close()
			_, xerr = io.Copy(rw, in)
			if xerr != nil {
				fmt.Printf("[Warn] io error in copy '%s' to 'PART%d://%s': %v\n", path, ptn, relPath, xerr)
				return nil
			}
		}
		return nil
	})
	if err != nil {
		panic(fmt.Sprint("error in copy files to partition 1: ", err))
	}
	_, ok := ifs.(*fat32.FileSystem)
	if !ok {
		panic(fmt.Errorf("create filesystem for partition %d failed: finally it is not a fat32 filesystem"))
	}
}

func CopyImgToPartition(d *disk.Disk, ptn int, imgfile string) {
	in, err := os.Open(imgfile)
	if err != nil {
		panic(fmt.Sprint("error in open img file '", imgfile, "': ", err))
	}
	defer in.Close()
	n, err := d.WritePartitionContents(ptn, in)
	if err != nil {
		panic(fmt.Sprintf("error in writing img '%s' to partition %d: %v '", imgfile, ptn, err))
	}
	fmt.Printf("write %s into partition %d.\n", humanize.IBytes(uint64(n)), ptn)
}

func FillPartition(d *disk.Disk, ptn int, size int) {
	in := NewSizableZeroReader(size)
	n, err := d.WritePartitionContents(ptn, in)
	if err != nil {
		panic(fmt.Sprintf("error in filling partition %d: %v '", ptn, err))
	}
	fmt.Printf("fill %s into partition %d.\n", humanize.IBytes(uint64(n)), ptn)
}

func WriteUboot() {
	fUboot, err := os.OpenFile(Cfg.UBootImg, os.O_RDONLY|os.O_EXCL, 0666)
	if err != nil {
		panic(fmt.Sprintf("error in open u-boot img: %v '", err))
	}
	defer fUboot.Close()
	fDiskImg, err := os.OpenFile(OutputImg, os.O_EXCL|os.O_RDWR, 0666)
	if err != nil {
		panic(fmt.Sprintf("error in open disk img for writting bootloader: %v '", err))
	}
	defer fUboot.Close()
	n, err := fDiskImg.Seek(UBOOT_OFFSET, 0)
	if err != nil {
		panic(fmt.Sprintf("error in seeking img for writting bootloader: %v '", err))
	}
	if n != UBOOT_OFFSET {
		panic(fmt.Sprintf("error in seeking img for writting bootloader: current pos %d != UBOOT_OFFSET(%d) '", n, UBOOT_OFFSET))
	}
	n, err = io.Copy(fDiskImg, fUboot)
	if err != nil {
		panic(fmt.Sprintf("io error in writting bootloader: %v '", err))
	}
	fmt.Printf("%d bytes written.\n", n)
}