package artifact

import (
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"bugabinga.net/luci/internal/ciconfig"
)

const (
	maxFiles = 100
	maxFile  = 10 << 20
	maxTotal = 50 << 20
)

type File struct {
	Name string
	Size int64
}

type Store struct{ DataDir string }

func (s Store) root(runID, childID, set string) (string, error) {
	for label, value := range map[string]string{"run": runID, "child": childID, "set": set} {
		if !safeSegment(value) || staging(value) {
			return "", fmt.Errorf("unsafe artifact %s %q", label, value)
		}
	}
	return filepath.Join(s.DataDir, "artifacts", runID, childID, set), nil
}

func (s Store) Collect(workspace, runID, childID string, declaration ciconfig.Artifact) error {
	destination, err := s.root(runID, childID, declaration.Name)
	if err != nil {
		return err
	}
	if _, err := os.Lstat(destination); err == nil {
		return fmt.Errorf("artifact set %q already exists", declaration.Name)
	} else if !errors.Is(err, os.ErrNotExist) {
		return err
	}
	parent := filepath.Dir(destination)
	if err := os.MkdirAll(parent, 0o755); err != nil {
		return err
	}
	stage, err := os.MkdirTemp(parent, ".stage-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(stage)
	info, err := os.Lstat(workspace)
	if err != nil || !info.IsDir() || info.Mode()&os.ModeSymlink != 0 {
		return fmt.Errorf("unsafe artifact workspace")
	}
	workspaceRoot, err := os.OpenRoot(workspace)
	if err != nil {
		return err
	}
	defer workspaceRoot.Close()
	existing, err := s.List(runID, childID)
	if err != nil {
		return err
	}
	files := int64(len(existing))
	var total int64
	for _, file := range existing {
		total += file.Size
	}
	for _, relative := range declaration.Paths {
		if !safePath(relative) || hasStagingPart(relative) {
			return fmt.Errorf("unsafe artifact path %q", relative)
		}
		if err := rejectSymlinkAncestors(workspaceRoot, relative); err != nil {
			if errors.Is(err, os.ErrNotExist) && !declaration.Required {
				continue
			}
			if errors.Is(err, os.ErrNotExist) {
				return fmt.Errorf("required artifact %q path %q missing", declaration.Name, relative)
			}
			return fmt.Errorf("artifact %q path %q: %w", declaration.Name, relative, err)
		}
		if err := copyPath(stage, workspaceRoot, relative, relative, &files, &total); err != nil {
			return fmt.Errorf("artifact %q path %q: %w", declaration.Name, relative, err)
		}
	}
	return os.Rename(stage, destination)
}

func copyPath(stage string, sourceRoot *os.Root, source, relative string, files, total *int64) error {
	return fs.WalkDir(sourceRoot.FS(), source, func(name string, entry fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		info, err := sourceRoot.Lstat(name)
		if err != nil {
			return err
		}
		if info.Mode()&os.ModeSymlink != 0 {
			return fmt.Errorf("symlink rejected")
		}
		rel, err := filepath.Rel(source, name)
		if err != nil {
			return err
		}
		destination := filepath.Join(stage, filepath.FromSlash(relative), filepath.FromSlash(rel))
		if info.IsDir() {
			return os.MkdirAll(destination, 0o755)
		}
		if !info.Mode().IsRegular() {
			return fmt.Errorf("non-regular file rejected")
		}
		*files++
		if *files > maxFiles {
			return fmt.Errorf("exceeds %d files", maxFiles)
		}
		if info.Size() > maxFile {
			return fmt.Errorf("exceeds %d byte file limit", maxFile)
		}
		input, err := sourceRoot.Open(name)
		if err != nil {
			return err
		}
		opened, err := input.Stat()
		if err != nil {
			input.Close()
			return err
		}
		if !opened.Mode().IsRegular() || !os.SameFile(info, opened) {
			input.Close()
			return fmt.Errorf("file changed while collecting")
		}
		output, err := os.OpenFile(destination, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o644)
		if err != nil {
			input.Close()
			return err
		}
		written, copyErr := io.Copy(output, io.LimitReader(input, maxFile+1))
		closeErr := errors.Join(input.Close(), output.Close())
		if copyErr != nil || closeErr != nil {
			return errors.Join(copyErr, closeErr)
		}
		if written > maxFile {
			return fmt.Errorf("exceeds %d byte file limit", maxFile)
		}
		*total += written
		if *total > maxTotal {
			return fmt.Errorf("exceeds %d byte total limit", maxTotal)
		}
		return nil
	})
}

func (s Store) List(runID, childID string) ([]File, error) {
	if !safeSegment(runID) || staging(runID) || (childID != "" && (!safeSegment(childID) || staging(childID))) {
		return nil, os.ErrNotExist
	}
	root, err := os.OpenRoot(filepath.Join(s.DataDir, "artifacts"))
	if errors.Is(err, os.ErrNotExist) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	defer root.Close()
	base := runID
	if childID != "" {
		base += "/" + childID
	}
	if err := rejectSymlinkAncestors(root, base); errors.Is(err, os.ErrNotExist) {
		return nil, nil
	} else if err != nil {
		return nil, err
	}
	var out []File
	err = fs.WalkDir(root.FS(), base, func(name string, entry fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if staging(entry.Name()) {
			if entry.IsDir() {
				return fs.SkipDir
			}
			return nil
		}
		info, err := root.Lstat(name)
		if err != nil {
			return err
		}
		if info.Mode()&os.ModeSymlink != 0 || (!info.IsDir() && !info.Mode().IsRegular()) {
			return fmt.Errorf("invalid artifact storage")
		}
		if info.IsDir() {
			return nil
		}
		rel, err := filepath.Rel(runID, name)
		if err != nil {
			return err
		}
		out = append(out, File{Name: rel, Size: info.Size()})
		return nil
	})
	sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name })
	return out, err
}

func (s Store) Open(runID, name string) (*os.File, error) {
	if !safeSegment(runID) || staging(runID) || !safePath(name) || hasStagingPart(name) {
		return nil, os.ErrNotExist
	}
	root, err := os.OpenRoot(filepath.Join(s.DataDir, "artifacts"))
	if err != nil {
		return nil, err
	}
	defer root.Close()
	artifactName := runID + "/" + name
	if err := rejectSymlinkAncestors(root, artifactName); err != nil {
		return nil, os.ErrNotExist
	}
	info, err := root.Lstat(artifactName)
	if err != nil || !info.Mode().IsRegular() {
		return nil, os.ErrNotExist
	}
	file, err := root.Open(artifactName)
	if err != nil {
		return nil, err
	}
	opened, err := file.Stat()
	if err != nil || !opened.Mode().IsRegular() || !os.SameFile(info, opened) {
		file.Close()
		return nil, os.ErrNotExist
	}
	return file, nil
}

func rejectSymlinkAncestors(root *os.Root, name string) error {
	for _, part := range strings.Split(name, "/") {
		if !safeSegment(part) || staging(part) {
			return os.ErrNotExist
		}
	}
	var current string
	for _, part := range strings.Split(name, "/") {
		if current == "" {
			current = part
		} else {
			current += "/" + part
		}
		info, err := root.Lstat(current)
		if err != nil {
			return err
		}
		if info.Mode()&os.ModeSymlink != 0 || (!info.IsDir() && current != name) {
			return fmt.Errorf("symlink or non-directory artifact ancestor")
		}
	}
	return nil
}

func staging(value string) bool { return strings.HasPrefix(value, ".stage-") }
func hasStagingPart(value string) bool {
	for _, part := range strings.Split(value, "/") {
		if staging(part) {
			return true
		}
	}
	return false
}
func safeSegment(value string) bool {
	return value != "" && value == filepath.Base(value) && value != "." && value != ".." && !strings.Contains(value, string(filepath.Separator))
}
func safePath(value string) bool {
	return value != "" && filepath.ToSlash(filepath.Clean(value)) == value && !strings.HasPrefix(value, "../") && !strings.HasPrefix(value, "/")
}