package archive

import (
	"archive/tar"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/user"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/klauspost/compress/zstd"
)

// Route отображает имя в архиве на путь назначения. ok=false — запись пропускается.
type Route func(name string) (dest string, ok bool)

type Extractor struct {
	Route    Route
	Warnings []string
	Files    int

	users  map[string]int
	groups map[string]int
	dirs   []dirMeta
}

type dirMeta struct {
	path  string
	mode  os.FileMode
	mtime time.Time
}

// Extract распаковывает zstd-tar из r.
func (e *Extractor) Extract(r io.Reader) error {
	zr, err := zstd.NewReader(r)
	if err != nil {
		return err
	}
	defer zr.Close()
	tr := tar.NewReader(zr)
	for {
		hdr, err := tr.Next()
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			return err
		}
		if strings.Contains(hdr.Name, "..") {
			e.warn("пропущено подозрительное имя %q", hdr.Name)
			continue
		}
		dest, ok := e.Route(hdr.Name)
		if !ok {
			continue
		}
		if err := e.extractOne(tr, hdr, filepath.Clean(dest)); err != nil {
			return fmt.Errorf("%s: %w", hdr.Name, err)
		}
	}
	// права и время каталогов — в конце, иначе запись файлов их перебьёт
	for i := len(e.dirs) - 1; i >= 0; i-- {
		d := e.dirs[i]
		_ = os.Chmod(d.path, d.mode)
		_ = os.Chtimes(d.path, d.mtime, d.mtime)
	}
	return nil
}

func (e *Extractor) warn(format string, a ...any) {
	e.Warnings = append(e.Warnings, fmt.Sprintf(format, a...))
}

func (e *Extractor) extractOne(tr *tar.Reader, hdr *tar.Header, dest string) error {
	mode := os.FileMode(hdr.Mode).Perm()
	// tar хранит suid/sgid/sticky в битах c_ISUID/c_ISGID/c_ISVTX
	if hdr.Mode&04000 != 0 {
		mode |= os.ModeSetuid
	}
	if hdr.Mode&02000 != 0 {
		mode |= os.ModeSetgid
	}
	if hdr.Mode&01000 != 0 {
		mode |= os.ModeSticky
	}

	switch hdr.Typeflag {
	case tar.TypeDir:
		if fi, err := os.Lstat(dest); err == nil && !fi.IsDir() {
			if err := os.Remove(dest); err != nil {
				return err
			}
		}
		if err := os.MkdirAll(dest, 0o755); err != nil {
			return err
		}
		e.chown(dest, hdr)
		e.dirs = append(e.dirs, dirMeta{dest, mode, hdr.ModTime})
	case tar.TypeReg:
		if err := os.MkdirAll(filepath.Dir(dest), 0o755); err != nil {
			return err
		}
		// пишем во временный файл и переименовываем: работает и для занятых бинарников
		tmp := dest + ".bm-tmp"
		f, err := os.OpenFile(tmp, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o600)
		if err != nil {
			return err
		}
		if _, err := io.Copy(f, tr); err != nil {
			f.Close()
			os.Remove(tmp)
			return err
		}
		if err := f.Close(); err != nil {
			os.Remove(tmp)
			return err
		}
		e.chown(tmp, hdr)
		if err := os.Chmod(tmp, mode); err != nil {
			e.warn("%s: chmod: %v", dest, err)
		}
		_ = os.Chtimes(tmp, hdr.ModTime, hdr.ModTime)
		if fi, err := os.Lstat(dest); err == nil && fi.IsDir() {
			if err := os.RemoveAll(dest); err != nil {
				os.Remove(tmp)
				return err
			}
		}
		if err := os.Rename(tmp, dest); err != nil {
			os.Remove(tmp)
			return err
		}
	case tar.TypeSymlink:
		if err := os.MkdirAll(filepath.Dir(dest), 0o755); err != nil {
			return err
		}
		if fi, err := os.Lstat(dest); err == nil {
			if fi.IsDir() && fi.Mode()&os.ModeSymlink == 0 {
				if err := os.RemoveAll(dest); err != nil {
					return err
				}
			} else if err := os.Remove(dest); err != nil {
				return err
			}
		}
		if err := os.Symlink(hdr.Linkname, dest); err != nil {
			return err
		}
		e.chown(dest, hdr)
	default:
		return nil
	}
	e.Files++
	return nil
}

// chown выставляет владельца по имени из архива (uid на новом сервере может отличаться),
// при отсутствии такого пользователя — по числовому id.
func (e *Extractor) chown(p string, hdr *tar.Header) {
	if runtime.GOOS == "windows" {
		return
	}
	uid := e.lookup(&e.users, hdr.Uname, hdr.Uid, func(n string) (string, error) {
		u, err := user.Lookup(n)
		if err != nil {
			return "", err
		}
		return u.Uid, nil
	})
	gid := e.lookup(&e.groups, hdr.Gname, hdr.Gid, func(n string) (string, error) {
		g, err := user.LookupGroup(n)
		if err != nil {
			return "", err
		}
		return g.Gid, nil
	})
	if err := os.Lchown(p, uid, gid); err != nil {
		e.warn("%s: chown: %v", p, err)
	}
}

func (e *Extractor) lookup(cache *map[string]int, name string, fallback int, find func(string) (string, error)) int {
	if name == "" {
		return fallback
	}
	if *cache == nil {
		*cache = map[string]int{}
	}
	if id, ok := (*cache)[name]; ok {
		return id
	}
	id := fallback
	if s, err := find(name); err == nil {
		if n, err := strconv.Atoi(s); err == nil {
			id = n
		}
	}
	(*cache)[name] = id
	return id
}

// ReadManifest читает manifest.json из начала архива (без распаковки остального).
func ReadManifest(r io.Reader, v any) error {
	zr, err := zstd.NewReader(r)
	if err != nil {
		return err
	}
	defer zr.Close()
	tr := tar.NewReader(zr)
	hdr, err := tr.Next()
	if err != nil {
		return err
	}
	if hdr.Name != "manifest.json" {
		return fmt.Errorf("первая запись архива %q, ожидался manifest.json", hdr.Name)
	}
	return json.NewDecoder(tr).Decode(v)
}