package cli

import (
	"archive/tar"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"strings"

	"banger/internal/config"
	"banger/internal/hostnat"
	"banger/internal/imagecat"
	"banger/internal/imagepull"
	"banger/internal/model"
	"banger/internal/paths"
	"banger/internal/system"

	"github.com/klauspost/compress/zstd"
	"github.com/spf13/cobra"
)

func (d *deps) newInternalCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:    "internal",
		Hidden: true,
		RunE:   helpNoArgs,
	}
	cmd.AddCommand(
		newInternalNATCommand(),
		newInternalWorkSeedCommand(),
		newInternalSSHKeyPathCommand(),
		newInternalFirecrackerPathCommand(),
		newInternalVSockAgentPathCommand(),
		newInternalMakeBundleCommand(),
	)
	return cmd
}

func newInternalSSHKeyPathCommand() *cobra.Command {
	return &cobra.Command{
		Use:    "ssh-key-path",
		Hidden: true,
		Args:   noArgsUsage("usage: banger internal ssh-key-path"),
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, err := paths.Resolve()
			if err != nil {
				return err
			}
			cfg, err := config.Load(layout)
			if err != nil {
				return err
			}
			_, err = fmt.Fprintln(cmd.OutOrStdout(), cfg.SSHKeyPath)
			return err
		},
	}
}

func newInternalFirecrackerPathCommand() *cobra.Command {
	return &cobra.Command{
		Use:    "firecracker-path",
		Hidden: true,
		Args:   noArgsUsage("usage: banger internal firecracker-path"),
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, err := paths.Resolve()
			if err != nil {
				return err
			}
			cfg, err := config.Load(layout)
			if err != nil {
				return err
			}
			if strings.TrimSpace(cfg.FirecrackerBin) == "" {
				return errors.New("firecracker binary not configured; install firecracker or set firecracker_bin")
			}
			_, err = fmt.Fprintln(cmd.OutOrStdout(), cfg.FirecrackerBin)
			return err
		},
	}
}

func newInternalVSockAgentPathCommand() *cobra.Command {
	return &cobra.Command{
		Use:    "vsock-agent-path",
		Hidden: true,
		Args:   noArgsUsage("usage: banger internal vsock-agent-path"),
		RunE: func(cmd *cobra.Command, args []string) error {
			path, err := paths.CompanionBinaryPath("banger-vsock-agent")
			if err != nil {
				return err
			}
			_, err = fmt.Fprintln(cmd.OutOrStdout(), path)
			return err
		},
	}
}

func newInternalMakeBundleCommand() *cobra.Command {
	var (
		rootfsTarPath string
		name          string
		distro        string
		arch          string
		kernelRef     string
		description   string
		sizeSpec      string
		outPath       string
	)
	cmd := &cobra.Command{
		Use:    "make-bundle",
		Hidden: true,
		Short:  "Build a banger image bundle (.tar.zst) from a flat rootfs tar",
		Args:   noArgsUsage("usage: banger internal make-bundle --rootfs-tar <file|-> --name <n> --out <bundle.tar.zst>"),
		RunE: func(cmd *cobra.Command, args []string) error {
			return runInternalMakeBundle(cmd, internalMakeBundleOpts{
				rootfsTarPath: rootfsTarPath,
				name:          name,
				distro:        distro,
				arch:          arch,
				kernelRef:     kernelRef,
				description:   description,
				sizeSpec:      sizeSpec,
				outPath:       outPath,
			})
		},
	}
	cmd.Flags().StringVar(&rootfsTarPath, "rootfs-tar", "", "flat rootfs tar file, or '-' for stdin")
	cmd.Flags().StringVar(&name, "name", "", "bundle name (filesystem-safe identifier)")
	cmd.Flags().StringVar(&distro, "distro", "", "distro label (e.g. debian)")
	cmd.Flags().StringVar(&arch, "arch", "x86_64", "architecture label")
	cmd.Flags().StringVar(&kernelRef, "kernel-ref", "", "kernelcat entry name this image pairs with")
	cmd.Flags().StringVar(&description, "description", "", "short description")
	cmd.Flags().StringVar(&sizeSpec, "size", "", "rootfs ext4 size (e.g. 4G); defaults to tree size + 25%")
	cmd.Flags().StringVar(&outPath, "out", "", "output bundle path (.tar.zst)")
	return cmd
}

type internalMakeBundleOpts struct {
	rootfsTarPath string
	name          string
	distro        string
	arch          string
	kernelRef     string
	description   string
	sizeSpec      string
	outPath       string
}

func runInternalMakeBundle(cmd *cobra.Command, opts internalMakeBundleOpts) error {
	if err := imagecat.ValidateName(opts.name); err != nil {
		return err
	}
	if strings.TrimSpace(opts.rootfsTarPath) == "" {
		return errors.New("--rootfs-tar is required")
	}
	if strings.TrimSpace(opts.outPath) == "" {
		return errors.New("--out is required")
	}
	if strings.TrimSpace(opts.arch) == "" {
		opts.arch = "x86_64"
	}

	var sizeBytes int64
	if s := strings.TrimSpace(opts.sizeSpec); s != "" {
		n, err := model.ParseSize(s)
		if err != nil {
			return fmt.Errorf("parse --size: %w", err)
		}
		sizeBytes = n
	}

	ctx := cmd.Context()
	stagingRoot, err := os.MkdirTemp("", "banger-mkbundle-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(stagingRoot)
	rootfsTree := filepath.Join(stagingRoot, "rootfs")
	if err := os.MkdirAll(rootfsTree, 0o755); err != nil {
		return err
	}

	var tarReader io.Reader
	if opts.rootfsTarPath == "-" {
		tarReader = cmd.InOrStdin()
	} else {
		f, err := os.Open(opts.rootfsTarPath)
		if err != nil {
			return fmt.Errorf("open rootfs tar: %w", err)
		}
		defer f.Close()
		tarReader = f
	}

	fmt.Fprintln(cmd.ErrOrStderr(), "[make-bundle] extracting rootfs")
	meta, err := imagepull.FlattenTar(ctx, tarReader, rootfsTree)
	if err != nil {
		return fmt.Errorf("flatten rootfs: %w", err)
	}

	// docker create drops /.dockerenv (and containerd drops
	// /run/.containerenv) into the container's writable layer, so
	// `docker export` includes them in the tar. systemd-detect-virt
	// reads those files and flags the boot as virtualization=docker,
	// which disables udev device-unit activation (including the work-
	// disk dev-vdb.device) and leaves systemd waiting forever. Strip
	// them before building the ext4.
	for _, marker := range []string{".dockerenv", "run/.containerenv"} {
		path := filepath.Join(rootfsTree, marker)
		if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
			return fmt.Errorf("strip %s: %w", marker, err)
		}
		delete(meta.Entries, marker)
	}

	if sizeBytes <= 0 {
		treeSize, err := dirSize(rootfsTree)
		if err != nil {
			return fmt.Errorf("size rootfs tree: %w", err)
		}
		// +50% headroom for ext4 overhead (inode tables, block-group
		// descriptors, journal, 5% reserved margin).
		sizeBytes = treeSize + treeSize/2
		if sizeBytes < imagepull.MinExt4Size {
			sizeBytes = imagepull.MinExt4Size
		}
	}

	ext4Path := filepath.Join(stagingRoot, imagecat.RootfsFilename)
	runner := system.NewRunner()
	fmt.Fprintf(cmd.ErrOrStderr(), "[make-bundle] building rootfs.ext4 (%d bytes)\n", sizeBytes)
	if err := imagepull.BuildExt4(ctx, runner, rootfsTree, ext4Path, sizeBytes); err != nil {
		return fmt.Errorf("build ext4: %w", err)
	}
	fmt.Fprintln(cmd.ErrOrStderr(), "[make-bundle] applying ownership fixup")
	if err := imagepull.ApplyOwnership(ctx, runner, ext4Path, meta); err != nil {
		return fmt.Errorf("apply ownership: %w", err)
	}
	fmt.Fprintln(cmd.ErrOrStderr(), "[make-bundle] injecting guest agents")
	vsockBin, err := paths.CompanionBinaryPath("banger-vsock-agent")
	if err != nil {
		return fmt.Errorf("locate vsock agent: %w", err)
	}
	if err := imagepull.InjectGuestAgents(ctx, runner, ext4Path, imagepull.GuestAgentAssets{VsockAgentBin: vsockBin}); err != nil {
		return fmt.Errorf("inject guest agents: %w", err)
	}

	manifest := imagecat.Manifest{
		Name:        opts.name,
		Distro:      strings.TrimSpace(opts.distro),
		Arch:        opts.arch,
		KernelRef:   strings.TrimSpace(opts.kernelRef),
		Description: strings.TrimSpace(opts.description),
	}
	manifestPath := filepath.Join(stagingRoot, imagecat.ManifestFilename)
	manifestData, err := json.MarshalIndent(manifest, "", "  ")
	if err != nil {
		return err
	}
	if err := os.WriteFile(manifestPath, append(manifestData, '\n'), 0o644); err != nil {
		return err
	}

	fmt.Fprintln(cmd.ErrOrStderr(), "[make-bundle] packaging bundle")
	if err := writeBundleTarZst(opts.outPath, ext4Path, manifestPath); err != nil {
		return fmt.Errorf("write bundle: %w", err)
	}

	sum, err := sha256HexFile(opts.outPath)
	if err != nil {
		return err
	}
	stat, err := os.Stat(opts.outPath)
	if err != nil {
		return err
	}
	fmt.Fprintf(cmd.OutOrStdout(), "bundle: %s\nsha256: %s\nsize:   %d\n", opts.outPath, sum, stat.Size())
	return nil
}

func dirSize(root string) (int64, error) {
	var total int64
	err := filepath.WalkDir(root, func(_ string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if !d.Type().IsRegular() {
			return nil
		}
		info, err := d.Info()
		if err != nil {
			return err
		}
		total += info.Size()
		return nil
	})
	return total, err
}

func writeBundleTarZst(outPath, rootfsPath, manifestPath string) error {
	if err := os.MkdirAll(filepath.Dir(outPath), 0o755); err != nil {
		return err
	}
	out, err := os.OpenFile(outPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644)
	if err != nil {
		return err
	}
	defer out.Close()
	zw, err := zstd.NewWriter(out, zstd.WithEncoderLevel(zstd.SpeedBestCompression))
	if err != nil {
		return err
	}
	tw := tar.NewWriter(zw)
	for _, src := range []struct{ path, name string }{
		{rootfsPath, imagecat.RootfsFilename},
		{manifestPath, imagecat.ManifestFilename},
	} {
		if err := writeBundleFile(tw, src.path, src.name); err != nil {
			_ = tw.Close()
			_ = zw.Close()
			return err
		}
	}
	if err := tw.Close(); err != nil {
		_ = zw.Close()
		return err
	}
	if err := zw.Close(); err != nil {
		return err
	}
	return out.Close()
}

func writeBundleFile(tw *tar.Writer, src, name string) error {
	f, err := os.Open(src)
	if err != nil {
		return err
	}
	defer f.Close()
	fi, err := f.Stat()
	if err != nil {
		return err
	}
	if err := tw.WriteHeader(&tar.Header{
		Name:     name,
		Size:     fi.Size(),
		Mode:     0o644,
		Typeflag: tar.TypeReg,
		ModTime:  fi.ModTime(),
	}); err != nil {
		return err
	}
	_, err = io.Copy(tw, f)
	return err
}

func sha256HexFile(path string) (string, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer f.Close()
	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}
	return hex.EncodeToString(h.Sum(nil)), nil
}

func newInternalWorkSeedCommand() *cobra.Command {
	var rootfsPath string
	var outPath string
	cmd := &cobra.Command{
		Use:    "work-seed",
		Hidden: true,
		Args:   noArgsUsage("usage: banger internal work-seed --rootfs <path> [--out <path>]"),
		RunE: func(cmd *cobra.Command, args []string) error {
			rootfsPath = strings.TrimSpace(rootfsPath)
			outPath = strings.TrimSpace(outPath)
			if rootfsPath == "" {
				return errors.New("rootfs path is required")
			}
			if outPath == "" {
				outPath = system.WorkSeedPath(rootfsPath)
			}
			if err := system.EnsureSudo(cmd.Context()); err != nil {
				return err
			}
			return system.BuildWorkSeedImage(cmd.Context(), system.NewRunner(), rootfsPath, outPath)
		},
	}
	cmd.Flags().StringVar(&rootfsPath, "rootfs", "", "rootfs image path")
	cmd.Flags().StringVar(&outPath, "out", "", "output work-seed image path")
	return cmd
}

func newInternalNATCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:    "nat",
		Hidden: true,
		RunE:   helpNoArgs,
	}
	cmd.AddCommand(
		newInternalNATActionCommand("up", true),
		newInternalNATActionCommand("down", false),
	)
	return cmd
}

func newInternalNATActionCommand(use string, enable bool) *cobra.Command {
	var guestIP string
	var tapDevice string
	cmd := &cobra.Command{
		Use:    use,
		Hidden: true,
		Args:   noArgsUsage("usage: banger internal nat " + use + " --guest-ip <ip> --tap <tap-device>"),
		RunE: func(cmd *cobra.Command, args []string) error {
			guestIP = strings.TrimSpace(guestIP)
			tapDevice = strings.TrimSpace(tapDevice)
			if guestIP == "" {
				return errors.New("guest IP is required")
			}
			if tapDevice == "" {
				return errors.New("tap device is required")
			}
			if err := system.EnsureSudo(cmd.Context()); err != nil {
				return err
			}
			return hostnat.Ensure(cmd.Context(), system.NewRunner(), guestIP, tapDevice, enable)
		},
	}
	cmd.Flags().StringVar(&guestIP, "guest-ip", "", "guest IPv4 address")
	cmd.Flags().StringVar(&tapDevice, "tap", "", "tap device name")
	return cmd
}