package runtimebundle

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func TestBootstrapExtractsBundleAndValidatesChecksum(t *testing.T) {
	manifestDir := t.TempDir()
	bundleData := buildArchive(t, map[string]string{
		"runtime/firecracker":                                       "fc",
		"runtime/customize.sh":                                      "#!/bin/bash\n",
		"runtime/packages.apt":                                      "vim\n",
		"runtime/rootfs-docker.ext4":                                "rootfs",
		"runtime/wtf/root/boot/vmlinux-6.8.0-94-generic":            "kernel",
		"runtime/wtf/root/boot/initrd.img-6.8.0-94-generic":         "initrd",
		"runtime/wtf/root/lib/modules/6.8.0-94-generic/modules.dep": "dep",
	})
	archivePath := filepath.Join(manifestDir, "bundle.tar.gz")
	if err := os.WriteFile(archivePath, bundleData, 0o644); err != nil {
		t.Fatalf("WriteFile: %v", err)
	}

	manifest := Manifest{
		URL:           "./bundle.tar.gz",
		SHA256:        sha256Hex(bundleData),
		BundleRoot:    "runtime",
		RequiredPaths: []string{"firecracker", "customize.sh", "packages.apt", "rootfs-docker.ext4", "wtf/root/boot/vmlinux-6.8.0-94-generic", "wtf/root/lib/modules/6.8.0-94-generic"},
	}
	outDir := filepath.Join(t.TempDir(), "runtime")
	if err := Bootstrap(context.Background(), manifest, filepath.Join(manifestDir, "runtime-bundle.toml"), outDir); err != nil {
		t.Fatalf("Bootstrap: %v", err)
	}
	for _, rel := range manifest.RequiredPaths {
		if _, err := os.Stat(filepath.Join(outDir, rel)); err != nil {
			t.Fatalf("runtime missing %s: %v", rel, err)
		}
	}
}

func TestBootstrapRejectsChecksumMismatch(t *testing.T) {
	manifestDir := t.TempDir()
	archivePath := filepath.Join(manifestDir, "bundle.tar.gz")
	if err := os.WriteFile(archivePath, []byte("not-a-tarball"), 0o644); err != nil {
		t.Fatalf("WriteFile: %v", err)
	}
	manifest := Manifest{
		URL:           "./bundle.tar.gz",
		SHA256:        strings.Repeat("0", 64),
		BundleRoot:    "runtime",
		RequiredPaths: []string{"firecracker"},
	}
	err := Bootstrap(context.Background(), manifest, filepath.Join(manifestDir, "runtime-bundle.toml"), filepath.Join(t.TempDir(), "runtime"))
	if err == nil || !strings.Contains(err.Error(), "checksum mismatch") {
		t.Fatalf("Bootstrap() error = %v, want checksum mismatch", err)
	}
}

func TestPackageWritesArchive(t *testing.T) {
	runtimeDir := t.TempDir()
	for _, rel := range []string{
		"firecracker",
		"customize.sh",
		"packages.apt",
		"rootfs-docker.ext4",
		"wtf/root/boot/vmlinux-6.8.0-94-generic",
		"wtf/root/boot/initrd.img-6.8.0-94-generic",
		"wtf/root/lib/modules/6.8.0-94-generic",
	} {
		path := filepath.Join(runtimeDir, rel)
		if rel == "wtf/root/lib/modules/6.8.0-94-generic" {
			if err := os.MkdirAll(path, 0o755); err != nil {
				t.Fatalf("MkdirAll: %v", err)
			}
			if err := os.WriteFile(filepath.Join(path, "modules.dep"), []byte(rel), 0o644); err != nil {
				t.Fatalf("WriteFile: %v", err)
			}
			continue
		}
		if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
			t.Fatalf("MkdirAll: %v", err)
		}
		if err := os.WriteFile(path, []byte(rel), 0o644); err != nil {
			t.Fatalf("WriteFile: %v", err)
		}
	}
	manifest := Manifest{
		BundleRoot: "runtime",
		RequiredPaths: []string{
			"firecracker",
			"customize.sh",
			"packages.apt",
			"rootfs-docker.ext4",
			"wtf/root/boot/vmlinux-6.8.0-94-generic",
			"wtf/root/boot/initrd.img-6.8.0-94-generic",
			"wtf/root/lib/modules/6.8.0-94-generic",
		},
	}
	outArchive := filepath.Join(t.TempDir(), "bundle.tar.gz")
	sum, err := Package(runtimeDir, outArchive, manifest)
	if err != nil {
		t.Fatalf("Package: %v", err)
	}
	if sum == "" {
		t.Fatalf("Package() returned empty checksum")
	}
	if _, err := os.Stat(outArchive); err != nil {
		t.Fatalf("archive missing: %v", err)
	}
}

func buildArchive(t *testing.T, files map[string]string) []byte {
	t.Helper()
	var buf bytes.Buffer
	gz := gzip.NewWriter(&buf)
	tw := tar.NewWriter(gz)
	for name, contents := range files {
		header := &tar.Header{
			Name: name,
			Mode: 0o644,
			Size: int64(len(contents)),
		}
		if err := tw.WriteHeader(header); err != nil {
			t.Fatalf("WriteHeader(%s): %v", name, err)
		}
		if _, err := tw.Write([]byte(contents)); err != nil {
			t.Fatalf("Write(%s): %v", name, err)
		}
	}
	if err := tw.Close(); err != nil {
		t.Fatalf("Close tar: %v", err)
	}
	if err := gz.Close(); err != nil {
		t.Fatalf("Close gzip: %v", err)
	}
	return buf.Bytes()
}

func sha256Hex(data []byte) string {
	sum := sha256.Sum256(data)
	return hex.EncodeToString(sum[:])
}