package system

import (
	"bytes"
	"context"
	"errors"
	"io"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

// stdinFuncRunner is funcRunner extended with a RunStdin hook so we
// can assert the exact debugfs batch script that callers stream in.
type stdinFuncRunner struct {
	funcRunner
	runStdin func(ctx context.Context, stdin io.Reader, name string, args ...string) ([]byte, error)
}

func (r stdinFuncRunner) RunStdin(ctx context.Context, stdin io.Reader, name string, args ...string) ([]byte, error) {
	if r.runStdin == nil {
		return nil, errors.New("unexpected RunStdin call")
	}
	return r.runStdin(ctx, stdin, name, args...)
}

// userOwnedImage writes a zero-length regular file at a tempdir and
// returns its path. Regular files trigger extfsRun's non-sudo branch,
// which is the whole point of the new toolkit.
func userOwnedImage(t *testing.T) string {
	t.Helper()
	path := filepath.Join(t.TempDir(), "work.ext4")
	if err := os.WriteFile(path, []byte{}, 0o644); err != nil {
		t.Fatalf("write image: %v", err)
	}
	return path
}

func TestExt4PathExists(t *testing.T) {
	image := userOwnedImage(t)

	t.Run("path found", func(t *testing.T) {
		r := funcRunner{
			run: func(_ context.Context, name string, args ...string) ([]byte, error) {
				if name != "debugfs" {
					t.Fatalf("name = %q, want debugfs", name)
				}
				want := []string{"-R", "stat /root/.ssh", image}
				for i := range want {
					if args[i] != want[i] {
						t.Fatalf("args[%d] = %q, want %q (full %v)", i, args[i], want[i], args)
					}
				}
				return []byte("Inode: 12 Type: directory"), nil
			},
		}
		ok, err := Ext4PathExists(context.Background(), r, image, "/root/.ssh")
		if err != nil {
			t.Fatalf("Ext4PathExists: %v", err)
		}
		if !ok {
			t.Fatal("expected exists = true")
		}
	})

	t.Run("path missing", func(t *testing.T) {
		r := funcRunner{
			run: func(context.Context, string, ...string) ([]byte, error) {
				// debugfs prints the "File not found" message to stdout
				// on lookup miss. No exit error (debugfs exits 0 for
				// soft misses on `stat`).
				return []byte("stat: File not found by ext2_lookup while starting pathname"), nil
			},
		}
		ok, err := Ext4PathExists(context.Background(), r, image, "/root/.ssh")
		if err != nil {
			t.Fatalf("Ext4PathExists: %v", err)
		}
		if ok {
			t.Fatal("expected exists = false")
		}
	})

	t.Run("rejects hostile path", func(t *testing.T) {
		r := funcRunner{}
		if _, err := Ext4PathExists(context.Background(), r, image, `/evil"path`); err == nil {
			t.Fatal("expected rejection for path containing double-quote")
		}
	})
}

func TestReadExt4File(t *testing.T) {
	image := userOwnedImage(t)
	r := funcRunner{
		run: func(_ context.Context, name string, args ...string) ([]byte, error) {
			if name != "debugfs" {
				t.Fatalf("name = %q, want debugfs", name)
			}
			if args[0] != "-R" || args[1] != "cat /etc/fstab" {
				t.Fatalf("args = %v, want -R \"cat /etc/fstab\" ...", args)
			}
			return []byte("tmpfs /tmp tmpfs defaults 0 0\n"), nil
		},
	}
	got, err := ReadExt4File(context.Background(), r, image, "/etc/fstab")
	if err != nil {
		t.Fatalf("ReadExt4File: %v", err)
	}
	if !bytes.Contains(got, []byte("tmpfs /tmp")) {
		t.Fatalf("got = %q, want contains tmpfs line", got)
	}
}

func TestMkdirExt4_BatchesStatMkdirAndMetadata(t *testing.T) {
	image := userOwnedImage(t)

	var capturedScript string
	r := stdinFuncRunner{
		funcRunner: funcRunner{
			run: func(_ context.Context, name string, args ...string) ([]byte, error) {
				// The only non-stdin call should be the existence check.
				if name == "debugfs" && len(args) >= 2 && args[0] == "-R" && strings.HasPrefix(args[1], "stat ") {
					return []byte("stat: File not found"), nil
				}
				t.Fatalf("unexpected Run(%q, %v)", name, args)
				return nil, nil
			},
		},
		runStdin: func(_ context.Context, stdin io.Reader, name string, args ...string) ([]byte, error) {
			if name != "debugfs" {
				t.Fatalf("stdin runner name = %q, want debugfs", name)
			}
			want := []string{"-w", "-f", "-", image}
			for i, w := range want {
				if args[i] != w {
					t.Fatalf("stdin args[%d] = %q, want %q", i, args[i], w)
				}
			}
			b, _ := io.ReadAll(stdin)
			capturedScript = string(b)
			return nil, nil
		},
	}

	if err := MkdirExt4(context.Background(), r, image, "/.ssh", 0o700, 0, 0); err != nil {
		t.Fatalf("MkdirExt4: %v", err)
	}

	// mkdir line must be present (path didn't exist).
	if !strings.Contains(capturedScript, "mkdir /.ssh") {
		t.Fatalf("script missing mkdir line:\n%s", capturedScript)
	}
	// Mode must include the directory file-type nibble (040000 | 0700 = 040700).
	if !strings.Contains(capturedScript, "set_inode_field /.ssh mode 040700") {
		t.Fatalf("script missing mode line with S_IFDIR+0700:\n%s", capturedScript)
	}
	if !strings.Contains(capturedScript, "set_inode_field /.ssh uid 0") {
		t.Fatalf("script missing uid line:\n%s", capturedScript)
	}
	if !strings.Contains(capturedScript, "set_inode_field /.ssh gid 0") {
		t.Fatalf("script missing gid line:\n%s", capturedScript)
	}
}

func TestMkdirExt4_SkipsMkdirWhenDirectoryExists(t *testing.T) {
	image := userOwnedImage(t)

	var capturedScript string
	r := stdinFuncRunner{
		funcRunner: funcRunner{
			run: func(_ context.Context, name string, args ...string) ([]byte, error) {
				// First call: existence probe. Return success.
				if name == "debugfs" && args[0] == "-R" && strings.HasPrefix(args[1], "stat ") {
					return []byte("Inode: 12 Type: directory Mode: 0700"), nil
				}
				t.Fatalf("unexpected Run(%q, %v)", name, args)
				return nil, nil
			},
		},
		runStdin: func(_ context.Context, stdin io.Reader, _ string, _ ...string) ([]byte, error) {
			b, _ := io.ReadAll(stdin)
			capturedScript = string(b)
			return nil, nil
		},
	}

	if err := MkdirExt4(context.Background(), r, image, "/.ssh", 0o700, 0, 0); err != nil {
		t.Fatalf("MkdirExt4: %v", err)
	}

	// Directory existed — no mkdir line, but metadata lines still fire.
	if strings.Contains(capturedScript, "mkdir ") {
		t.Fatalf("script should not contain mkdir for pre-existing path:\n%s", capturedScript)
	}
	if !strings.Contains(capturedScript, "set_inode_field /.ssh mode") {
		t.Fatalf("script missing metadata reset for pre-existing dir:\n%s", capturedScript)
	}
}

func TestWriteExt4FileOwned_StagesTempfileAndBatchesOwnership(t *testing.T) {
	image := userOwnedImage(t)

	var observedTemp string
	var capturedScript string
	r := stdinFuncRunner{
		funcRunner: funcRunner{
			run: func(_ context.Context, name string, args ...string) ([]byte, error) {
				switch name {
				case "e2rm":
					// First non-stdin call — best-effort, we don't
					// verify the target since e2rm on a missing path
					// returns a visible error but the caller ignores it.
					return nil, nil
				case "e2cp":
					if len(args) != 2 {
						t.Fatalf("e2cp args = %v, want 2 (src, dst)", args)
					}
					observedTemp = args[0]
					// Assert the dst uses the image:path form.
					if args[1] != image+":/root/.ssh/authorized_keys" {
						t.Fatalf("e2cp dst = %q, want image:path", args[1])
					}
					// Assert the temp file was populated with our data
					// BEFORE e2cp was called.
					data, err := os.ReadFile(args[0])
					if err != nil {
						t.Fatalf("temp missing at e2cp time: %v", err)
					}
					if !bytes.Equal(data, []byte("managed-key\n")) {
						t.Fatalf("temp contents = %q, want managed-key", data)
					}
					return nil, nil
				}
				t.Fatalf("unexpected Run(%q, %v)", name, args)
				return nil, nil
			},
		},
		runStdin: func(_ context.Context, stdin io.Reader, _ string, _ ...string) ([]byte, error) {
			b, _ := io.ReadAll(stdin)
			capturedScript = string(b)
			return nil, nil
		},
	}

	err := WriteExt4FileOwned(
		context.Background(), r, image,
		"/root/.ssh/authorized_keys",
		0o600, 0, 0,
		[]byte("managed-key\n"),
	)
	if err != nil {
		t.Fatalf("WriteExt4FileOwned: %v", err)
	}

	// Temp cleanup ran — we saved observedTemp while it still existed;
	// by now it should be gone.
	if observedTemp == "" {
		t.Fatal("e2cp source path was never captured")
	}
	if _, err := os.Stat(observedTemp); !os.IsNotExist(err) {
		t.Fatalf("temp file not cleaned up: stat err = %v", err)
	}

	// Mode line must bake in S_IFREG (0100000) + 0600 = 0100600.
	if !strings.Contains(capturedScript, "set_inode_field /root/.ssh/authorized_keys mode 0100600") {
		t.Fatalf("script missing regular-file mode line:\n%s", capturedScript)
	}
}

func TestEnsureExt4RootPerms_UsesRootInodeLiteral(t *testing.T) {
	image := userOwnedImage(t)

	var capturedScript string
	r := stdinFuncRunner{
		funcRunner: funcRunner{},
		runStdin: func(_ context.Context, stdin io.Reader, _ string, _ ...string) ([]byte, error) {
			b, _ := io.ReadAll(stdin)
			capturedScript = string(b)
			return nil, nil
		},
	}

	if err := EnsureExt4RootPerms(context.Background(), r, image, 0o755, 0, 0); err != nil {
		t.Fatalf("EnsureExt4RootPerms: %v", err)
	}

	// Must address inode 2 — the ext4 root directory — with the
	// FULL i_mode word (S_IFDIR | 0755 = 040755). debugfs's
	// set_inode_field doesn't preserve the type nibble, so passing
	// just the permission bits (0755) would reset the root inode
	// to regular-file shape and break the next kernel mount.
	if !strings.Contains(capturedScript, "set_inode_field <2> mode 040755") {
		t.Fatalf("script missing root-inode mode line with S_IFDIR+0755:\n%s", capturedScript)
	}
	if !strings.Contains(capturedScript, "set_inode_field <2> uid 0") {
		t.Fatalf("script missing root-inode uid line:\n%s", capturedScript)
	}
}

func TestRejectDebugfsUnsafePath(t *testing.T) {
	for _, tc := range []struct {
		name    string
		path    string
		wantErr bool
	}{
		{"empty", "", true},
		{"relative", "relative/path", true},
		{"absolute plain", "/ok", false},
		{"absolute with space", "/ok path", false},
		{"contains double-quote", `/a"b`, true},
		{"contains backslash", `/a\b`, true},
		{"contains newline", "/a\nb", true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			err := rejectDebugfsUnsafePath(tc.path)
			if (err != nil) != tc.wantErr {
				t.Fatalf("rejectDebugfsUnsafePath(%q) err = %v, wantErr = %v", tc.path, err, tc.wantErr)
			}
		})
	}
}