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, nil); err == nil {
		t.Fatal("expected replacePartition size mismatch")
	}
}

func TestRootfsAdditionalTopLevelErrorBranches(t *testing.T) {
	dir := t.TempDir()
	imagePath := filepath.Join(dir, "image.img")
	if err := os.WriteFile(imagePath, make([]byte, 1024), 0o644); err != nil {
		t.Fatal(err)
	}
	rootFiles := []inject.FileSpec{{Path: "etc/metis/node.json", Content: []byte("{}"), Mode: 0o600, RootFS: true}}

	t.Run("progress and partition error", func(t *testing.T) {
		scripts := fakeImageCommandDir(t, map[string]string{"sfdisk": `exit 3`})
		t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))
		var steps []string
		err := InjectRootFSWithProgress(imagePath, rootFiles, func(update RootFSProgressUpdate) {
			steps = append(steps, update.Step)
		})
		if err == nil || len(steps) != 1 || steps[0] != RootFSProgressFindingPartition {
			t.Fatalf("expected partition error after progress, steps=%v err=%v", steps, err)
		}
	})

	t.Run("invalid partition json", func(t *testing.T) {
		scripts := fakeImageCommandDir(t, map[string]string{"sfdisk": `printf '{bad-json'`})
		t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))
		if _, _, err := findLinuxPartition(imagePath); err == nil {
			t.Fatal("expected partition JSON decode failure")
		}
	})

	t.Run("default sector size", func(t *testing.T) {
		scripts := fakeImageCommandDir(t, map[string]string{"sfdisk": `cat <<'JSON'
{"partitiontable":{"sectorsize":0,"partitions":[{"start":1,"size":1,"type":"83"}]}}
JSON`})
		t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))
		_, sector, err := findLinuxPartition(imagePath)
		if err != nil || sector != 512 {
			t.Fatalf("expected default sector size, sector=%d err=%v", sector, err)
		}
	})

	t.Run("temp dir failure", func(t *testing.T) {
		scripts := fakeImageCommandDir(t, map[string]string{"sfdisk": `cat <<'JSON'
{"partitiontable":{"sectorsize":512,"partitions":[{"start":1,"size":1,"type":"83"}]}}
JSON`})
		t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))
		tmpFile := filepath.Join(dir, "tmp-file")
		if err := os.WriteFile(tmpFile, []byte("not-a-dir"), 0o644); err != nil {
			t.Fatal(err)
		}
		t.Setenv("METIS_ROOTFS_TMP_DIR", tmpFile)
		if err := InjectRootFSWithProgress(imagePath, rootFiles, nil); err == nil {
			t.Fatal("expected rootfs temp-dir failure")
		}
	})

	t.Run("extract failure", func(t *testing.T) {
		shortImage := filepath.Join(dir, "short.img")
		if err := os.WriteFile(shortImage, make([]byte, 512), 0o644); err != nil {
			t.Fatal(err)
		}
		scripts := fakeImageCommandDir(t, map[string]string{"sfdisk": `cat <<'JSON'
{"partitiontable":{"sectorsize":512,"partitions":[{"start":1,"size":4,"type":"83"}]}}
JSON`})
		t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))
		if err := InjectRootFSWithProgress(shortImage, rootFiles, nil); err == nil {
			t.Fatal("expected partition extraction failure")
		}
	})

	t.Run("write files failure", func(t *testing.T) {
		scripts := fakeImageCommandDir(t, map[string]string{
			"sfdisk": `cat <<'JSON'
{"partitiontable":{"sectorsize":512,"partitions":[{"start":1,"size":1,"type":"83"}]}}
JSON`,
			"debugfs": `exit 5`,
		})
		t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))
		if err := InjectRootFSWithProgress(imagePath, rootFiles, nil); err == nil {
			t.Fatal("expected debugfs write failure")
		}
	})
}

func TestRootfsDirectHelperErrorBranches(t *testing.T) {
	dir := t.TempDir()
	part := partitionTablePart{Start: 1, Size: 1}
	if err := extractPartition(filepath.Join(dir, "missing.img"), filepath.Join(dir, "out.img"), part, 512); err == nil {
		t.Fatal("expected extract source open failure")
	}

	srcDir := filepath.Join(dir, "source-dir")
	if err := os.Mkdir(srcDir, 0o755); err != nil {
		t.Fatal(err)
	}
	if err := extractPartition(srcDir, filepath.Join(dir, "seek.img"), part, 512); err == nil {
		t.Fatal("expected extract seek/read failure")
	}

	src := filepath.Join(dir, "src.img")
	if err := os.WriteFile(src, make([]byte, 1024), 0o644); err != nil {
		t.Fatal(err)
	}
	if err := extractPartition(src, dir, partitionTablePart{Start: 0, Size: 1}, 512); err == nil {
		t.Fatal("expected extract destination create failure")
	}

	if err := replacePartition(src, filepath.Join(dir, "missing-root.img"), partitionTablePart{Start: 0, Size: 1}, 512, nil); err == nil {
		t.Fatal("expected missing root image failure")
	}
	root := filepath.Join(dir, "root.img")
	if err := os.WriteFile(root, make([]byte, 512), 0o644); err != nil {
		t.Fatal(err)
	}
	if err := replacePartition(filepath.Join(dir, "missing-image.img"), root, partitionTablePart{Start: 0, Size: 1}, 512, nil); err == nil {
		t.Fatal("expected destination image open failure")
	}
	if _, err := os.Stat("/dev/full"); err == nil {
		if err := replacePartition("/dev/full", root, partitionTablePart{Start: 0, Size: 1}, 512, nil); err == nil {
			t.Fatal("expected /dev/full replace failure")
		}
	}

	tmpFile := filepath.Join(dir, "tmp-file")
	if err := os.WriteFile(tmpFile, []byte("not-a-dir"), 0o644); err != nil {
		t.Fatal(err)
	}
	t.Setenv("METIS_ROOTFS_TMP_DIR", tmpFile)
	if err := writeExt4Files(src, []inject.FileSpec{{Path: "etc/node.json", Content: []byte("{}"), Mode: 0o600, RootFS: true}}); err == nil {
		t.Fatal("expected writeExt4Files temp-dir failure")
	}
}

func TestRootfsWriteAndVerifyErrorBranches(t *testing.T) {
	dir := t.TempDir()
	fsPath := filepath.Join(dir, "root.ext4")
	if err := os.WriteFile(fsPath, make([]byte, 512), 0o644); err != nil {
		t.Fatal(err)
	}

	t.Run("stage mkdir", func(t *testing.T) {
		if err := writeExt4Files(fsPath, []inject.FileSpec{{Path: "bad\x00path", Content: []byte("x"), Mode: 0o600, RootFS: true}}); err == nil {
			t.Fatal("expected stage mkdir failure")
		}
	})

	t.Run("debugfs write", func(t *testing.T) {
		scripts := fakeImageCommandDir(t, map[string]string{"debugfs": `exit 6`})
		t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))
		if err := writeExt4Files(fsPath, []inject.FileSpec{{Path: "etc/node.json", Content: []byte("{}"), Mode: 0o600, RootFS: true}}); err == nil {
			t.Fatal("expected debugfs write failure")
		}
	})

	t.Run("verify stat", func(t *testing.T) {
		scripts := fakeImageCommandDir(t, map[string]string{"debugfs": `exit 7`})
		t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))
		if err := verifyExt4File(fsPath, inject.FileSpec{Path: "etc/node.json", Content: []byte("{}"), Mode: 0o600}, t.TempDir()); err == nil {
			t.Fatal("expected debugfs stat failure")
		}
	})

	t.Run("verify mode", func(t *testing.T) {
		scripts := fakeImageCommandDir(t, map[string]string{"debugfs": `printf 'Mode:  0644\n'`})
		t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))
		if err := verifyExt4File(fsPath, inject.FileSpec{Path: "etc/node.json", Content: []byte("{}"), Mode: 0o600}, t.TempDir()); err == nil {
			t.Fatal("expected mode mismatch")
		}
	})

	t.Run("readback mkdir", func(t *testing.T) {
		scripts := fakeImageCommandDir(t, map[string]string{"debugfs": `printf 'Mode:  0600\n'`})
		t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))
		workDir := t.TempDir()
		if err := os.WriteFile(filepath.Join(workDir, "etc"), []byte("not-a-dir"), 0o644); err != nil {
			t.Fatal(err)
		}
		if err := verifyExt4File(fsPath, inject.FileSpec{Path: "etc/node.json", Content: []byte("{}"), Mode: 0o600}, workDir); err == nil {
			t.Fatal("expected readback mkdir failure")
		}
	})

	t.Run("dump failure", func(t *testing.T) {
		scripts := fakeImageCommandDir(t, map[string]string{"debugfs": `set -- ${2:-}
if [[ "${1:-}" == "stat" ]]; then printf 'Mode:  0600\n'; exit 0; fi
exit 8`})
		t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))
		if err := verifyExt4File(fsPath, inject.FileSpec{Path: "etc/node.json", Content: []byte("{}"), Mode: 0o600}, t.TempDir()); err == nil {
			t.Fatal("expected dump failure")
		}
	})

	t.Run("readback missing", func(t *testing.T) {
		scripts := fakeImageCommandDir(t, map[string]string{"debugfs": `set -- ${2:-}
if [[ "${1:-}" == "stat" ]]; then printf 'Mode:  0600\n'; exit 0; fi
exit 0`})
		t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))
		if err := verifyExt4File(fsPath, inject.FileSpec{Path: "etc/node.json", Content: []byte("{}"), Mode: 0o600}, t.TempDir()); err == nil {
			t.Fatal("expected missing readback file")
		}
	})

	t.Run("content mismatch", func(t *testing.T) {
		scripts := fakeImageCommandDir(t, map[string]string{"debugfs": `set -- ${2:-}
if [[ "${1:-}" == "stat" ]]; then printf 'Mode:  0600\n'; exit 0; fi
if [[ "${1:-}" == "dump" ]]; then printf 'wrong' > "${3:-}"; exit 0; fi
exit 0`})
		t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))
		if err := verifyExt4File(fsPath, inject.FileSpec{Path: "etc/node.json", Content: []byte("{}"), Mode: 0o600}, t.TempDir()); err == nil {
			t.Fatal("expected content mismatch")
		}
	})
}

func TestInjectBootFilesWithFakes(t *testing.T) {
	dumpDir := filepath.Join(t.TempDir(), "dump")
	t.Setenv("METIS_FAKE_UMOUNT_DUMP_DIR", dumpDir)
	scripts := fakeImageMountCommands(t, nil)
	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: "user-data", Content: []byte("#cloud-config\n"), Mode: 0o644, RootFS: false}}
	if err := InjectBootFiles(imagePath, files); err != nil {
		t.Fatalf("InjectBootFiles: %v", err)
	}
	if got := readDumpedFile(t, dumpDir, "user-data"); got != "#cloud-config\n" {
		t.Fatalf("dumped user-data = %q", got)
	}
}

func TestInjectBootFilesNoopsWithoutBootFiles(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 := InjectBootFiles(imagePath, []inject.FileSpec{{Path: "etc/metis/node.json", Content: []byte("{}"), Mode: 0o600, RootFS: true}}); err != nil {
		t.Fatalf("InjectBootFiles root-only: %v", err)
	}
}

func TestInjectBootFilesReportsMountFailures(t *testing.T) {
	scripts := fakeImageMountCommands(t, map[string]string{"mount": `exit 9`})
	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)
	}
	err := InjectBootFiles(imagePath, []inject.FileSpec{{Path: "user-data", Content: []byte("boot"), Mode: 0o644, RootFS: false}})
	if err == nil {
		t.Fatal("expected mount failure")
	}
}

func TestInjectImageWithFakes(t *testing.T) {
	rootScripts := fakeRootfsCommands(t, true)
	mountScripts := fakeImageMountCommands(t, nil)
	t.Setenv("PATH", rootScripts+string(os.PathListSeparator)+mountScripts+string(os.PathListSeparator)+os.Getenv("PATH"))

	dumpDir := filepath.Join(t.TempDir(), "dump")
	t.Setenv("METIS_FAKE_UMOUNT_DUMP_DIR", dumpDir)
	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'
`), Mode: 0o600, RootFS: true},
		{Path: "user-data", Content: []byte(`#cloud-config
`), Mode: 0o644, RootFS: false},
	}
	if err := InjectImage(imagePath, files); err != nil {
		t.Fatalf("InjectImage: %v", err)
	}
	if got := readDumpedFile(t, dumpDir, "user-data"); got != `#cloud-config
` {
		t.Fatalf("dumped user-data = %q", got)
	}
}

func readDumpedFile(t *testing.T, dumpDir, fileName string) string {
	t.Helper()
	var content string
	err := filepath.Walk(dumpDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() || filepath.Base(path) != fileName {
			return nil
		}
		data, err := os.ReadFile(path)
		if err != nil {
			return err
		}
		content = string(data)
		return filepath.SkipAll
	})
	if err != nil {
		t.Fatalf("walk dump dir: %v", err)
	}
	if content == "" {
		t.Fatalf("did not find %s under %s", fileName, dumpDir)
	}
	return content
}

func fakeImageMountCommands(t *testing.T, overrides map[string]string) string {
	t.Helper()
	dir := t.TempDir()
	scripts := map[string]string{
		"losetup": `if [[ "${1:-}" == "-Pf" && "${2:-}" == "--show" ]]; then
  printf '/dev/loop9\n'
  exit 0
fi
exit 0`,
		"mount": `mkdir -p "${2:-}"
exit 0`,
		"umount": `target="${1:-}"
dump="${METIS_FAKE_UMOUNT_DUMP_DIR:-}"
if [[ -n "$dump" && -d "$target" ]]; then
  base="$(basename "$target")"
  mkdir -p "$dump/$base"
  cp -a "$target"/. "$dump/$base/" 2>/dev/null || true
fi
exit 0`,
	}
	for name, body := range overrides {
		scripts[name] = body
	}
	for name, body := range scripts {
		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)
		}
	}
	return dir
}

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
}

func fakeImageCommandDir(t *testing.T, scripts map[string]string) string {
	t.Helper()
	dir := t.TempDir()
	for name, body := range scripts {
		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)
		}
	}
	return dir
}