//go:build unix

package promptdebug

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/sys/unix"
)

// secureDirectory is an opened directory descriptor. Every operation remains
// relative to that descriptor so later pathname swaps cannot redirect writes.
type secureDirectory struct {
	fd int
}

func openSecureDirectory(path string) (*secureDirectory, error) {
	if !filepath.IsAbs(path) {
		return nil, fmt.Errorf("directory must be absolute")
	}
	cleaned := filepath.Clean(path)
	current, err := unix.Open(string(filepath.Separator), unix.O_RDONLY|unix.O_DIRECTORY|unix.O_CLOEXEC, 0)
	if err != nil {
		return nil, err
	}
	defer func() {
		if current != -1 {
			_ = unix.Close(current)
		}
	}()
	for _, component := range strings.Split(strings.TrimPrefix(cleaned, string(filepath.Separator)), string(filepath.Separator)) {
		if component == "" {
			continue
		}
		next, created, err := openOrCreateSecureDirectory(current, component)
		if err != nil {
			return nil, err
		}
		if created {
			if err := unix.Fchmod(next, debugDirectoryMode); err != nil {
				_ = unix.Close(next)
				return nil, err
			}
		}
		_ = unix.Close(current)
		current = next
	}
	if err := unix.Fchmod(current, debugDirectoryMode); err != nil {
		return nil, err
	}
	result := &secureDirectory{fd: current}
	current = -1
	return result, nil
}

func (directory *secureDirectory) Close() error {
	if directory == nil || directory.fd < 0 {
		return nil
	}
	fd := directory.fd
	directory.fd = -1
	return unix.Close(fd)
}

func (directory *secureDirectory) openDirectory(components ...string) (*secureDirectory, error) {
	if directory == nil || directory.fd < 0 {
		return nil, fmt.Errorf("prompt debug directory is closed")
	}
	current := directory.fd
	owned := false
	defer func() {
		if owned {
			_ = unix.Close(current)
		}
	}()
	for _, component := range components {
		if err := validateSecureDirectoryName(component); err != nil {
			return nil, err
		}
		next, _, err := openOrCreateSecureDirectory(current, component)
		if err != nil {
			return nil, err
		}
		if err := unix.Fchmod(next, debugDirectoryMode); err != nil {
			_ = unix.Close(next)
			return nil, err
		}
		if owned {
			_ = unix.Close(current)
		}
		current = next
		owned = true
	}
	result := &secureDirectory{fd: current}
	owned = false
	return result, nil
}

func openOrCreateSecureDirectory(parent int, name string) (int, bool, error) {
	fd, err := unix.Openat(parent, name, unix.O_RDONLY|unix.O_DIRECTORY|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0)
	if err == nil {
		return fd, false, nil
	}
	if !errors.Is(err, unix.ENOENT) {
		return -1, false, fmt.Errorf("open secure directory %q: %w", name, err)
	}
	if err := unix.Mkdirat(parent, name, debugDirectoryMode); err != nil && !errors.Is(err, unix.EEXIST) {
		return -1, false, fmt.Errorf("create secure directory %q: %w", name, err)
	}
	fd, err = unix.Openat(parent, name, unix.O_RDONLY|unix.O_DIRECTORY|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0)
	if err != nil {
		return -1, false, fmt.Errorf("open created secure directory %q: %w", name, err)
	}
	return fd, true, nil
}

func validateSecureDirectoryName(name string) error {
	if name == "" || name == "." || name == ".." || strings.ContainsRune(name, filepath.Separator) {
		return fmt.Errorf("prompt debug directory component %q is invalid", name)
	}
	return nil
}

func (directory *secureDirectory) writeJSON(name string, value any) error {
	if directory == nil || directory.fd < 0 {
		return fmt.Errorf("prompt debug directory is closed")
	}
	if err := validateSecureDirectoryName(name); err != nil {
		return err
	}
	data, err := json.MarshalIndent(value, "", "  ")
	if err != nil {
		return fmt.Errorf("marshal prompt debug artifact: %w", err)
	}
	var information unix.Stat_t
	if err := unix.Fstatat(directory.fd, name, &information, unix.AT_SYMLINK_NOFOLLOW); err == nil {
		if information.Mode&unix.S_IFMT != unix.S_IFREG {
			return fmt.Errorf("prompt debug file %q is not a regular file", name)
		}
	} else if !errors.Is(err, unix.ENOENT) {
		return fmt.Errorf("inspect prompt debug file %q: %w", name, err)
	}
	temporaryName, temporary, err := directory.createTemporaryFile(name)
	if err != nil {
		return err
	}
	defer func() {
		if temporary != nil {
			_ = temporary.Close()
		}
		_ = unix.Unlinkat(directory.fd, temporaryName, 0)
	}()
	if err := temporary.Chmod(debugFileMode); err != nil {
		return err
	}
	if written, err := temporary.Write(data); err != nil {
		return err
	} else if written != len(data) {
		return io.ErrShortWrite
	}
	if err := temporary.Close(); err != nil {
		return err
	}
	temporary = nil
	if err := unix.Renameat(directory.fd, temporaryName, directory.fd, name); err != nil {
		return fmt.Errorf("replace prompt debug file %q: %w", name, err)
	}
	return nil
}

func (directory *secureDirectory) createTemporaryFile(name string) (string, *os.File, error) {
	for attempt := 0; attempt < 16; attempt++ {
		random := make([]byte, 12)
		if _, err := rand.Read(random); err != nil {
			return "", nil, err
		}
		temporaryName := "." + name + "." + hex.EncodeToString(random) + ".tmp"
		fd, err := unix.Openat(directory.fd, temporaryName, unix.O_WRONLY|unix.O_CREAT|unix.O_EXCL|unix.O_CLOEXEC|unix.O_NOFOLLOW, debugFileMode)
		if errors.Is(err, unix.EEXIST) {
			continue
		}
		if err != nil {
			return "", nil, fmt.Errorf("create temporary prompt debug file: %w", err)
		}
		return temporaryName, os.NewFile(uintptr(fd), temporaryName), nil
	}
	return "", nil, fmt.Errorf("create temporary prompt debug file: too many name collisions")
}