package api

import (
	"archive/zip"
	"context"
	"encoding/json"
	"errors"
	"git.misaka.ren/m1saka/dsh-web/internal/store"
	"io"
	"mime"
	"net/http"
	"os"
	"path"
	"sort"
	"strings"
	"syscall"
	"time"

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

const maxUpload = 32 << 20

var uploadSlots = make(chan struct{}, 4)

const (
	maxZipEntries       = 10000
	maxZipDepth         = 64
	maxZipBytes   int64 = 10 << 30
	zipTimeout          = 2 * time.Minute
)

var zipSlots = make(chan struct{}, 4)
var errZipLimit = errors.New("archive limit exceeded")

func renameNoReplace(root *os.Root, old, next string) error {
	anchor, err := root.Open(".")
	if err != nil {
		return err
	}
	defer anchor.Close()
	openParent := func(p string) (int, error) {
		return unix.Openat2(int(anchor.Fd()), path.Dir(p), &unix.OpenHow{
			Flags:   unix.O_RDONLY | unix.O_DIRECTORY | unix.O_CLOEXEC,
			Resolve: unix.RESOLVE_BENEATH | unix.RESOLVE_NO_SYMLINKS,
		})
	}
	src, err := openParent(old)
	if err != nil {
		return err
	}
	defer unix.Close(src)
	dst, err := openParent(next)
	if err != nil {
		return err
	}
	defer unix.Close(dst)
	var st unix.Stat_t
	if err = unix.Fstatat(src, path.Base(old), &st, unix.AT_SYMLINK_NOFOLLOW); err != nil {
		return err
	}
	if st.Mode&unix.S_IFMT == unix.S_IFLNK {
		return errors.New("unsafe path")
	}
	return unix.Renameat2(src, path.Base(old), dst, path.Base(next), unix.RENAME_NOREPLACE)
}

func safeRel(p string) (string, error) {
	p = strings.TrimSpace(strings.ReplaceAll(p, `\`, `/`))
	if p == "" {
		return ".", nil
	}
	if strings.HasPrefix(p, "/") {
		return "", errors.New("invalid path")
	}
	c := path.Clean(p)
	if c == ".." || strings.HasPrefix(c, "../") {
		return "", errors.New("invalid path")
	}
	return c, nil
}

func workspacePath(uid string) string { return path.Join("workspaces", uid, "files") }

func (a *API) openWorkspace(u store.User, instanceIDs ...string) (*os.Root, error) {
	instanceID := ""
	if len(instanceIDs) > 0 {
		instanceID = instanceIDs[0]
	}
	base, err := os.OpenRoot(a.root)
	if err != nil {
		return nil, err
	}
	defer base.Close()
	parts := []string{"workspaces", path.Join("workspaces", u.ID), workspacePath(u.ID)}
	if instanceID != "" {
		rel := path.Join("workspaces", u.ID, "instances", instanceID, "home")
		if err := rootSafe(base, rel, false); err != nil {
			return nil, err
		}
		return base.OpenRoot(rel)
	}
	for _, p := range parts {
		if err = rootSafe(base, p, true); err != nil {
			return nil, err
		}
		if err = base.Mkdir(p, 0750); err != nil && !errors.Is(err, os.ErrExist) {
			return nil, err
		}
	}
	if err = base.Chown(workspacePath(u.ID), 1000, 1000); err != nil {
		return nil, err
	}
	return base.OpenRoot(workspacePath(u.ID))
}

func rootSafe(root *os.Root, rel string, finalMayMissing bool) error {
	if rel == "." {
		return nil
	}
	parts := strings.Split(rel, "/")
	for i := range parts {
		q := strings.Join(parts[:i+1], "/")
		st, err := root.Lstat(q)
		if err != nil {
			if finalMayMissing && i == len(parts)-1 && errors.Is(err, os.ErrNotExist) {
				return nil
			}
			return err
		}
		if st.Mode()&os.ModeSymlink != 0 || (i < len(parts)-1 && !st.IsDir()) {
			return errors.New("unsafe path")
		}
	}
	return nil
}

func (a *API) workspace(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/api/ws/download" {
		var flush func()
		w, flush = deferJSONResponse(w)
		defer flush()
	}
	u, _, err := a.user(r)
	if err != nil {
		fail(w, 401, "authentication required")
		return
	}
	instanceID := strings.TrimSpace(r.URL.Query().Get("instance"))
	if instanceID == "" && r.URL.Path != "/api/ws/usage" {
		fail(w, http.StatusGone, "shared files removed; instance required")
		return
	}
	if instanceID != "" {
		// Pin this instance's home without holding the instance lifecycle lock
		// across client reads/writes. Lifecycle mutations reject active leases.
		inst := a.instanceLocks.Acquire(instanceID)
		i, e := a.s.GetInstance(instanceID)
		if e != nil || i.UID != u.ID {
			inst.Release()
			fail(w, 404, "instance not found")
			return
		}
		if r.Method != http.MethodGet && (i.Status == "error" || a.upgradeRecoveryRequired(i)) {
			inst.Release()
			fail(w, 409, "instance requires recovery")
			return
		}
		a.readersMu.Lock()
		if a.workspaceReaders == nil {
			a.workspaceReaders = make(map[string]int)
		}
		a.workspaceReaders[instanceID]++
		a.readersMu.Unlock()
		inst.Release()
		defer func() {
			a.readersMu.Lock()
			if a.workspaceReaders != nil {
				a.workspaceReaders[instanceID]--
				if a.workspaceReaders[instanceID] == 0 {
					delete(a.workspaceReaders, instanceID)
				}
			}
			a.readersMu.Unlock()
		}()
	}
	user := a.userLocks.Acquire(u.ID)
	root, err := a.openWorkspace(u, instanceID)
	user.Release()
	if err != nil {
		fail(w, 500, "storage unavailable")
		return
	}
	defer root.Close()
	switch strings.TrimPrefix(r.URL.Path, "/api/ws/") {
	case "usage":
		a.wsUsage(w, r, root, u)
	case "tree":
		a.wsTree(w, r, root, u)
	case "mkdir":
		a.wsMkdir(w, r, root, u)
	case "rename":
		a.wsRename(w, r, root, u)
	case "delete":
		a.wsDelete(w, r, root, u)
	case "upload":
		a.wsUpload(w, r, root, u)
	case "download":
		a.wsDownload(w, r, root)
	default:
		fail(w, 404, "not found")
	}
}

func entry(root *os.Root, rel string) (os.FileInfo, error) {
	if err := rootSafe(root, rel, false); err != nil {
		return nil, err
	}
	return root.Lstat(rel)
}

func (a *API) wsUsage(w http.ResponseWriter, r *http.Request, root *os.Root, u store.User) {
	if r.Method != "GET" {
		fail(w, 405, "method not allowed")
		return
	}
	user := a.userLocks.Acquire(u.ID)
	n, e := a.userUsage(u)
	user.Release()
	if e != nil {
		fail(w, 500, "usage unavailable")
		return
	}
	writeJSON(w, map[string]int64{"used_bytes": n, "quota_bytes": u.Quota})
}

func (a *API) wsTree(w http.ResponseWriter, r *http.Request, root *os.Root, u store.User) {
	user := a.userLocks.Acquire(u.ID)
	locked := user
	defer func() {
		if locked != nil {
			locked.Release()
		}
	}()
	if r.Method != http.MethodGet {
		fail(w, 405, "method not allowed")
		return
	}
	rel, err := safeRel(r.URL.Query().Get("path"))
	if err != nil {
		fail(w, 400, "invalid path")
		return
	}
	st, err := entry(root, rel)
	if err != nil || !st.IsDir() {
		fail(w, 404, "not found")
		return
	}
	es, err := readRootDir(root, rel)
	if err != nil {
		if errors.Is(err, errDirectoryLimit) {
			fail(w, 413, "directory entry limit exceeded")
			return
		}
		fail(w, 404, "not found")
		return
	}
	out := make([]map[string]any, 0, len(es))
	for _, de := range es {
		p := de.Name()
		if rel != "." {
			p = rel + "/" + p
		}
		st, err := entry(root, p)
		if err != nil || (!st.IsDir() && !st.Mode().IsRegular()) {
			continue
		}
		sz := int64(0)
		if !st.IsDir() {
			sz = st.Size()
		}
		out = append(out, map[string]any{"name": de.Name(), "path": p, "is_dir": st.IsDir(), "size": sz, "mod_time": st.ModTime()})
	}
	sort.Slice(out, func(i, j int) bool { return out[i]["name"].(string) < out[j]["name"].(string) })
	locked.Release()
	locked = nil
	writeJSON(w, out)
}

func (a *API) wsMkdir(w http.ResponseWriter, r *http.Request, root *os.Root, u store.User) {
	if r.Method != http.MethodPost {
		fail(w, 405, "method not allowed")
		return
	}
	var x struct {
		Path string `json:"path"`
	}
	if json.NewDecoder(r.Body).Decode(&x) != nil {
		fail(w, 400, "invalid json")
		return
	}
	p, e := safeRel(x.Path)
	user := a.userLocks.Acquire(u.ID)
	locked := user
	defer func() {
		if locked != nil {
			locked.Release()
		}
	}()
	if e != nil || p == "." || rootSafe(root, path.Dir(p), false) != nil {
		fail(w, 400, "cannot create directory")
		return
	}
	if e = root.Mkdir(p, 0750); e != nil {
		fail(w, 400, "cannot create directory")
		return
	}
	if f, e := root.OpenFile(p, os.O_RDONLY|syscall.O_NOFOLLOW|syscall.O_NONBLOCK, 0); e == nil {
		_ = f.Chown(1000, 1000)
		f.Close()
	}
	locked.Release()
	locked = nil
	a.s.AddAudit(u.ID, "workspace.mkdir", p, "")
	writeJSON(w, map[string]bool{"ok": true})
}

func (a *API) wsRename(w http.ResponseWriter, r *http.Request, root *os.Root, u store.User) {
	if r.Method != http.MethodPost {
		fail(w, 405, "method not allowed")
		return
	}
	var x struct {
		Path    string `json:"path"`
		NewPath string `json:"new_path"`
	}
	if json.NewDecoder(r.Body).Decode(&x) != nil {
		fail(w, 400, "invalid json")
		return
	}
	p, e := safeRel(x.Path)
	q, e2 := safeRel(x.NewPath)
	user := a.userLocks.Acquire(u.ID)
	locked := user
	defer func() {
		if locked != nil {
			locked.Release()
		}
	}()
	if e != nil || e2 != nil || p == "." || q == "." || rootSafe(root, p, false) != nil || rootSafe(root, path.Dir(q), false) != nil {
		fail(w, 400, "rename failed")
		return
	}
	st, e := root.Lstat(p)
	if e != nil || st.Mode()&os.ModeSymlink != 0 {
		fail(w, 400, "rename failed")
		return
	}
	if _, e = root.Lstat(q); e == nil {
		fail(w, 400, "rename failed")
		return
	}
	if e = renameNoReplace(root, p, q); e != nil {
		fail(w, 400, "rename failed")
		return
	}
	locked.Release()
	locked = nil
	a.s.AddAudit(u.ID, "workspace.rename", p, q)
	writeJSON(w, map[string]bool{"ok": true})
}

func (a *API) wsDelete(w http.ResponseWriter, r *http.Request, root *os.Root, u store.User) {
	user := a.userLocks.Acquire(u.ID)
	locked := user
	defer func() {
		if locked != nil {
			locked.Release()
		}
	}()
	if r.Method != http.MethodDelete {
		fail(w, 405, "method not allowed")
		return
	}
	p, e := safeRel(r.URL.Query().Get("path"))
	if e != nil || p == "." {
		fail(w, 400, "delete failed")
		return
	}
	st, e := entry(root, p)
	if e != nil || st.Mode()&os.ModeSymlink != 0 {
		fail(w, 400, "delete failed")
		return
	}
	if st.IsDir() {
		es, e := readRootDir(root, p)
		if e != nil || len(es) != 0 {
			fail(w, 400, "delete failed")
			return
		}
	}
	if e = root.Remove(p); e != nil {
		fail(w, 400, "delete failed")
		return
	}
	locked.Release()
	locked = nil
	a.s.AddAudit(u.ID, "workspace.delete", p, "")
	writeJSON(w, map[string]bool{"ok": true})
}

func (a *API) wsUpload(w http.ResponseWriter, r *http.Request, root *os.Root, u store.User) {
	if r.Method != "POST" {
		fail(w, 405, "method not allowed")
		return
	}
	select {
	case uploadSlots <- struct{}{}:
		defer func() { <-uploadSlots }()
	default:
		fail(w, 429, "too many concurrent uploads")
		return
	}
	usedBefore, usageErr := a.userUsage(u)
	if usageErr != nil || usedBefore >= u.Quota {
		fail(w, 413, "quota exceeded or usage unavailable")
		return
	}
	r.Body = http.MaxBytesReader(w, r.Body, maxUpload+(1<<20))
	controller := http.NewResponseController(w)
	if err := controller.SetReadDeadline(time.Now().Add(30 * time.Second)); err != nil && !errors.Is(err, http.ErrNotSupported) {
		fail(w, 400, "upload unavailable")
		return
	}
	defer controller.SetReadDeadline(time.Time{})
	if r.ParseMultipartForm(1<<20) != nil {
		fail(w, 400, "invalid upload")
		return
	}
	defer r.MultipartForm.RemoveAll()
	user := a.userLocks.Acquire(u.ID)
	locked := user
	defer func() {
		if locked != nil {
			locked.Release()
		}
	}()
	dir, e := safeRel(r.FormValue("path"))
	if e != nil || rootSafe(root, dir, false) != nil {
		fail(w, 400, "invalid directory")
		return
	}
	input, h, e := r.FormFile("file")
	if e != nil {
		fail(w, 400, "file required")
		return
	}
	defer input.Close()
	name, e := safeRel(h.Filename)
	if e != nil || name == "." || path.Base(name) != name {
		fail(w, 400, "invalid filename")
		return
	}
	if !uploadExtensionAllowed(name, os.Getenv("UPLOAD_ALLOWED_EXTENSIONS")) {
		fail(w, 415, "upload file type denied")
		return
	}
	target := path.Join(dir, name)
	if _, e = root.Lstat(target); !errors.Is(e, os.ErrNotExist) {
		fail(w, 409, "target exists")
		return
	}
	if err := a.s.DB.QueryRow("SELECT quota FROM users WHERE id=?", u.ID).Scan(&u.Quota); err != nil {
		fail(w, 413, "quota unavailable")
		return
	}
	used, e := a.userUsage(u)
	if e != nil || h.Size > maxUpload || used+h.Size > u.Quota {
		fail(w, 413, "quota or upload limit exceeded")
		return
	}
	tmp := path.Join(dir, ".upload-"+store.ID())
	f, e := root.OpenFile(tmp, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
	if e != nil {
		fail(w, 400, "upload directory unavailable")
		return
	}
	defer root.Remove(tmp)
	n, e := io.Copy(f, io.LimitReader(input, maxUpload+1))
	if e == nil && (n > maxUpload || used+n > u.Quota) {
		e = errors.New("quota exceeded")
	}
	if e == nil {
		e = f.Chown(1000, 1000)
	}
	if e == nil {
		e = f.Sync()
	}
	f.Close()
	if e == nil {
		e = renameNoReplace(root, tmp, target)
	}
	if errors.Is(e, os.ErrExist) {
		fail(w, 409, "target exists")
		return
	}
	if e != nil {
		fail(w, 400, "upload failed")
		return
	}
	locked.Release()
	locked = nil
	a.s.AddAudit(u.ID, "workspace.upload", target, "")
	writeJSON(w, map[string]bool{"ok": true})
}

// Empty policy permits arbitrary workspace files; configured entries are exact
// case-insensitive extensions, not MIME claims supplied by the client.
func uploadExtensionAllowed(name, policy string) bool {
	if strings.TrimSpace(policy) == "" {
		return true
	}
	ext := strings.ToLower(path.Ext(name))
	for _, entry := range strings.Split(policy, ",") {
		entry = strings.ToLower(strings.TrimSpace(entry))
		if strings.HasPrefix(entry, ".") && len(entry) > 1 && ext == entry {
			return true
		}
	}
	return false
}

func (a *API) userUsage(u store.User) (int64, error) {
	base, e := os.OpenRoot(a.root)
	if e != nil {
		return 0, e
	}
	defer base.Close()
	root, e := base.OpenRoot(path.Join("workspaces", u.ID))
	if e != nil {
		return 0, e
	}
	defer root.Close()
	var total int64
	count := 0
	var walk func(string, int) error
	walk = func(p string, depth int) error {
		count++
		if count > 100000 || depth > 128 {
			return errors.New("tree limit")
		}
		st, e := root.Lstat(p)
		if e != nil {
			return e
		}
		if st.Mode().IsRegular() {
			total += st.Size()
			return nil
		}
		if !st.IsDir() {
			return nil
		}
		f, e := root.OpenFile(p, os.O_RDONLY|syscall.O_DIRECTORY|syscall.O_NOFOLLOW, 0)
		if e != nil {
			return e
		}
		defer f.Close()
		for {
			ds, readErr := f.ReadDir(128)
			for _, d := range ds {
				if e = walk(path.Join(p, d.Name()), depth+1); e != nil {
					return e
				}
			}
			if readErr == io.EOF {
				return nil
			}
			if readErr != nil {
				return readErr
			}
			if count >= 100000 {
				return errors.New("tree limit")
			}
		}
	}
	e = walk(".", 0)
	return total, e
}

func (a *API) wsDownload(w http.ResponseWriter, r *http.Request, root *os.Root) {
	if r.Method != http.MethodGet {
		fail(w, 405, "method not allowed")
		return
	}
	p, e := safeRel(r.URL.Query().Get("path"))
	if e != nil {
		fail(w, 400, "invalid path")
		return
	}
	anchor, e := root.Open(".")
	if e != nil {
		fail(w, 404, "not found")
		return
	}
	defer anchor.Close()
	f, e := openDownload(anchor, p)
	if e != nil {
		fail(w, 404, "not found")
		return
	}
	defer f.Close()
	st, e := f.Stat()
	if e == nil && st.IsDir() {
		a.wsDownloadZip(w, r, f, p)
		return
	}
	if e != nil || !st.Mode().IsRegular() {
		fail(w, 400, "invalid file")
		return
	}
	w.Header().Set("Content-Disposition", "attachment")
	http.ServeContent(w, r, path.Base(p), st.ModTime(), f)
}

func openDownload(anchor *os.File, p string) (*os.File, error) {
	fd, err := unix.Openat2(int(anchor.Fd()), p, &unix.OpenHow{
		Flags:   unix.O_RDONLY | unix.O_NONBLOCK | unix.O_CLOEXEC,
		Resolve: unix.RESOLVE_BENEATH | unix.RESOLVE_NO_SYMLINKS | unix.RESOLVE_NO_XDEV,
	})
	if err != nil {
		return nil, err
	}
	f := os.NewFile(uintptr(fd), p)
	st, err := f.Stat()
	if err == nil && (!st.IsDir() && !st.Mode().IsRegular() || st.Mode().IsRegular() && st.Sys().(*syscall.Stat_t).Nlink != 1) {
		err = errors.New("unsafe file")
	}
	if err != nil {
		f.Close()
		return nil, err
	}
	return f, nil
}

type zipEntry struct {
	name string
	info os.FileInfo
}

func zipManifest(ctx context.Context, anchor *os.File) ([]zipEntry, error) {
	var entries []zipEntry
	var total int64
	var walk func(string, int) error
	walk = func(p string, depth int) error {
		if err := ctx.Err(); err != nil {
			return err
		}
		if depth > maxZipDepth || len(entries) >= maxZipEntries {
			return errZipLimit
		}
		f, err := openDownload(anchor, p)
		if err != nil {
			return err
		}
		defer f.Close()
		st, err := f.Stat()
		if err != nil {
			return err
		}
		if p != "." {
			if st.Size() < 0 || st.Mode().IsRegular() && st.Size() > maxZipBytes-total {
				return errZipLimit
			}
			if st.Mode().IsRegular() {
				total += st.Size()
			}
			entries = append(entries, zipEntry{p, st})
		}
		if !st.IsDir() {
			return nil
		}
		for {
			ds, err := f.ReadDir(128)
			if err != nil && err != io.EOF {
				return err
			}
			for _, d := range ds {
				if strings.ContainsAny(d.Name(), "\\:") {
					return errors.New("unsafe archive name")
				}
				if e := walk(path.Join(p, d.Name()), depth+1); e != nil {
					return e
				}
			}
			if err == io.EOF {
				return nil
			}
		}
	}
	err := walk(".", 0)
	return entries, err
}

type zipContextWriter struct {
	ctx context.Context
	w   io.Writer
}

func (w zipContextWriter) Write(p []byte) (int, error) {
	if err := w.ctx.Err(); err != nil {
		return 0, err
	}
	return w.w.Write(p)
}

func streamZip(ctx context.Context, w io.Writer, anchor *os.File, entries []zipEntry) error {
	z := zip.NewWriter(zipContextWriter{ctx, w})
	buf := make([]byte, 32<<10)
	for _, item := range entries {
		if err := ctx.Err(); err != nil {
			return err
		}
		f, err := openDownload(anchor, item.name)
		if err != nil {
			return err
		}
		err = func() error {
			defer f.Close()
			st, err := f.Stat()
			if err != nil {
				return err
			}
			if !os.SameFile(st, item.info) || st.Mode() != item.info.Mode() || st.Size() != item.info.Size() || !st.ModTime().Equal(item.info.ModTime()) {
				return errors.New("file changed")
			}
			h, err := zip.FileInfoHeader(st)
			if err != nil {
				return err
			}
			h.Name = item.name
			if st.IsDir() {
				h.Name += "/"
			}
			dst, err := z.CreateHeader(h)
			if err != nil || st.IsDir() {
				return err
			}
			n, err := io.CopyBuffer(dst, io.LimitReader(f, item.info.Size()), buf)
			if err == nil && n != item.info.Size() {
				err = io.ErrUnexpectedEOF
			}
			if err == nil {
				after, statErr := f.Stat()
				if statErr != nil {
					return statErr
				}
				if after.Size() != st.Size() || !after.ModTime().Equal(st.ModTime()) {
					return errors.New("file changed")
				}
			}
			return err
		}()
		if err != nil {
			return err
		}
	}
	return z.Close()
}

func (a *API) wsDownloadZip(w http.ResponseWriter, r *http.Request, anchor *os.File, p string) {
	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()
	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{})
	entries, err := zipManifest(ctx, anchor)
	if err != nil {
		code := 400
		if errors.Is(err, errZipLimit) {
			code = 413
		}
		fail(w, code, "cannot archive directory")
		return
	}
	name := path.Base(p)
	if p == "." {
		name = "workspace"
	}
	w.Header().Set("Content-Type", "application/zip")
	w.Header().Set("Content-Disposition", mime.FormatMediaType("attachment", map[string]string{"filename": name + ".zip"}))
	if err := streamZip(ctx, w, anchor, entries); err != nil {
		// A failed stream must not become a successful, silently incomplete archive.
		panic(http.ErrAbortHandler)
	}
}

const maxDirectoryEntries = 10000

var errDirectoryLimit = errors.New("directory entry limit exceeded")

func readRootDir(root *os.Root, p string) ([]os.DirEntry, error) {
	f, e := root.OpenFile(p, os.O_RDONLY|syscall.O_DIRECTORY|syscall.O_NOFOLLOW, 0)
	if e != nil {
		return nil, e
	}
	defer f.Close()
	entries, err := f.ReadDir(maxDirectoryEntries + 1)
	if len(entries) > maxDirectoryEntries {
		return nil, errDirectoryLimit
	}
	if err == io.EOF {
		err = nil
	}
	return entries, err
}