package main

import (
	"archive/tar"
	"bytes"
	"crypto"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strings"

	"github.com/cavaliergopher/cpio"
	"github.com/klauspost/compress/zstd"

	"novit.tech/direktil/pkg/base64"
	"novit.tech/direktil/pkg/config"
	"novit.tech/direktil/pkg/cpiocat"
	"novit.tech/direktil/pkg/localconfig"
)

func renderBootstrapConfig(w http.ResponseWriter, ctx *renderContext) (err error) {
	log.Printf("sending bootstrap config for %q", ctx.Host.Name)

	ba, err := ctx.BootstrapConfig()
	if err != nil {
		return err
	}

	_, err = w.Write(ba)
	return
}

func buildInitrd(out io.Writer, ctx *renderContext) (err error) {
	_, cfg, err := ctx.Config()
	if err != nil {
		return fmt.Errorf("render config failed: %w", err)
	}

	zout, err := zstd.NewWriter(out, zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(12)))
	if err != nil {
		return fmt.Errorf("zstd writer setup failed: %w", err)
	}

	cat := cpiocat.New(zout)

	// initrd
	initrdPath, err := distFetch("initrd", ctx.Host.Initrd)
	if err != nil {
		return fmt.Errorf("fetch initrd failed: %w", err)
	}
	cat.AppendArchFile(initrdPath)

	// copy config's udev rules (FIXME: too much logic for dls)
	for _, file := range cfg.Files {
		if !strings.HasPrefix(file.Path, "/etc/udev/rules.d/") {
			continue
		}

		content := []byte(file.Content)
		if c := file.Content64; c != "" {
			content, err = base64.Decode(c)
			if err != nil {
				return fmt.Errorf("decode %s failed: %w", file.Path, err)
			}
		}

		mode := file.Mode
		if mode == 0 {
			mode = 0o644
		}

		cat.AppendBytes(content, file.Path, cpio.FileMode(mode))
	}

	// embedded layers (modules)
	for _, layer := range cfg.Layers {
		switch layer {
		case "modules":
			layerVersion := ctx.Host.Versions[layer]
			if layerVersion == "" {
				layerVersion = ctx.Host.Kernel
			}
			modulesPath, err := distFetch("layers", layer, layerVersion)
			if err != nil {
				return fmt.Errorf("fetch layer %s failed: %w", layer, err)
			}

			cat.AppendFile(modulesPath, "modules.sqfs")
		}
	}

	// config
	cfgBytes, err := ctx.BootstrapConfig()
	if err != nil {
		return fmt.Errorf("render bootstrap config failed: %w", err)
	}
	cat.AppendBytes(cfgBytes, "config.yaml", 0o600)

	// ssh keys
	cat.AppendDir("/etc", 0o755)
	cat.AppendDir("/etc/ssh", 0o700)

	// XXX do we want bootstrap-stage keys instead of the real host key?
	for _, format := range sshKeyTypes {
		keyPath := "/etc/ssh/ssh_host_" + format + "_key"
		cat.AppendBytes(cfg.FileContent(keyPath), keyPath, 0o600)
	}

	// ssh user CA
	userCA, err := sshCAPubKey(ctx.Host.ClusterName)
	if err != nil {
		return fmt.Errorf("failed to get SSH user CA: %w", err)
	}

	cat.AppendBytes(userCA, "user_ca.pub", 0600)

	if err = cat.Close(); err != nil {
		return fmt.Errorf("cpio close failed: %w", err)
	}

	if err = zout.Close(); err != nil {
		return fmt.Errorf("zstd close failed: %w", err)
	}

	return
}

func getSigner(clusterName string) (signer crypto.Signer, err error) {
	ca, err := getUsableClusterCA(clusterName, "boot-signer")
	if err != nil {
		return
	}
	return ca.ParseKey()
}

func buildBootstrap(out io.Writer, ctx *renderContext) (err error) {
	arch := tar.NewWriter(out)
	defer arch.Close()

	signer, err := getSigner(ctx.Host.ClusterName)
	if err != nil {
		return
	}

	hash := crypto.SHA512

	sign := func(name string, digest []byte) (err error) {
		sigBytes, err := signer.Sign(nil, digest, hash)
		if err != nil {
			err = fmt.Errorf("signing to %s failed: %w", name, err)
			return err
		}

		if err = arch.WriteHeader(&tar.Header{
			Name: name,
			Size: int64(len(sigBytes)),
			Mode: 0o644,
		}); err != nil {
			return
		}

		_, err = io.Copy(arch, bytes.NewReader(sigBytes))
		return
	}

	// config
	cfgBytes, cfg, err := ctx.Config()
	if err != nil {
		return err
	}

	err = arch.WriteHeader(&tar.Header{
		Name: "config.yaml",
		Size: int64(len(cfgBytes)),
		Mode: 0o600,
	})
	if err != nil {
		return
	}

	_, err = arch.Write(cfgBytes)
	if err != nil {
		return
	}

	{
		h := hash.New()
		h.Write(cfgBytes)
		err = sign("config.yaml.sig", h.Sum(nil))
		if err != nil {
			return
		}
	}

	// layers
	appendSignedLayer := func(layer, layerPath string) (err error) {
		f, err := os.Open(layerPath)
		if err != nil {
			return err
		}

		defer f.Close()

		stat, err := f.Stat()
		if err != nil {
			return err
		}

		h := hash.New()
		reader := io.TeeReader(f, h)

		if err = arch.WriteHeader(&tar.Header{
			Name: layer,
			Size: stat.Size(),
			Mode: 0o600,
		}); err != nil {
			return err
		}

		_, err = io.Copy(arch, reader)
		if err != nil {
			return err
		}

		digest := h.Sum(nil)
		err = sign(layer+".sig", digest)

		return
	}

	allErofs := true

	for _, layer := range cfg.Layers {
		if layer == "modules" {
			continue // modules are in the initrd with boot v2
		}

		if !strings.HasSuffix(ctx.Host.Versions[layer], ".erofs") {
			allErofs = false
			break
		}
	}

	if allErofs {
		layerPath, e := layersCombo(ctx.Host, cfg, signer)
		if e != nil {
			err = e
			return
		}

		if err = appendSignedLayer("merged", layerPath); err != nil {
			return
		}
	} else {
		for _, layer := range cfg.Layers {
			if layer == "modules" {
				continue // modules are in the initrd with boot v2
			}

			layerPath, e := fetchHostLayer(ctx.Host, layer)
			if e != nil {
				err = e
				return
			}

			if err = appendSignedLayer(layer+".fs", layerPath); err != nil {
				return
			}
		}
	}

	return nil
}

func layersCombo(host *localconfig.Host, cfg *config.Config, signer crypto.Signer) (path string, err error) {
	key := layersComboKey(host, cfg)
	return opMutex(key, func() (path string, err error) {
		return buildLayersCombo(host, cfg, signer)
	})
}

func layersComboKey(host *localconfig.Host, cfg *config.Config) string {
	key := new(strings.Builder)
	key.WriteString("layers")
	for _, layer := range cfg.Layers {
		if layer == "modules" {
			continue
		}
		key.WriteByte(':')
		key.WriteString(layer)
		if v, ok := host.Versions[layer]; ok {
			key.WriteByte('@')
			key.WriteString(v)
		}
	}

	return key.String()
}

func buildLayersCombo(host *localconfig.Host, cfg *config.Config, signer crypto.Signer) (path string, err error) {
	path = filepath.Join(*dataDir, "cache")
	if err = os.MkdirAll(path, 0o700); err != nil {
		return
	}

	key := layersComboKey(host, cfg)
	path = filepath.Join(path, key) + ".fs"

	if _, statErr := os.Stat(path); statErr == nil {
		return // exists -> already done
	}

	workdir, err := os.MkdirTemp("/tmp", "layers")
	if err != nil {
		return
	}

	defer os.RemoveAll(workdir)

	tmpTar := filepath.Join(workdir, "output.tar")

	layers := slices.Clone(cfg.Layers)
	slices.Reverse(layers)

	cmdOut := new(bytes.Buffer)

	for i, layer := range layers {
		if layer == "modules" {
			continue // modules are in the initrd with boot v2
		}

		layerFile, e := fetchHostLayer(host, layer)
		if e != nil {
			err = e
			return
		}

		mountPoint := filepath.Join(workdir, layer)
		os.MkdirAll(mountPoint, 0700)

		if e := exec.Command("erofsfuse", layerFile, mountPoint).Run(); e != nil {
			err = fmt.Errorf("erofsfuse %s %s failed: %w", layerFile, mountPoint, e)
			return
		}

		defer func() {
			if err := exec.Command("umount", mountPoint).Run(); err != nil {
				log.Printf("umount %s failed: %v", mountPoint, err)
			}
		}()

		mode := "--append"
		if i == 0 {
			mode = "--create"
		}

		if err = runCmd(cmdOut, "tar", mode, "-p", "-f", tmpTar, "-C", mountPoint, "."); err != nil {
			return
		}
	}

	fsOut := filepath.Join(workdir, "output.fs")
	if err = runCmd(cmdOut, "mkfs.erofs", "-z", "lzma", "-C131072", "-Efragments,ztailpacking",
		"-T0", "--all-time", "--ignore-mtime", "--tar=f", fsOut, tmpTar); err != nil {
		return
	}

	hashOut := filepath.Join(workdir, "output.hash")

	if err = runCmd(cmdOut, "veritysetup", "format", fsOut, hashOut); err != nil {
		return
	}

	var rootHash []byte
	for line := range strings.SplitSeq(cmdOut.String(), "\n") {
		v, ok := strings.CutPrefix(line, "Root hash:")
		if !ok {
			continue
		}
		v = strings.TrimSpace(v)

		b, e := hex.DecodeString(v)
		if e != nil {
			err = fmt.Errorf("invalid root hash: %w", e)
			return
		}

		rootHash = b
		break
	}

	if len(rootHash) == 0 {
		err = fmt.Errorf("root hash not found in output")
		return
	}

	sigBytes, err := signer.Sign(nil, rootHash, crypto.SHA256)
	if err != nil {
		err = fmt.Errorf("root hash signature failed: %w", err)
		return
	}

	outPath := path + ".tmp"
	err = func() (err error) {
		fsRd, e := os.Open(fsOut)
		if e != nil {
			return e
		}
		defer fsRd.Close()
		hashRd, e := os.Open(hashOut)
		if e != nil {
			return e
		}
		defer hashRd.Close()

		fsStat, e := fsRd.Stat()
		if e != nil {
			return e
		}
		hashStat, e := hashRd.Stat()
		if e != nil {
			return e
		}

		out, err := os.Create(outPath)
		if err != nil {
			return
		}
		defer out.Close()

		append := func(sz uint64, rd io.Reader) {
			if err != nil {
				return
			}
			szB := make([]byte, 8)
			binary.BigEndian.PutUint64(szB, sz)
			_, err = out.Write(szB)
			if err != nil {
				return
			}

			_, err = io.Copy(out, rd)
		}

		append(uint64(len(sigBytes)), bytes.NewBuffer(sigBytes))
		append(uint64(len(rootHash)), bytes.NewBuffer(rootHash))
		append(uint64(fsStat.Size()), fsRd)
		append(uint64(hashStat.Size()), hashRd)

		if err != nil {
			return
		}
		err = out.Close()
		return
	}()
	if err != nil {
		err = fmt.Errorf("assembly failed: %w", err)
		return
	}

	err = os.Rename(outPath, path)
	return
}

func fetchHostLayer(host *localconfig.Host, layer string) (path string, err error) {
	layerVersion := host.Versions[layer]
	if layerVersion == "" {
		return "", fmt.Errorf("layer %q not mapped to a version", layer)
	}

	return distFetch("layers", layer, layerVersion)
}

func runCmd(cmdOut *bytes.Buffer, prog string, arg ...string) error {
	cmdOut.Reset()

	cmd := exec.Command(prog, arg...)
	cmd.Stdout = cmdOut
	cmd.Stderr = os.Stderr
	if e := cmd.Run(); e != nil {
		return fmt.Errorf("%s %q failed: %w", cmd.Path, cmd.Args, e)
	}
	return nil
}
