package system

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"strings"
)

// ext4 mode bitmasks that debugfs's `set_inode_field ... mode` expects.
// debugfs wants the full file-type + permission word, not just the
// permission bits. Callers pass the permission portion; these constants
// OR it into the right file type.
const (
	ext4ModeRegularFile = 0o100000 // S_IFREG
	ext4ModeDirectory   = 0o040000 // S_IFDIR
)

// MkdirExt4 creates a directory inside the ext4 image, setting its
// owner/group/mode to root:root:<mode> by default or whatever the
// caller passes. Idempotent: if the directory already exists, it's
// left alone and only the metadata (uid/gid/mode) is reset to what
// was requested. Runs a single `debugfs -w` invocation so ~all the
// state transitions land in one fs-lock window.
//
// guestPath must be an absolute path inside the ext4 image (e.g.
// "/.ssh"). The function escapes the path for debugfs before sending
// it down the wire.
func MkdirExt4(ctx context.Context, runner CommandRunner, imagePath, guestPath string, mode os.FileMode, uid, gid int) error {
	escaped, err := escapeDebugfsGuestPath(guestPath)
	if err != nil {
		return err
	}
	var script bytes.Buffer
	// `mkdir` errors if the entry already exists. Tolerate that by
	// running `stat` first: on "exists" we skip the mkdir line and
	// fall through to the metadata resets, which are idempotent.
	exists, err := Ext4PathExists(ctx, runner, imagePath, guestPath)
	if err != nil {
		return err
	}
	if !exists {
		fmt.Fprintf(&script, "mkdir %s\n", escaped)
	}
	fmt.Fprintf(&script, "set_inode_field %s mode 0%o\n", escaped, ext4ModeDirectory|(uint32(mode.Perm())&0o7777))
	fmt.Fprintf(&script, "set_inode_field %s uid %d\n", escaped, uid)
	fmt.Fprintf(&script, "set_inode_field %s gid %d\n", escaped, gid)
	return debugfsScript(ctx, runner, imagePath, &script)
}

// MkdirAllExt4 creates each intermediate directory in guestPath that
// doesn't already exist, with the given mode/uid/gid. Mirrors
// os.MkdirAll's shape, not mkdir(1) -p: existing directories are left
// with their current metadata untouched (we don't reset mode/uid/gid
// on pre-existing parents, only on the final segment). Paths starting
// at "/" are allowed — the root is treated as pre-existing.
func MkdirAllExt4(ctx context.Context, runner CommandRunner, imagePath, guestPath string, mode os.FileMode, uid, gid int) error {
	if err := rejectDebugfsUnsafePath(guestPath); err != nil {
		return err
	}
	segments := strings.Split(strings.Trim(guestPath, "/"), "/")
	cur := ""
	for i, seg := range segments {
		if seg == "" {
			continue
		}
		cur = cur + "/" + seg
		exists, err := Ext4PathExists(ctx, runner, imagePath, cur)
		if err != nil {
			return err
		}
		if exists {
			continue
		}
		// Intermediate dirs inherit the requested mode/uid/gid too —
		// callers that want a different mode on parents should create
		// them explicitly. Matches the most common use (mkdir -p a
		// config tree where every hop is root-owned).
		if i < len(segments)-1 || !exists {
			if err := MkdirExt4(ctx, runner, imagePath, cur, mode, uid, gid); err != nil {
				return err
			}
		}
	}
	return nil
}

// WriteExt4FileOwned copies `data` into <imagePath>:<guestPath> and
// forces the inode's uid/gid/mode to the requested values. Unlike
// WriteExt4FileMode, this helper does NOT assume the image is a
// root-owned block device: if the image is a regular file the daemon
// user owns, every call runs without sudo. That's the common case for
// work-disk writes (vm_authsync, image_seed, runFileSync).
//
// Safety: always remove the destination first so e2cp sees a clean
// target (avoids copy-into-existing-file quirks on older e2tools).
func WriteExt4FileOwned(ctx context.Context, runner CommandRunner, imagePath, guestPath string, mode os.FileMode, uid, gid int, data []byte) error {
	tmp, err := stageDataTempfile(data, mode)
	if err != nil {
		return err
	}
	defer os.Remove(tmp)

	_, _ = extfsRun(ctx, runner, imagePath, "e2rm", imagePath+":"+guestPath)
	if _, err := extfsRun(ctx, runner, imagePath, "e2cp", tmp, imagePath+":"+guestPath); err != nil {
		return err
	}

	// Fix per-file uid/gid/mode in a debugfs batch. e2cp -O/-G exist
	// but ship inconsistently across distros; driving the inode via
	// set_inode_field matches how imagepull.ApplyOwnership has worked
	// reliably in production.
	escaped, err := escapeDebugfsGuestPath(guestPath)
	if err != nil {
		return err
	}
	var script bytes.Buffer
	fmt.Fprintf(&script, "set_inode_field %s mode 0%o\n", escaped, ext4ModeRegularFile|(uint32(mode.Perm())&0o7777))
	fmt.Fprintf(&script, "set_inode_field %s uid %d\n", escaped, uid)
	fmt.Fprintf(&script, "set_inode_field %s gid %d\n", escaped, gid)
	return debugfsScript(ctx, runner, imagePath, &script)
}

// EnsureExt4RootPerms sets the filesystem root inode (inode <2>,
// which is what `/` resolves to) to the given directory mode + owner.
// sshd's StrictModes inside the guest walks the home directory's
// ownership; the work disk is mounted at /root in the guest, so its
// root inode is /root as far as sshd is concerned. Default-safe
// value: 0755 root:root.
//
// Note on debugfs mode semantics: `set_inode_field <path> mode N`
// OVERWRITES the full i_mode word — it does NOT preserve the type
// nibble. Passing just the permission bits (e.g. 0755) would reset
// the root inode to a regular-file shape, and the next kernel mount
// would fail with "Structure needs cleaning." The constant ORed
// below restores the S_IFDIR type bits explicitly.
func EnsureExt4RootPerms(ctx context.Context, runner CommandRunner, imagePath string, mode os.FileMode, uid, gid int) error {
	fullMode := ext4ModeDirectory | (uint32(mode.Perm()) & 0o7777)
	var script bytes.Buffer
	fmt.Fprintf(&script, "set_inode_field <2> mode 0%o\n", fullMode)
	fmt.Fprintf(&script, "set_inode_field <2> uid %d\n", uid)
	fmt.Fprintf(&script, "set_inode_field <2> gid %d\n", gid)
	return debugfsScript(ctx, runner, imagePath, &script)
}

// Ext4PathExists reports whether guestPath resolves inside imagePath.
// Missing-path is NOT an error — the boolean distinguishes them.
// Uses `debugfs -R "stat <path>"` and inspects stderr for the
// standard "File not found" message e2fsprogs emits.
func Ext4PathExists(ctx context.Context, runner CommandRunner, imagePath, guestPath string) (bool, error) {
	// debugfs stat wants the path without any extra quoting beyond
	// what debugfs already does; we still reject quoting-hostile
	// chars up front.
	if err := rejectDebugfsUnsafePath(guestPath); err != nil {
		return false, err
	}
	out, err := extfsRun(ctx, runner, imagePath, "debugfs", "-R", "stat "+guestPath, imagePath)
	combined := strings.ToLower(string(out) + " " + fmt.Sprint(err))
	if strings.Contains(combined, "file not found") {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	return true, nil
}

// ReadExt4File reads guestPath from imagePath as raw bytes. Wraps the
// older ReadDebugFSText with a []byte return and the same unsafe-path
// rejection the write helpers use.
func ReadExt4File(ctx context.Context, runner CommandRunner, imagePath, guestPath string) ([]byte, error) {
	if err := rejectDebugfsUnsafePath(guestPath); err != nil {
		return nil, err
	}
	out, err := extfsRun(ctx, runner, imagePath, "debugfs", "-R", "cat "+guestPath, imagePath)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// ---- internal helpers ----

// extfsRun executes an ext4-toolkit command against imagePath,
// auto-elevating to sudo when imagePath is a block device (dm-snapshot
// targets, raw loop devices) and staying as the invoking user when
// it's a regular file (the user-owned .ext4 files under StateDir that
// this refactor targets). Tests that don't care can pass any runner
// that satisfies CommandRunner.
func extfsRun(ctx context.Context, runner CommandRunner, imagePath, name string, args ...string) ([]byte, error) {
	if needsElevation(imagePath) {
		all := append([]string{name}, args...)
		return runner.RunSudo(ctx, all...)
	}
	return runner.Run(ctx, name, args...)
}

// needsElevation returns true when imagePath is something only root
// can write to (block devices owned root:disk). For regular files
// the invoking user owns, returns false. On stat failure we err on
// the side of NOT elevating — the subsequent tool invocation will
// surface a clearer error than a bogus sudo escalation would.
func needsElevation(imagePath string) bool {
	info, err := os.Stat(imagePath)
	if err != nil {
		return false
	}
	return !info.Mode().IsRegular()
}

// debugfsScript streams a scripted batch to `debugfs -w -f -
// <imagePath>`. Requires the runner to implement StdinRunner — every
// production runner in banger does, but test doubles may not, in
// which case we fall back to one debugfs invocation per line. The
// fallback is a correctness net; production always gets the batched
// single-invocation path.
func debugfsScript(ctx context.Context, runner CommandRunner, imagePath string, script *bytes.Buffer) error {
	if script.Len() == 0 {
		return nil
	}
	stdinRunner, ok := runner.(StdinRunner)
	if ok {
		// StdinRunner's interface always runs un-elevated (it's a
		// Runner method, not RunSudo). For block devices we need sudo.
		// When elevation is required, fall through to the per-line
		// path which routes through extfsRun.
		if !needsElevation(imagePath) {
			out, err := stdinRunner.RunStdin(ctx, script, "debugfs", "-w", "-f", "-", imagePath)
			if err != nil {
				return fmt.Errorf("debugfs batch: %w: %s", err, bytes.TrimSpace(out))
			}
			return nil
		}
	}
	// Per-line fallback. Not ideal for throughput but preserves
	// semantics in tests and in the rare case we run against a
	// block device via this toolkit.
	for _, line := range strings.Split(script.String(), "\n") {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		if _, err := extfsRun(ctx, runner, imagePath, "debugfs", "-w", "-R", line, imagePath); err != nil {
			return fmt.Errorf("debugfs %q: %w", line, err)
		}
	}
	return nil
}

// escapeDebugfsGuestPath produces a debugfs-safe rendition of the
// guest path. debugfs tokenises on whitespace by default; paths with
// spaces must be double-quoted. Paths containing the double-quote
// itself, backslashes, or newlines are rejected outright — quoting
// those reliably in debugfs's hand-rolled parser is lore we don't
// want to inherit.
func escapeDebugfsGuestPath(guestPath string) (string, error) {
	if err := rejectDebugfsUnsafePath(guestPath); err != nil {
		return "", err
	}
	if strings.ContainsAny(guestPath, " \t") {
		return `"` + guestPath + `"`, nil
	}
	return guestPath, nil
}

func rejectDebugfsUnsafePath(guestPath string) error {
	if guestPath == "" {
		return fmt.Errorf("guest path is required")
	}
	if !strings.HasPrefix(guestPath, "/") {
		return fmt.Errorf("guest path %q must be absolute", guestPath)
	}
	if strings.ContainsAny(guestPath, "\"\\\n\r") {
		return fmt.Errorf("guest path %q contains characters debugfs cannot safely encode", guestPath)
	}
	return nil
}

func stageDataTempfile(data []byte, mode os.FileMode) (string, error) {
	tmp, err := os.CreateTemp("", "banger-ext4-*")
	if err != nil {
		return "", err
	}
	path := tmp.Name()
	if _, err := tmp.Write(data); err != nil {
		_ = tmp.Close()
		_ = os.Remove(path)
		return "", err
	}
	if err := tmp.Close(); err != nil {
		_ = os.Remove(path)
		return "", err
	}
	if err := os.Chmod(path, mode.Perm()); err != nil {
		_ = os.Remove(path)
		return "", err
	}
	return path, nil
}

// RdumpExt4Dir shells out to `debugfs -R "rdump <src> <dst>" image`
// to spill a tree from the ext4 image into a host directory. Used by
// ensureWorkDisk's no-seed path to extract /root from the base rootfs
// without mounting. Content is preserved; per-entry metadata (uid,
// gid, mode) is captured via a subsequent stat walk inside debugfs.
// Returns the destination directory (same as dst on success).
func RdumpExt4Dir(ctx context.Context, runner CommandRunner, imagePath, srcPath, dstDir string) error {
	if err := rejectDebugfsUnsafePath(srcPath); err != nil {
		return err
	}
	if err := os.MkdirAll(dstDir, 0o755); err != nil {
		return err
	}
	_, err := extfsRun(ctx, runner, imagePath, "debugfs", "-R", "rdump "+srcPath+" "+dstDir, imagePath)
	return err
}