package publish

import (
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"

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

type Request struct {
	Adapter   string
	Image     string
	From      string
	To        string
	Workspace string
}

type Executor interface {
	Publish(context.Context, Request, io.Writer) error
}

type Service struct {
	RegistryAuthFile string
	PkgRoot          string
	SiteRoot         string
	Run              func(context.Context, []string, []string, io.Writer) error
}

const maxRegistryArchiveBytes int64 = 10 << 30

var registryReference = regexp.MustCompile(`^(?:localhost|[a-z0-9](?:[a-z0-9.-]*[a-z0-9])?(?::[0-9]+)?)/[a-z0-9]+(?:[._-][a-z0-9]+)*(?:/[a-z0-9]+(?:[._-][a-z0-9]+)*)*(?::[A-Za-z0-9_][A-Za-z0-9_.-]{0,127})?(?:@sha256:[a-f0-9]{64})?$`)
var manifestDigest = regexp.MustCompile(`^sha256:[a-f0-9]{64}$`)

var renamePublish = os.Rename
var renameExchange = func(oldPath, newPath string) error {
	return unix.Renameat2(unix.AT_FDCWD, oldPath, unix.AT_FDCWD, newPath, unix.RENAME_EXCHANGE)
}

func (s Service) Publish(ctx context.Context, request Request, output io.Writer) error {
	switch request.Adapter {
	case "registry":
		if s.RegistryAuthFile == "" {
			return fmt.Errorf("publish registry requires LUCI_PUBLISH_REGISTRY_AUTH_FILE")
		}
		if request.Image != "" {
			return fmt.Errorf("publish registry image is no longer supported; write an OCI archive in the job and use from")
		}
		if request.From == "" {
			return fmt.Errorf("publish registry requires from")
		}
		if !registryReference.MatchString(request.To) {
			return fmt.Errorf("publish registry requires a docker registry target, not %q", request.To)
		}
		archive, cleanup, err := stageRegistryArchive(ctx, request.Workspace, request.From)
		if err != nil {
			return err
		}
		defer cleanup()
		digestFile, err := os.CreateTemp(filepath.Dir(archive), "digest-")
		if err != nil {
			return err
		}
		digestPath := digestFile.Name()
		if err := digestFile.Close(); err != nil {
			_ = os.Remove(digestPath)
			return err
		}
		defer os.Remove(digestPath)
		run := s.Run
		if run == nil {
			run = runSkopeo
		}
		if err := run(ctx, []string{"copy", "--authfile", s.RegistryAuthFile, "--digestfile", digestPath, "oci-archive:" + archive, "docker://" + request.To}, os.Environ(), output); err != nil {
			return err
		}
		digest, err := os.ReadFile(digestPath)
		if err != nil {
			return fmt.Errorf("read published registry digest: %w", err)
		}
		value := strings.TrimSpace(string(digest))
		if !manifestDigest.MatchString(value) {
			return fmt.Errorf("invalid published registry digest %q", value)
		}
		reference, err := canonicalRegistryReference(request.To, value)
		if err != nil {
			return err
		}
		_, err = fmt.Fprintf(output, "[luci] registry-published destination=%s digest=%s reference=%s\n", request.To, value, reference)
		return err
	case "pkg":
		return publishFiles(ctx, s.PkgRoot, "LUCI_PUBLISH_PKG_ROOT", request)
	case "site":
		return publishFiles(ctx, s.SiteRoot, "LUCI_PUBLISH_SITE_ROOT", request)
	default:
		return fmt.Errorf("unsupported publish adapter %q", request.Adapter)
	}
}

func canonicalRegistryReference(destination, digest string) (string, error) {
	repository, pinned, hasPinned := strings.Cut(destination, "@")
	if hasPinned && pinned != digest {
		return "", fmt.Errorf("published registry digest %s conflicts with pinned destination %s", digest, destination)
	}
	if separator := strings.LastIndex(repository, "/"); strings.LastIndex(repository, ":") > separator {
		repository = repository[:strings.LastIndex(repository, ":")]
	}
	return repository + "@" + digest, nil
}

func runSkopeo(ctx context.Context, args, env []string, output io.Writer) error {
	command := exec.CommandContext(ctx, "skopeo", args...)
	command.Env = env
	command.Stdout = output
	command.Stderr = output
	return command.Run()
}

func stageRegistryArchive(ctx context.Context, workspace, relative string) (string, func(), error) {
	if err := ctx.Err(); err != nil {
		return "", nil, err
	}
	source, err := safeSource(workspace, relative)
	if err != nil {
		return "", nil, err
	}
	inputFD, err := unix.Open(source, unix.O_RDONLY|unix.O_CLOEXEC|unix.O_NOFOLLOW, 0)
	if err != nil {
		return "", nil, fmt.Errorf("open registry archive: %w", err)
	}
	input := os.NewFile(uintptr(inputFD), source)
	var stat unix.Stat_t
	if err := unix.Fstat(inputFD, &stat); err != nil {
		_ = input.Close()
		return "", nil, err
	}
	if stat.Mode&unix.S_IFMT != unix.S_IFREG {
		_ = input.Close()
		return "", nil, fmt.Errorf("registry archive must be a regular file")
	}
	if stat.Size <= 0 || stat.Size > maxRegistryArchiveBytes {
		_ = input.Close()
		return "", nil, fmt.Errorf("registry archive size %d exceeds allowed range", stat.Size)
	}
	var fs unix.Statfs_t
	if err := unix.Statfs(workspace, &fs); err != nil {
		_ = input.Close()
		return "", nil, err
	}
	if fs.Bsize <= 0 || (uint64(stat.Size)+uint64(fs.Bsize)-1)/uint64(fs.Bsize) > uint64(fs.Bavail) {
		_ = input.Close()
		return "", nil, fmt.Errorf("insufficient workspace space to stage registry archive")
	}
	stage, err := os.MkdirTemp(workspace, ".luci-registry-")
	if err != nil {
		_ = input.Close()
		return "", nil, err
	}
	archive := filepath.Join(stage, "archive")
	output, err := os.OpenFile(archive, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
	if err == nil {
		_, err = io.Copy(output, contextReader{ctx: ctx, reader: input})
		closeErr := output.Close()
		if err == nil {
			err = closeErr
		}
	}
	inputErr := input.Close()
	if err == nil {
		err = inputErr
	}
	if err != nil {
		_ = os.RemoveAll(stage)
		return "", nil, err
	}
	return archive, func() { _ = os.RemoveAll(stage) }, nil
}

func publishFiles(ctx context.Context, root, envName string, request Request) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	if root == "" {
		return fmt.Errorf("publish %s requires %s", request.Adapter, envName)
	}
	if request.From == "" || request.Image != "" {
		return fmt.Errorf("publish %s requires from", request.Adapter)
	}
	source, err := safeSource(request.Workspace, request.From)
	if err != nil {
		return err
	}
	if err := rejectSymlinks(source); err != nil {
		return err
	}
	target, parent, err := safeTarget(root, request.To)
	if err != nil {
		return err
	}
	if err := os.MkdirAll(parent, 0o755); err != nil {
		return err
	}
	if err := safeTargetParent(root, request.To); err != nil {
		return err
	}
	stage, err := os.MkdirTemp(parent, ".luci-publish-")
	if err != nil {
		return err
	}
	cleanupStage := true
	defer func() {
		if cleanupStage {
			_ = os.RemoveAll(stage)
		}
	}()
	staged := filepath.Join(stage, "value")
	if err := copyTree(ctx, source, staged); err != nil {
		return err
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	lock, err := lockPublishParent(parent)
	if err != nil {
		return err
	}
	defer lock.Close()
	info, err := os.Lstat(target)
	if err == nil {
		if info.Mode()&os.ModeSymlink != 0 {
			return fmt.Errorf("publish target is symlink %s", target)
		}
		if err := renameExchange(staged, target); err != nil {
			return fmt.Errorf("publish replacement requires atomic rename exchange: %w", err)
		}
		return nil
	}
	if !os.IsNotExist(err) {
		return err
	}
	if err := renamePublish(staged, target); err != nil {
		return err
	}
	return nil
}

func safeJoin(root, relative string) (string, error) {
	clean := filepath.Clean(relative)
	if strings.TrimSpace(root) == "" || strings.TrimSpace(relative) == "" || clean == "." || filepath.IsAbs(relative) || clean == ".." || strings.HasPrefix(clean, ".."+string(filepath.Separator)) {
		return "", fmt.Errorf("unsafe publish path %q", relative)
	}
	return filepath.Join(root, clean), nil
}

func safeSource(root, relative string) (string, error) {
	path, err := safeJoin(root, relative)
	if err != nil {
		return "", err
	}
	current := root
	for _, part := range strings.Split(filepath.Clean(relative), string(filepath.Separator)) {
		current = filepath.Join(current, part)
		info, err := os.Lstat(current)
		if err != nil {
			return "", err
		}
		if info.Mode()&os.ModeSymlink != 0 {
			return "", fmt.Errorf("publish source traverses symlink %s", current)
		}
	}
	return path, nil
}

func safeTarget(root, relative string) (string, string, error) {
	path, err := safeJoin(root, relative)
	if err != nil {
		return "", "", err
	}
	if err := safeTargetParent(root, relative); err != nil && !os.IsNotExist(err) {
		return "", "", err
	}
	return path, filepath.Dir(path), nil
}

func safeTargetParent(root, relative string) error {
	current := root
	parts := strings.Split(filepath.Clean(relative), string(filepath.Separator))
	for _, part := range append([]string{"."}, parts[:len(parts)-1]...) {
		if part != "." {
			current = filepath.Join(current, part)
		}
		info, err := os.Lstat(current)
		if err != nil {
			return err
		}
		if !info.IsDir() || info.Mode()&os.ModeSymlink != 0 {
			return fmt.Errorf("publish target traverses unsafe parent %s", current)
		}
	}
	return nil
}

func lockPublishParent(parent string) (*os.File, error) {
	path := filepath.Join(parent, ".luci-publish.lock")
	fd, err := unix.Open(path, unix.O_CREAT|unix.O_RDWR|unix.O_CLOEXEC|unix.O_NOFOLLOW, 0o644)
	if err != nil {
		return nil, err
	}
	file := os.NewFile(uintptr(fd), path)
	var stat unix.Stat_t
	if err := unix.Fstat(fd, &stat); err != nil {
		_ = file.Close()
		return nil, err
	}
	if stat.Mode&unix.S_IFMT != unix.S_IFREG {
		_ = file.Close()
		return nil, fmt.Errorf("publish lock is not a regular file %s", path)
	}
	if err := unix.Flock(fd, unix.LOCK_EX); err != nil {
		_ = file.Close()
		return nil, err
	}
	return file, nil
}

func rejectSymlinks(path string) error {
	info, err := os.Lstat(path)
	if err != nil {
		return err
	}
	if info.Mode()&os.ModeSymlink != 0 {
		return fmt.Errorf("publish source contains symlink %s", path)
	}
	if !info.IsDir() {
		if !info.Mode().IsRegular() {
			return fmt.Errorf("unsupported publish source %s", path)
		}
		return nil
	}
	entries, err := os.ReadDir(path)
	if err != nil {
		return err
	}
	for _, entry := range entries {
		if err := rejectSymlinks(filepath.Join(path, entry.Name())); err != nil {
			return err
		}
	}
	return nil
}

func copyTree(ctx context.Context, source, target string) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	info, err := os.Stat(source)
	if err != nil {
		return err
	}
	if info.IsDir() {
		if err := os.MkdirAll(target, info.Mode().Perm()); err != nil {
			return err
		}
		entries, err := os.ReadDir(source)
		if err != nil {
			return err
		}
		for _, entry := range entries {
			if err := copyTree(ctx, filepath.Join(source, entry.Name()), filepath.Join(target, entry.Name())); err != nil {
				return err
			}
		}
		return nil
	}
	if !info.Mode().IsRegular() {
		return fmt.Errorf("unsupported publish source %s", source)
	}
	if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
		return err
	}
	input, err := os.Open(source)
	if err != nil {
		return err
	}
	defer input.Close()
	output, err := os.OpenFile(target, os.O_CREATE|os.O_EXCL|os.O_WRONLY, info.Mode().Perm())
	if err != nil {
		return err
	}
	_, copyErr := io.Copy(output, contextReader{ctx: ctx, reader: input})
	closeErr := output.Close()
	if copyErr != nil {
		return copyErr
	}
	return closeErr
}

type contextReader struct {
	ctx    context.Context
	reader io.Reader
}

func (r contextReader) Read(data []byte) (int, error) {
	if err := r.ctx.Err(); err != nil {
		return 0, err
	}
	return r.reader.Read(data)
}