package api

import (
	"archive/tar"
	"compress/gzip"
	"context"
	"errors"
	"io"
	"mime"
	"net/http"
	"os"
	"path/filepath"
	"syscall"
	"time"

	"git.misaka.ren/m1saka/dsh-web/internal/store"
)

var errSnapshotDownloadLimit = errors.New("snapshot archive exceeds size, entry or depth limit")
var errUnsafeSnapshotArchive = errors.New("unsafe snapshot link or unsupported file type")

type snapshotArchiveWriter struct {
	ctx       context.Context
	w         io.Writer
	remaining int64
}

func (w *snapshotArchiveWriter) Write(p []byte) (int, error) {
	if err := w.ctx.Err(); err != nil {
		return 0, err
	}
	if int64(len(p)) > w.remaining {
		return 0, errSnapshotDownloadLimit
	}
	n, err := w.w.Write(p)
	w.remaining -= int64(n)
	return n, err
}

func snapshotArchive(ctx context.Context, root *os.Root, out io.Writer, limit int64) error {
	// Bound headers, padding and gzip expansion independently of file content.
	gz := gzip.NewWriter(&snapshotArchiveWriter{ctx, out, limit + (256 << 20)})
	tw := tar.NewWriter(&snapshotArchiveWriter{ctx, gz, limit + (128 << 20)})
	remaining := limit
	entries := 100000
	var walk func(*os.Root, string, string, int) error
	walk = func(parent *os.Root, name, rel string, depth int) error {
		if err := ctx.Err(); err != nil {
			return err
		}
		entries--
		if entries < 0 || depth > 128 {
			return errSnapshotDownloadLimit
		}
		st, err := parent.Lstat(name)
		if err != nil {
			return err
		}
		link := ""
		if st.Mode()&os.ModeSymlink != 0 {
			link, err = parent.Readlink(name)
			if err != nil {
				return err
			}
			if filepath.IsAbs(link) || !filepath.IsLocal(filepath.Join(filepath.Dir(rel), link)) {
				return errUnsafeSnapshotArchive
			}
			// Root resolution also catches escapes through chains of relative links.
			if _, err = root.Stat(rel); err != nil && !errors.Is(err, os.ErrNotExist) {
				return errUnsafeSnapshotArchive
			}
		} else if !st.IsDir() && !st.Mode().IsRegular() {
			return errUnsafeSnapshotArchive
		}
		if st.Mode().IsRegular() && st.Size() > remaining {
			return errSnapshotDownloadLimit
		}
		if st.Mode().IsRegular() {
			remaining -= st.Size()
		}
		h, err := tar.FileInfoHeader(st, link)
		if err != nil {
			return err
		}
		h.Name = filepath.ToSlash(rel)
		if st.IsDir() {
			h.Name += "/"
		}
		if err = tw.WriteHeader(h); err != nil {
			return err
		}
		if st.Mode()&os.ModeSymlink != 0 {
			return nil
		}
		if st.IsDir() {
			child, err := parent.OpenRoot(name)
			if err != nil {
				return err
			}
			defer child.Close()
			actual, err := child.Stat(".")
			if err != nil || !os.SameFile(st, actual) {
				return errors.New("snapshot directory changed")
			}
			f, err := child.Open(".")
			if err != nil {
				return err
			}
			defer f.Close()
			for {
				children, err := f.ReadDir(128)
				if err != nil && err != io.EOF {
					return err
				}
				for _, entry := range children {
					if err := walk(child, entry.Name(), filepath.Join(rel, entry.Name()), depth+1); err != nil {
						return err
					}
				}
				if err == io.EOF {
					return nil
				}
			}
		}
		f, err := parent.OpenFile(name, os.O_RDONLY|syscall.O_NOFOLLOW|syscall.O_NONBLOCK, 0)
		if err != nil {
			return err
		}
		defer f.Close()
		actual, err := f.Stat()
		if err != nil || !actual.Mode().IsRegular() || !os.SameFile(st, actual) || actual.Size() != st.Size() {
			return errors.New("snapshot file changed")
		}
		if _, err = io.CopyN(tw, f, st.Size()); err != nil {
			return err
		}
		after, err := f.Stat()
		if err != nil || after.Size() != st.Size() || !after.ModTime().Equal(st.ModTime()) {
			return errors.New("snapshot file changed")
		}
		return nil
	}
	if err := walk(root, ".", ".", 0); err != nil {
		return err
	}
	if err := tw.Close(); err != nil {
		return err
	}
	return gz.Close()
}

func (a *API) snapshotDownload(w http.ResponseWriter, r *http.Request, i store.Instance) {
	if r.Method != http.MethodGet {
		fail(w, 405, "method not allowed")
		return
	}
	ids := r.URL.Query()["download"]
	if len(ids) != 1 || !snapshotIDPattern.MatchString(ids[0]) {
		fail(w, 400, "invalid snapshot ID")
		return
	}
	id := ids[0]
	select {
	case zipSlots <- struct{}{}:
		defer func() { <-zipSlots }()
	default:
		fail(w, 429, "too many archive downloads")
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), zipTimeout)
	defer cancel()
	archive, err := os.CreateTemp("", "dsh-snapshot-*.tar.gz")
	if err != nil {
		fail(w, 500, "temporary archive unavailable")
		return
	}
	defer os.Remove(archive.Name())
	defer archive.Close()
	code := 500
	err = func() error {
		lock := a.instanceLocks.Acquire(i.ID)
		defer lock.Release()
		if err := ctx.Err(); err != nil {
			return err
		}
		current, err := a.s.GetInstance(i.ID)
		if err != nil || current.UID != i.UID {
			code = 404
			return errors.New("not found")
		}
		root, err := managementRoot(a.root, "workspaces", current.UID, "instances", current.ID, id)
		if err != nil {
			if errors.Is(err, os.ErrNotExist) {
				code = 404
			}
			if errors.Is(err, errUnsafeManagementPath) {
				code = 400
			}
			return errors.New("snapshot unavailable or unsafe path")
		}
		defer root.Close()
		return snapshotArchive(ctx, root, archive, upgradeSnapshotLimit)
	}()
	if err != nil {
		if errors.Is(err, errUnsafeSnapshotArchive) {
			code = 400
		}
		if errors.Is(err, errSnapshotDownloadLimit) {
			code = 413
		}
		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
			code = 408
		}
		fail(w, code, err.Error())
		return
	}
	if ctx.Err() != nil {
		fail(w, 408, "snapshot download timed out or canceled")
		return
	}
	if _, err = archive.Seek(0, io.SeekStart); err != nil {
		fail(w, 500, "archive unavailable")
		return
	}
	controller := http.NewResponseController(w)
	if err = controller.SetWriteDeadline(time.Now().Add(zipTimeout)); err != nil && !errors.Is(err, http.ErrNotSupported) {
		fail(w, 500, "download unavailable")
		return
	}
	defer controller.SetWriteDeadline(time.Time{})
	a.s.AddAudit(i.UID, "snapshot.download", i.ID, id)
	w.Header().Set("Content-Type", "application/gzip")
	w.Header().Set("Cache-Control", "no-store")
	w.Header().Set("Content-Disposition", mime.FormatMediaType("attachment", map[string]string{"filename": id + ".tar.gz"}))
	http.ServeContent(w, r, id+".tar.gz", time.Time{}, archive)
}