package image

import (
	"os"
	"path/filepath"
	"testing"

	"metis/pkg/inject"
)

func TestInjectRootFSWithFakes(t *testing.T) {
	scripts := fakeRootfsCommands(t, true)
	t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))

	imagePath := filepath.Join(t.TempDir(), "image.img")
	if err := os.WriteFile(imagePath, make([]byte, 4096), 0o644); err != nil {
		t.Fatal(err)
	}
	files := []inject.FileSpec{
		{Path: "etc/metis/firstboot.env", Content: []byte("METIS_HOSTNAME='titan-13'\n"), Mode: 0o600, RootFS: true},
		{Path: "usr/local/sbin/test.sh", Content: []byte("#!/usr/bin/env bash\nexit 0\n"), Mode: 0o755, RootFS: true},
	}
	if err := InjectRootFS(imagePath, files); err != nil {
		t.Fatalf("InjectRootFS: %v", err)
	}
}

func TestInjectRootFSSkipsBootOnlyFiles(t *testing.T) {
	imagePath := filepath.Join(t.TempDir(), "image.img")
	if err := os.WriteFile(imagePath, make([]byte, 4096), 0o644); err != nil {
		t.Fatal(err)
	}
	if err := InjectRootFS(imagePath, []inject.FileSpec{{Path: "user-data", Content: []byte("boot"), RootFS: false}}); err != nil {
		t.Fatalf("InjectRootFS boot-only: %v", err)
	}
}

func TestFindLinuxPartitionAndTypeChecks(t *testing.T) {
	scripts := fakeRootfsCommands(t, true)
	t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))
	imagePath := filepath.Join(t.TempDir(), "image.img")
	if err := os.WriteFile(imagePath, make([]byte, 4096), 0o644); err != nil {
		t.Fatal(err)
	}
	part, sector, err := findLinuxPartition(imagePath)
	if err != nil {
		t.Fatalf("findLinuxPartition: %v", err)
	}
	if sector != 512 || part.Start != 1 || part.Size != 2 {
		t.Fatalf("unexpected partition info: %+v sector=%d", part, sector)
	}
	if !isLinuxPartitionType("83") || !isLinuxPartitionType("8300") || !isLinuxPartitionType("0fc63daf-8483-4772-8e79-3d69d8477de4") {
		t.Fatal("expected linux partition types to match")
	}
	if isLinuxPartitionType("ef") {
		t.Fatal("did not expect non-linux type to match")
	}
}

func TestFindLinuxPartitionReturnsErrorWhenNoLinuxPartitionExists(t *testing.T) {
	scripts := fakeRootfsCommands(t, false)
	t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))
	imagePath := filepath.Join(t.TempDir(), "image.img")
	if err := os.WriteFile(imagePath, make([]byte, 4096), 0o644); err != nil {
		t.Fatal(err)
	}
	if _, _, err := findLinuxPartition(imagePath); err == nil {
		t.Fatal("expected error without Linux partition")
	}
}

func TestParentDirs(t *testing.T) {
	got := parentDirs("etc/metis/firstboot.env")
	want := []string{"/etc", "/etc/metis"}
	if len(got) != len(want) {
		t.Fatalf("parentDirs length mismatch: got %v want %v", got, want)
	}
	for i := range want {
		if got[i] != want[i] {
			t.Fatalf("parentDirs[%d] = %q want %q", i, got[i], want[i])
		}
	}
}

func TestRootfsErrorBranches(t *testing.T) {
	part := partitionTablePart{Start: 1, Size: 2}
	dir := t.TempDir()
	src := filepath.Join(dir, "src.img")
	dst := filepath.Join(dir, "dst.img")
	if err := os.WriteFile(src, make([]byte, 512), 0o644); err != nil {
		t.Fatal(err)
	}
	if err := extractPartition(src, dst, part, 512); err == nil {
		t.Fatal("expected extractPartition to fail on short source image")
	}
	if err := os.WriteFile(dst, make([]byte, 512), 0o644); err != nil {
		t.Fatal(err)
	}
	if err := replacePartition(src, dst, part, 512); err == nil {
		t.Fatal("expected replacePartition size mismatch")
	}
}

func fakeRootfsCommands(t *testing.T, includeLinux bool) string {
	t.Helper()
	dir := t.TempDir()
	write := func(name, body string) {
		path := filepath.Join(dir, name)
		if err := os.WriteFile(path, []byte("#!/usr/bin/env bash\nset -eu\n"+body+"\n"), 0o755); err != nil {
			t.Fatalf("write %s: %v", name, err)
		}
	}
	partitions := `{"partitiontable":{"sectorsize":512,"partitions":[{"start":3,"size":1,"type":"ef"},{"start":1,"size":2,"type":"ef"}]}}`
	if includeLinux {
		partitions = `{"partitiontable":{"sectorsize":512,"partitions":[{"start":3,"size":1,"type":"ef"},{"start":1,"size":2,"type":"83"}]}}`
	}
	write("sfdisk", "cat <<'JSON'\n"+partitions+"\nJSON")
	write("debugfs", `if [[ "${1:-}" == "-w" ]]; then
  exit 0
fi
if [[ "${1:-}" == "-R" ]]; then
  set -- $2
  case "${1:-}" in
    stat)
      case "${2:-}" in
        /etc/metis/firstboot.env) printf 'Mode:  0600\n' ;;
        /usr/local/sbin/test.sh) printf 'Mode:  0755\n' ;;
      esac
      exit 0
      ;;
    dump)
      dest="${3:-}"
      case "${2:-}" in
        /etc/metis/firstboot.env) printf "METIS_HOSTNAME='titan-13'\n" > "${dest}" ;;
        /usr/local/sbin/test.sh) printf '#!/usr/bin/env bash\nexit 0\n' > "${dest}" ;;
      esac
      exit 0
      ;;
  esac
fi
exit 0`)
	return dir
}
