package imagepull

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"banger/internal/guestnet"
	"banger/internal/system"
	"banger/internal/vsockagent"
)

// GuestAgentAssets bundles everything the guest side of banger needs in a
// rootfs that doesn't already have it. Callers (the daemon's PullImage)
// resolve the vsock-agent binary path via paths.CompanionBinaryPath and
// hand it in; the rest comes from the respective asset packages.
type GuestAgentAssets struct {
	VsockAgentBin string // absolute path on the host, copied verbatim
}

// InjectGuestAgents writes banger's guest-side assets (vsock agent
// binary + systemd unit, network bootstrap script + unit, vsock modules-
// load config, symlinks that enable the units at boot) into ext4File.
// All entries land with uid=0, gid=0 and appropriate modes.
//
// Runs in one debugfs -w invocation: dirs, files, sif (uid/gid/mode),
// and symlinks all in one scripted batch. No sudo required because the
// ext4 is owned by the runner.
func InjectGuestAgents(ctx context.Context, runner system.CommandRunner, ext4File string, assets GuestAgentAssets) error {
	if assets.VsockAgentBin == "" {
		return fmt.Errorf("vsock-agent binary path is required")
	}
	if _, err := os.Stat(assets.VsockAgentBin); err != nil {
		return fmt.Errorf("vsock-agent binary %q missing: %w", assets.VsockAgentBin, err)
	}

	// Stage content blobs as temp files so debugfs `write` can pick
	// them up. All other commands (mkdir/sif/symlink) are inline.
	stage, err := os.MkdirTemp("", "banger-inject-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(stage)

	steps := []injectFile{
		{
			hostSrc:   assets.VsockAgentBin,
			guestPath: vsockagent.GuestInstallPath, // /usr/local/bin/banger-vsock-agent
			mode:      0o755,
		},
		{
			content:   []byte(guestnet.BootstrapScript()),
			guestPath: guestnet.GuestScriptPath, // /usr/local/libexec/banger-network-bootstrap
			mode:      0o755,
		},
		{
			content:   []byte(guestnet.SystemdServiceUnit()),
			guestPath: "/etc/systemd/system/" + guestnet.SystemdServiceName, // banger-network.service
			mode:      0o644,
		},
		{
			content:   []byte(vsockagent.ServiceUnit()),
			guestPath: "/etc/systemd/system/" + vsockagent.ServiceName, // banger-vsock-agent.service
			mode:      0o644,
		},
		{
			content:   []byte(vsockagent.ModulesLoadConfig()),
			guestPath: "/etc/modules-load.d/banger-vsock.conf",
			mode:      0o644,
		},
		{
			content:   []byte(FirstBootScript()),
			guestPath: FirstBootScriptPath, // /usr/local/libexec/banger-first-boot
			mode:      0o755,
		},
		{
			content:   []byte(FirstBootUnit()),
			guestPath: "/etc/systemd/system/" + FirstBootUnitName,
			mode:      0o644,
		},
		{
			content:   nil, // empty marker file — its existence triggers the service
			guestPath: FirstBootMarkerPath,
			mode:      0o644,
		},
	}

	// Resolve content-backed steps to on-disk temp files.
	for i := range steps {
		if steps[i].hostSrc != "" {
			continue
		}
		tmp := filepath.Join(stage, fmt.Sprintf("blob-%d", i))
		if err := os.WriteFile(tmp, steps[i].content, 0o644); err != nil {
			return err
		}
		steps[i].hostSrc = tmp
	}

	symlinks := []injectSymlink{
		{
			target: "/etc/systemd/system/" + guestnet.SystemdServiceName,
			link:   "/etc/systemd/system/multi-user.target.wants/" + guestnet.SystemdServiceName,
		},
		{
			target: "/etc/systemd/system/" + vsockagent.ServiceName,
			link:   "/etc/systemd/system/multi-user.target.wants/" + vsockagent.ServiceName,
		},
		{
			target: "/etc/systemd/system/" + FirstBootUnitName,
			link:   "/etc/systemd/system/multi-user.target.wants/" + FirstBootUnitName,
		},
	}

	script := buildInjectScript(steps, symlinks)

	stdinRunner, ok := runner.(system.StdinRunner)
	if !ok {
		return fmt.Errorf("inject requires a runner that supports stdin (got %T)", runner)
	}
	out, err := stdinRunner.RunStdin(ctx, script, "debugfs", "-w", "-f", "-", ext4File)
	if err != nil {
		return fmt.Errorf("debugfs inject: %w: %s", err, string(out))
	}
	// Scan output for hard errors — debugfs keeps going past errors
	// with -f, so we need to look at stdout/stderr-as-stdout for bad
	// signs. mkdir errors on already-present dirs are expected; we
	// ignore "File exists" and "Is a directory". Other errors bubble.
	if bad := scanInjectOutput(out); bad != "" {
		return fmt.Errorf("debugfs inject: %s", bad)
	}
	return nil
}

type injectFile struct {
	content   []byte
	hostSrc   string // set by InjectGuestAgents after staging
	guestPath string
	mode      uint32 // perm bits; type bits added by buildInjectScript
}

type injectSymlink struct {
	target string
	link   string
}

// buildInjectScript emits the debugfs command stream.
func buildInjectScript(files []injectFile, symlinks []injectSymlink) *bytes.Buffer {
	var buf bytes.Buffer

	// Create every ancestor directory of every file/symlink path. mkdir
	// on an already-existing dir is benign (debugfs continues past the
	// error), but we prune duplicates to keep the script clean.
	dirs := collectAncestors(files, symlinks)
	for _, d := range dirs {
		fmt.Fprintf(&buf, "mkdir %s\n", d)
	}

	// Write each file content.
	for _, f := range files {
		fmt.Fprintf(&buf, "write %s %s\n", f.hostSrc, f.guestPath)
	}

	// Fix ownership + mode on every written file (uid=0, gid=0).
	for _, f := range files {
		fmt.Fprintf(&buf, "set_inode_field %s uid 0\n", f.guestPath)
		fmt.Fprintf(&buf, "set_inode_field %s gid 0\n", f.guestPath)
		fmt.Fprintf(&buf, "set_inode_field %s mode 0%o\n", f.guestPath, 0o100000|f.mode)
	}

	// Fix dir ownership. Don't touch modes — mkdir's default 0755 is fine.
	for _, d := range dirs {
		fmt.Fprintf(&buf, "set_inode_field %s uid 0\n", d)
		fmt.Fprintf(&buf, "set_inode_field %s gid 0\n", d)
	}

	// Finally, create the enable-at-boot symlinks.
	for _, s := range symlinks {
		fmt.Fprintf(&buf, "symlink %s %s\n", s.link, s.target)
	}

	return &buf
}

// collectAncestors walks every file + symlink path and returns the unique
// set of parent directories, sorted shallowest first so mkdir ordering
// is valid.
func collectAncestors(files []injectFile, symlinks []injectSymlink) []string {
	set := map[string]struct{}{}
	add := func(p string) {
		dir := filepath.Dir(p)
		for dir != "" && dir != "/" {
			set[dir] = struct{}{}
			dir = filepath.Dir(dir)
		}
	}
	for _, f := range files {
		add(f.guestPath)
	}
	for _, s := range symlinks {
		add(s.link)
	}
	out := make([]string, 0, len(set))
	for d := range set {
		out = append(out, d)
	}
	// Shallow-first by depth, then lexicographic.
	sort.Slice(out, func(i, j int) bool {
		di := strings.Count(out[i], "/")
		dj := strings.Count(out[j], "/")
		if di != dj {
			return di < dj
		}
		return out[i] < out[j]
	})
	return out
}

// scanInjectOutput returns a non-empty string if debugfs reported an
// error that's not a benign "File exists" from mkdir on an already-
// present directory. Debugfs emits errors on stderr AND stdout (which
// we capture together); we look for known failure signatures.
func scanInjectOutput(out []byte) string {
	lines := strings.Split(string(out), "\n")
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		// Benign: mkdir on existing dir.
		if strings.Contains(line, "File exists") {
			continue
		}
		// Failure signatures we care about.
		if strings.Contains(line, "error writing file") ||
			strings.Contains(line, "couldn't find") ||
			strings.Contains(line, "No such file") ||
			strings.Contains(line, "Unrecognized command") ||
			strings.Contains(line, "symlink:") {
			return line
		}
	}
	return ""
}
