package runtimebundle

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

func TestBootstrapExtractsBundleAndValidatesChecksum(t *testing.T) {
	manifestDir := t.TempDir()
	bundleData := buildArchive(t, map[string]string{
		"runtime/firecracker":                                       "fc",
		"runtime/id_ed25519":                                        "key",
		"runtime/namegen":                                           "namegen",
		"runtime/customize.sh":                                      "#!/bin/bash\n",
		"runtime/packages.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",
		"runtime/bundle.json":                                       mustJSON(t, BundleMetadata{FirecrackerBin: "firecracker", SSHKeyPath: "id_ed25519", NamegenPath: "namegen", CustomizeScript: "customize.sh", DefaultPackages: "packages.apt", DefaultRootfs: "rootfs-docker.ext4", DefaultKernel: "wtf/root/boot/vmlinux-6.8.0-94-generic", DefaultInitrd: "wtf/root/boot/initrd.img-6.8.0-94-generic", DefaultModulesDir: "wtf/root/lib/modules/6.8.0-94-generic"}),
	})
	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 TestBootstrapRejectsMissingURLWithLocalManifestGuidance(t *testing.T) {
	manifest := Manifest{
		SHA256:        strings.Repeat("0", 64),
		BundleRoot:    "runtime",
		RequiredPaths: []string{"firecracker"},
	}
	err := Bootstrap(context.Background(), manifest, filepath.Join(t.TempDir(), "runtime-bundle.toml"), filepath.Join(t.TempDir(), "runtime"))
	if err == nil || !strings.Contains(err.Error(), "local manifest copy") {
		t.Fatalf("Bootstrap() error = %v, want local manifest guidance", err)
	}
}

func TestBootstrapRejectsMissingSHAWithArchiveGuidance(t *testing.T) {
	manifest := Manifest{
		URL:           "./bundle.tar.gz",
		BundleRoot:    "runtime",
		RequiredPaths: []string{"firecracker"},
	}
	err := Bootstrap(context.Background(), manifest, filepath.Join(t.TempDir(), "runtime-bundle.toml"), filepath.Join(t.TempDir(), "runtime"))
	if err == nil || !strings.Contains(err.Error(), "staged or published runtime bundle archive") {
		t.Fatalf("Bootstrap() error = %v, want archive guidance", err)
	}
}

func TestPackageWritesArchive(t *testing.T) {
	runtimeDir := t.TempDir()
	for _, rel := range []string{
		"firecracker",
		"id_ed25519",
		"namegen",
		"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",
		BundleMeta: BundleMetadata{
			FirecrackerBin:    "firecracker",
			SSHKeyPath:        "id_ed25519",
			NamegenPath:       "namegen",
			CustomizeScript:   "customize.sh",
			DefaultPackages:   "packages.apt",
			DefaultRootfs:     "rootfs-docker.ext4",
			DefaultKernel:     "wtf/root/boot/vmlinux-6.8.0-94-generic",
			DefaultInitrd:     "wtf/root/boot/initrd.img-6.8.0-94-generic",
			DefaultModulesDir: "wtf/root/lib/modules/6.8.0-94-generic",
		},
		RequiredPaths: []string{
			"firecracker",
			"id_ed25519",
			"namegen",
			"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)
	}
	runtimeOut := filepath.Join(t.TempDir(), "runtime")
	if err := Bootstrap(context.Background(), Manifest{
		URL:           outArchive,
		SHA256:        sum,
		BundleRoot:    "runtime",
		RequiredPaths: manifest.RequiredPaths,
	}, filepath.Join(t.TempDir(), "runtime-bundle.toml"), runtimeOut); err != nil {
		t.Fatalf("Bootstrap packaged archive: %v", err)
	}
	if _, err := os.Stat(filepath.Join(runtimeOut, BundleMetadataFile)); err != nil {
		t.Fatalf("bundle metadata missing after bootstrap: %v", err)
	}
	meta, err := LoadBundleMetadata(runtimeOut)
	if err != nil {
		t.Fatalf("LoadBundleMetadata: %v", err)
	}
	if meta.DefaultRootfs != manifest.BundleMeta.DefaultRootfs {
		t.Fatalf("DefaultRootfs = %q, want %q", meta.DefaultRootfs, manifest.BundleMeta.DefaultRootfs)
	}
}

func TestLoadBundleMetadataRejectsMissingRequiredPath(t *testing.T) {
	runtimeDir := t.TempDir()
	for _, rel := range []string{"firecracker", "id_ed25519", "namegen", "customize.sh", "packages.apt", "rootfs-docker.ext4"} {
		path := filepath.Join(runtimeDir, rel)
		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)
		}
	}
	data := mustJSON(t, BundleMetadata{
		FirecrackerBin:  "firecracker",
		SSHKeyPath:      "id_ed25519",
		NamegenPath:     "namegen",
		CustomizeScript: "customize.sh",
		DefaultPackages: "packages.apt",
		DefaultRootfs:   "rootfs-docker.ext4",
		DefaultKernel:   "missing-kernel",
	})
	if err := os.WriteFile(filepath.Join(runtimeDir, BundleMetadataFile), []byte(data), 0o644); err != nil {
		t.Fatalf("WriteFile: %v", err)
	}
	if _, err := LoadBundleMetadata(runtimeDir); err == nil || !strings.Contains(err.Error(), "default_kernel") {
		t.Fatalf("LoadBundleMetadata() error = %v, want default_kernel failure", 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[:])
}

func mustJSON(t *testing.T, value any) string {
	t.Helper()
	data, err := json.Marshal(value)
	if err != nil {
		t.Fatalf("Marshal: %v", err)
	}
	return string(data)
}