package archive

import (
	"archive/tar"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"sync"
)

// Extract writes an opened archive into filesDir and envPath, refusing
// anything a tar could smuggle: absolute paths, "..", the protocol marker,
// links, devices and FIFOs, two entries naming one destination, more
// entries than MaxEntries, content over Cap and a decoded stream over
// MaxDecoded. Every directory must appear before anything inside it (Build's
// sorted order always does), so each one carries its archived mode.
//
// Nothing reaches the destination unless the whole archive is accepted:
// entries are written into a private staging directory beside filesDir
// (and env into a temporary file beside envPath), made durable, and moved
// into place only after the stream has been read to its end. filesDir must be absent
// or an empty directory (ErrNotEmpty otherwise; a non-empty directory
// would keep files the archive does not name), and its parent and
// envPath's parent must exist. envPath is replaced when the archive
// carries env and removed when it does not, so afterwards PlainDigest of
// the pair equals the archive's. The files move first and env second; a
// failure between the two renames leaves the files in place and is
// returned.
func Extract(r io.Reader, filesDir, envPath string) error {
	if filesDir == "" {
		return errors.New("archive: no files directory named")
	}
	filesDir = filepath.Clean(filesDir)
	if envPath != "" {
		envPath = filepath.Clean(envPath)
		if rel, err := filepath.Rel(filesDir, envPath); err == nil && filepath.IsLocal(rel) {
			return fmt.Errorf("%w: the env file %s is inside the files directory", ErrBadPath, envPath)
		}
	}
	rootMode, existed, err := checkDestination(filesDir)
	if err != nil {
		return err
	}
	stage, err := os.MkdirTemp(filepath.Dir(filesDir), "."+filepath.Base(filesDir)+".extract-")
	if err != nil {
		return err
	}
	x := &extraction{stage: stage, envPath: envPath, files: newSyncPool(syncWorkers)}
	defer x.cleanup()
	if x.root, err = os.OpenRoot(stage); err != nil {
		return err
	}
	if err := x.read(r); err != nil {
		return err
	}
	return x.commit(filesDir, rootMode, existed)
}

// checkDestination requires filesDir to be absent or an empty directory
// (not a symlink to one) and returns the mode the extracted root takes:
// the existing directory's, else 0755.
func checkDestination(filesDir string) (fs.FileMode, bool, error) {
	info, err := os.Lstat(filesDir)
	if errors.Is(err, fs.ErrNotExist) {
		return 0o755, false, nil
	}
	if err != nil {
		return 0, false, err
	}
	if !info.IsDir() {
		return 0, false, fmt.Errorf("%w: the destination %s is not a directory", ErrBadPath, filesDir)
	}
	f, err := os.Open(filesDir)
	if err != nil {
		return 0, false, err
	}
	names, err := f.Readdirnames(1)
	f.Close()
	if len(names) > 0 {
		return 0, false, fmt.Errorf("%w: %s", ErrNotEmpty, filesDir)
	}
	if err != nil && err != io.EOF {
		return 0, false, err
	}
	return info.Mode().Perm(), true, nil
}

type modeFix struct {
	rel  string
	mode fs.FileMode
}

// extraction is one Extract in progress.
type extraction struct {
	stage   string   // the staging directory, removed unless committed
	root    *os.Root // stage, for every write
	envPath string
	envTmp  string // the staged env file, if the archive carried one
	dirs    []modeFix
	files   *syncPool // fsyncs and closes each staged file as it is written
}

func (x *extraction) cleanup() {
	x.files.wait()
	if x.root != nil {
		x.root.Close()
	}
	if x.stage != "" {
		removeStage(x.stage)
	}
	if x.envTmp != "" {
		os.Remove(x.envTmp)
	}
}

// removeStage removes a staging directory whose directories may already
// carry their final, restrictive modes. It holds only what Extract
// created: regular files and directories.
func removeStage(stage string) {
	filepath.WalkDir(stage, func(p string, d fs.DirEntry, err error) error {
		if err == nil && d.IsDir() {
			os.Chmod(p, 0o700)
		}
		return nil
	})
	os.RemoveAll(stage)
}

// read writes every entry into the staging area and reads the stream to
// its end.
func (x *extraction) read(r io.Reader) error {
	zr, err := newDecoder(r)
	if err != nil {
		return err
	}
	defer zr.Close()
	dec := &boundedReader{r: zr, left: decodedLimit}
	tr := tar.NewReader(dec)
	rules := newStreamRules()
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		_, rel, isEnv, err := rules.admit(hdr)
		if err != nil {
			return err
		}
		mode := fs.FileMode(hdr.Mode).Perm()
		if hdr.Typeflag == tar.TypeDir {
			// Created writable so children can land; the archived mode is
			// set in commit, before the directory is synced.
			if err := x.root.Mkdir(rel, 0o700); err != nil {
				return err
			}
			x.dirs = append(x.dirs, modeFix{rel, mode})
			continue
		}
		if isEnv {
			err = x.writeEnv(tr, hdr.Size, mode)
		} else {
			err = x.writeFile(tr, rel, hdr.Size, mode)
		}
		if err != nil {
			return err
		}
	}
	if _, err := io.Copy(io.Discard, dec); err != nil {
		return err
	}
	return requireEOF(r)
}

func (x *extraction) writeFile(src io.Reader, rel string, size int64, mode fs.FileMode) error {
	// O_EXCL: a name that already exists in the fresh staging tree (on a
	// case-insensitive filesystem, "A" after "a") is refused, not merged.
	f, err := x.root.OpenFile(rel, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
	if err != nil {
		return err
	}
	return x.finishFile(f, src, size, mode)
}

func (x *extraction) writeEnv(src io.Reader, size int64, mode fs.FileMode) error {
	if x.envPath == "" {
		return fmt.Errorf("%w: the archive carries env and no env path was given", ErrBadPath)
	}
	f, err := os.CreateTemp(filepath.Dir(x.envPath), "."+filepath.Base(x.envPath)+".extract-")
	if err != nil {
		return err
	}
	x.envTmp = f.Name()
	return x.finishFile(f, src, size, mode)
}

// finishFile copies exactly size bytes and sets the archived mode (fchmod,
// so the umask does not apply; §8.1), then hands the still-open file to
// the sync pool, which fsyncs and closes it. Syncing the open descriptor
// means a mode that refuses every open (0000) cannot stop the sync, and
// the pool overlaps the syncs with the rest of the extraction.
func (x *extraction) finishFile(f *os.File, src io.Reader, size int64, mode fs.FileMode) error {
	if _, err := io.CopyN(f, src, size); err != nil {
		f.Close()
		return err
	}
	if err := f.Chmod(mode); err != nil {
		f.Close()
		return err
	}
	x.files.submit(func() error { return syncAndClose(f, fsyncFile) })
	return nil
}

// commit makes the staged tree durable with its archived modes, then
// publishes it.
//
// Crash ordering. Every staged file is fsynced with its archived mode
// (finishFile) and every staged directory is fsynced after its archived
// mode is set, deepest level first, then the staged root with its final
// mode: all before the rename that publishes the tree. The rename is made
// durable by fsyncing the destination's parent, then env is published the
// same way. So a crash at any point leaves either no published tree (a
// hidden .extract- directory to sweep) or a complete one carrying the
// archived modes, never a published tree whose directories still have
// their staging modes. The landing writes its marker only after Extract
// returns nil.
//
// Cost. Scoped fsyncs only, never a filesystem-wide flush: other tenants'
// dirty data on the box is not this landing's to wait for. They run on
// syncWorkers workers, since concurrent fsyncs on ext4 and xfs share
// journal commits. The cut budget prices this time under its Archive and
// RestoreStaged terms (moves.Estimate).
func (x *extraction) commit(filesDir string, rootMode fs.FileMode, existed bool) error {
	if err := x.files.wait(); err != nil {
		return err
	}
	if err := x.settleDirs(); err != nil {
		return err
	}
	if err := settleDir(func() (*os.File, error) { return os.Open(x.stage) }, rootMode); err != nil {
		return err
	}
	if existed {
		// os.Remove refuses a directory that is no longer empty.
		if err := os.Remove(filesDir); err != nil {
			return err
		}
	}
	if err := os.Rename(x.stage, filesDir); err != nil {
		return err
	}
	x.stage = ""
	if err := syncDir(filepath.Dir(filesDir)); err != nil {
		return err
	}
	if x.envPath == "" {
		return nil
	}
	if x.envTmp == "" {
		if err := os.Remove(x.envPath); err != nil && !errors.Is(err, fs.ErrNotExist) {
			return err
		}
	} else {
		if err := os.Rename(x.envTmp, x.envPath); err != nil {
			return err
		}
		x.envTmp = ""
	}
	return syncDir(filepath.Dir(x.envPath))
}

// settleDirs sets every staged directory's archived mode through an open
// descriptor and fsyncs it, one depth level at a time, deepest first: a
// directory is opened by a path through its ancestors, which are still
// 0700 until their own level, so a restrictive mode (0300, 0000) never
// blocks a deeper directory, and the descriptor syncs whatever the mode.
func (x *extraction) settleDirs() error {
	levels := map[int][]modeFix{}
	for _, d := range x.dirs {
		depth := strings.Count(d.rel, "/")
		levels[depth] = append(levels[depth], d)
	}
	depths := make([]int, 0, len(levels))
	for depth := range levels {
		depths = append(depths, depth)
	}
	sort.Sort(sort.Reverse(sort.IntSlice(depths)))
	for _, depth := range depths {
		p := newSyncPool(syncWorkers)
		for _, d := range levels[depth] {
			p.submit(func() error {
				return settleDir(func() (*os.File, error) { return x.root.Open(d.rel) }, d.mode)
			})
		}
		if err := p.wait(); err != nil {
			return err
		}
	}
	return nil
}

func settleDir(open func() (*os.File, error), mode fs.FileMode) error {
	f, err := open()
	if err != nil {
		return err
	}
	if err := f.Chmod(mode); err != nil {
		f.Close()
		return err
	}
	return syncAndClose(f, syncDirFile)
}

// fsyncFile is the one fsync every staged file and directory goes
// through. A variable so a test can see each one and the mode it carried.
var fsyncFile = func(f *os.File) error { return f.Sync() }

// syncDirFile fsyncs a directory. Windows cannot sync a directory handle;
// the platform's boxes are Linux, and this only keeps that build whole.
func syncDirFile(f *os.File) error {
	if runtime.GOOS == "windows" {
		return nil
	}
	return fsyncFile(f)
}

func syncAndClose(f *os.File, sync func(*os.File) error) error {
	err := sync(f)
	if cerr := f.Close(); err == nil {
		err = cerr
	}
	return err
}

// syncDir makes a rename in dir durable.
func syncDir(dir string) error {
	d, err := os.Open(dir)
	if err != nil {
		return err
	}
	return syncAndClose(d, syncDirFile)
}

// syncWorkers is how many fsyncs run at once (defaultSyncWorkers, per
// platform). A variable so the timing test can compare serial and parallel.
var syncWorkers = defaultSyncWorkers

// syncPool runs jobs on a bounded set of workers and keeps the first
// error. At most workers jobs run and workers more wait, so the staged
// files held open for syncing stay bounded.
type syncPool struct {
	jobs chan func() error
	wg   sync.WaitGroup
	once sync.Once
	mu   sync.Mutex
	err  error
}

func newSyncPool(workers int) *syncPool {
	workers = max(workers, 1)
	p := &syncPool{jobs: make(chan func() error, workers)}
	p.wg.Add(workers)
	for range workers {
		go func() {
			defer p.wg.Done()
			for job := range p.jobs {
				if err := job(); err != nil {
					p.mu.Lock()
					if p.err == nil {
						p.err = err
					}
					p.mu.Unlock()
				}
			}
		}()
	}
	return p
}

func (p *syncPool) submit(job func() error) { p.jobs <- job }

// wait closes the pool, waits for every job and returns the first error.
// Safe to call more than once.
func (p *syncPool) wait() error {
	p.once.Do(func() {
		close(p.jobs)
		p.wg.Wait()
	})
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.err
}