package store

import (
	"archive/tar"
	"bufio"
	"bytes"
	"compress/gzip"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
)

// applyLayer extracts one OCI/docker layer tar onto rootfs, honoring the
// whiteout convention: `.wh.NAME` removes NAME from the lower layers,
// `.wh..wh..opq` empties the directory it sits in. Whiteouts are applied
// before the layer's own entries, so an opaque directory refilled by the
// same layer ends up with exactly that layer's content.
//
// Every path is resolved inside rootfs the way the sandbox will see it:
// a lower layer's symlink (/bin -> usr/bin) is followed within rootfs,
// never on the host, and anything escaping rootfs is refused. Ownership
// is dropped, device nodes are skipped; the sandbox maps uid 0 to the
// caller anyway.
func applyLayer(rootfs, layerPath string) error {
	if err := forEachEntry(layerPath, func(hdr *tar.Header, _ io.Reader) error {
		return applyWhiteout(rootfs, hdr.Name)
	}); err != nil {
		return err
	}
	return forEachEntry(layerPath, func(hdr *tar.Header, body io.Reader) error {
		if strings.HasPrefix(path.Base(hdr.Name), whiteoutPrefix) {
			return nil
		}
		return extractEntry(rootfs, hdr, body)
	})
}

const (
	whiteoutPrefix = ".wh."
	opaqueWhiteout = ".wh..wh..opq"
)

func applyWhiteout(rootfs, name string) error {
	base := path.Base(name)
	if !strings.HasPrefix(base, whiteoutPrefix) {
		return nil
	}
	dir, err := resolveInRootfs(rootfs, path.Dir(name))
	if err != nil {
		return err
	}
	if base == opaqueWhiteout {
		entries, err := os.ReadDir(dir)
		if os.IsNotExist(err) {
			return nil
		}
		if err != nil {
			return err
		}
		for _, e := range entries {
			if err := os.RemoveAll(filepath.Join(dir, e.Name())); err != nil {
				return err
			}
		}
		return nil
	}
	return os.RemoveAll(filepath.Join(dir, strings.TrimPrefix(base, whiteoutPrefix)))
}

func extractEntry(rootfs string, hdr *tar.Header, body io.Reader) error {
	name := path.Clean("/" + hdr.Name)
	if name == "/" {
		return nil
	}
	parent, err := resolveInRootfs(rootfs, path.Dir(name))
	if err != nil {
		return err
	}
	if err := os.MkdirAll(parent, 0o755); err != nil {
		return err
	}
	target := filepath.Join(parent, path.Base(name))
	mode := os.FileMode(hdr.Mode) & (os.ModePerm | os.ModeSetuid | os.ModeSetgid | os.ModeSticky)
	switch hdr.Typeflag {
	case tar.TypeDir:
		if fi, err := os.Lstat(target); err == nil && !fi.IsDir() {
			if err := os.RemoveAll(target); err != nil {
				return err
			}
		}
		if err := os.MkdirAll(target, 0o755); err != nil {
			return err
		}
		// Directories stay owner-writable so later layers can add to
		// them; the rootfs is bound read-only into the sandbox anyway.
		return os.Chmod(target, mode|0o700)
	case tar.TypeReg, tar.TypeRegA:
		if err := removeIfExists(target); err != nil {
			return err
		}
		f, err := os.OpenFile(target, os.O_CREATE|os.O_WRONLY|os.O_EXCL, 0o600)
		if err != nil {
			return err
		}
		if _, err := io.Copy(f, body); err != nil {
			f.Close()
			return err
		}
		if err := f.Close(); err != nil {
			return err
		}
		return os.Chmod(target, mode)
	case tar.TypeSymlink:
		if err := removeIfExists(target); err != nil {
			return err
		}
		return os.Symlink(hdr.Linkname, target)
	case tar.TypeLink:
		src, err := resolveInRootfs(rootfs, "/"+hdr.Linkname)
		if err != nil {
			return err
		}
		if err := removeIfExists(target); err != nil {
			return err
		}
		return os.Link(src, target)
	case tar.TypeChar, tar.TypeBlock, tar.TypeFifo:
		return nil // device nodes: /dev comes from the host; fifos are not image content
	case tar.TypeXGlobalHeader:
		return nil
	default:
		return fmt.Errorf("%s: unsupported tar entry type %q", hdr.Name, hdr.Typeflag)
	}
}

func removeIfExists(target string) error {
	if _, err := os.Lstat(target); err == nil {
		return os.RemoveAll(target)
	}
	return nil
}

// resolveInRootfs maps an absolute in-image path to its host location
// under rootfs, following symlinks the way the sandbox would: relative
// to the link's directory, absolute ones relative to rootfs, and ".."
// above the root stays at the root, exactly as in a chroot — so a
// hostile layer can rewrite the image, never the host. Components that
// do not exist yet are appended verbatim.
func resolveInRootfs(rootfs, name string) (string, error) {
	root := filepath.Clean(rootfs)
	resolved := root
	remaining := strings.Split(strings.Trim(path.Clean("/"+name), "/"), "/")
	for hops := 0; len(remaining) > 0; {
		component := remaining[0]
		remaining = remaining[1:]
		switch component {
		case "", ".":
			continue
		case "..":
			if resolved != root {
				resolved = filepath.Dir(resolved)
			}
			continue
		}
		next := filepath.Join(resolved, component)
		fi, err := os.Lstat(next)
		if err != nil || fi.Mode()&os.ModeSymlink == 0 {
			resolved = next
			continue
		}
		if hops++; hops > 64 {
			return "", fmt.Errorf("%s: too many levels of symbolic links", name)
		}
		link, err := os.Readlink(next)
		if err != nil {
			return "", err
		}
		linkParts := strings.Split(strings.Trim(path.Clean(link), "/"), "/")
		if path.IsAbs(link) {
			resolved = root
		}
		remaining = append(linkParts, remaining...)
	}
	if resolved != root && !strings.HasPrefix(resolved, root+string(filepath.Separator)) {
		return "", fmt.Errorf("%s: escapes the rootfs", name)
	}
	return resolved, nil
}

// forEachEntry streams a layer, transparently decompressing gzip
// in-process and zstd through the zstd binary.
func forEachEntry(layerPath string, fn func(*tar.Header, io.Reader) error) error {
	f, err := os.Open(layerPath)
	if err != nil {
		return err
	}
	defer f.Close()
	r, closeFn, err := decompress(f)
	if err != nil {
		return fmt.Errorf("%s: %w", layerPath, err)
	}
	defer closeFn()
	tr := tar.NewReader(r)
	for {
		hdr, err := tr.Next()
		if errors.Is(err, io.EOF) {
			return nil
		}
		if err != nil {
			return fmt.Errorf("%s: %w", layerPath, err)
		}
		if err := fn(hdr, tr); err != nil {
			return fmt.Errorf("%s: %s: %w", filepath.Base(layerPath), hdr.Name, err)
		}
	}
}

var (
	gzipMagic = []byte{0x1f, 0x8b}
	zstdMagic = []byte{0x28, 0xb5, 0x2f, 0xfd}
)

func decompress(f *os.File) (io.Reader, func(), error) {
	br := bufio.NewReader(f)
	head, _ := br.Peek(4)
	switch {
	case bytes.HasPrefix(head, gzipMagic):
		gz, err := gzip.NewReader(br)
		if err != nil {
			return nil, nil, err
		}
		return gz, func() { gz.Close() }, nil
	case bytes.HasPrefix(head, zstdMagic):
		cmd := exec.Command("zstd", "-dc")
		cmd.Stdin = br
		cmd.Stderr = os.Stderr
		out, err := cmd.StdoutPipe()
		if err != nil {
			return nil, nil, err
		}
		if err := cmd.Start(); err != nil {
			return nil, nil, fmt.Errorf("zstd layer needs the zstd binary: %w", err)
		}
		return out, func() { out.Close(); _ = cmd.Wait() }, nil
	default:
		return br, func() {}, nil
	}
}