package main

import (
	"bufio"
	"compress/gzip"
	"errors"
	"fmt"
	"io"
	"os"
	"path"
	"sort"
	"strings"
	"time"
)

type archiveEntry struct {
	name string
	mode uint32
	data []byte
}

func writeInitramfs(output string, initBinary []byte) error {
	return writeInitramfsWithEntries(output, initBinary, nil)
}

func writeInitramfsWithEntries(output string, initBinary []byte, extra []archiveEntry) error {
	if len(initBinary) == 0 {
		return errors.New("init binary is empty")
	}
	entries := []archiveEntry{
		{name: "dev", mode: 0o040755},
		{name: "etc", mode: 0o040755},
		{name: "etc/tera", mode: 0o040755},
		{name: "etc/tera/update-trust", mode: 0o040555},
		{name: "proc", mode: 0o040555},
		{name: "sys", mode: 0o040555},
		{name: "userdata", mode: 0o040755},
		{name: "init", mode: 0o100755, data: initBinary},
	}
	entries = append(entries, extra...)
	seen := make(map[string]struct{}, len(entries))
	for _, entry := range entries {
		if _, duplicate := seen[entry.name]; duplicate {
			return fmt.Errorf("duplicate initramfs path %q", entry.name)
		}
		seen[entry.name] = struct{}{}
	}
	sort.Slice(entries, func(i, j int) bool { return entries[i].name < entries[j].name })

	file, err := os.OpenFile(output, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
	if err != nil {
		return err
	}
	committed := false
	defer func() {
		_ = file.Close()
		if !committed {
			_ = os.Remove(output)
		}
	}()
	gzipWriter, err := gzip.NewWriterLevel(file, gzip.BestCompression)
	if err != nil {
		return err
	}
	gzipWriter.Header.ModTime = time.Unix(0, 0).UTC()
	gzipWriter.Header.OS = 255
	buffered := bufio.NewWriterSize(gzipWriter, 1024*1024)
	for index, entry := range entries {
		if err := writeNewcEntry(buffered, uint32(index+1), entry); err != nil {
			return err
		}
	}
	if err := writeNewcEntry(buffered, uint32(len(entries)+1), archiveEntry{name: "TRAILER!!!"}); err != nil {
		return err
	}
	if err := buffered.Flush(); err != nil {
		return err
	}
	if err := gzipWriter.Close(); err != nil {
		return err
	}
	if err := file.Sync(); err != nil {
		return err
	}
	if err := file.Close(); err != nil {
		return err
	}
	committed = true
	return nil
}

func writeNewcEntry(writer io.Writer, inode uint32, entry archiveEntry) error {
	clean := path.Clean(strings.TrimPrefix(entry.name, "/"))
	if clean == "." || clean != entry.name || strings.Contains(clean, "..") {
		return fmt.Errorf("unsafe initramfs path %q", entry.name)
	}
	name := append([]byte(clean), 0)
	header := fmt.Sprintf("070701%08x%08x%08x%08x%08x%08x%08x%08x%08x%08x%08x%08x%08x",
		inode, entry.mode, 0, 0, 1, 0, len(entry.data), 0, 0, 0, 0, len(name), 0)
	if len(header) != 110 {
		return errors.New("internal newc header length error")
	}
	if _, err := io.WriteString(writer, header); err != nil {
		return err
	}
	if _, err := writer.Write(name); err != nil {
		return err
	}
	if err := pad4(writer, len(header)+len(name)); err != nil {
		return err
	}
	if _, err := writer.Write(entry.data); err != nil {
		return err
	}
	return pad4(writer, len(entry.data))
}

func pad4(writer io.Writer, length int) error {
	padding := (4 - length%4) % 4
	if padding == 0 {
		return nil
	}
	_, err := writer.Write(make([]byte, padding))
	return err
}