package zfs

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os/user"
	"regexp"
	"strings"

	"github.com/bitfield/script"
	"github.com/go-errors/errors"
	"golang.org/x/exp/maps"
	"golang.org/x/exp/slices"

	"gensokyo.cafe/xmnt/cfg"
	"gensokyo.cafe/xmnt/mnt"
	"gensokyo.cafe/xmnt/msg"
	"gensokyo.cafe/xmnt/util"
)

const ZfsBin = "/usr/bin/zfs"

type execFunc func(string, io.Reader, ...string) ([]byte, error)

type Permissions struct {
	Mount   bool
	LoadKey bool
}

var allPermission = Permissions{Mount: true, LoadKey: true}

type KeyStatus string

const (
	KeyStatusAvailable   KeyStatus = "available"
	KeyStatusUnavailable KeyStatus = "unavailable"
)

type CanMount string

const (
	CanMountOn     CanMount = "on"
	CanMountOff    CanMount = "off"
	CanMountNoAuto CanMount = "noauto"
)

type Dataset struct {
	Name           string
	GUID           string
	CanMount       CanMount
	MountPoint     string
	KeyStatus      KeyStatus
	EncryptionRoot string
	Mounted        bool

	permissions *Permissions
}

func (d *Dataset) loadPermissions() error {
	if d.permissions != nil {
		return nil
	}
	currentUser, err := user.Current()
	if err != nil {
		return errors.WrapPrefix(err, "cannot obtain current user", 0)
	}
	if currentUser.Uid == "0" {
		d.permissions = &allPermission
		return nil
	}

	permissions := &Permissions{}

	output, err := util.RunCommand(ZfsBin, nil, "allow", d.Name)
	if err != nil {
		return errors.WrapPrefix(err, fmt.Sprintf("cannot obtain permission list for zfs dataset %s", d.Name), 0)
	}
	pattern := regexp.MustCompile(`^\s*user ` + currentUser.Username + ` ([\w-]+(,[\w-]+)*)$`)
	script.Echo(string(output)).FilterLine(func(line string) string {
		match := pattern.FindStringSubmatch(line)
		if match != nil {
			for _, perm := range strings.Split(match[1], ",") {
				switch perm {
				case "mount":
					permissions.Mount = true
				case "load-key":
					permissions.LoadKey = true
				}
			}
		}
		return ""
	}).Wait()
	d.permissions = permissions

	return nil
}

func (d *Dataset) Permissions() (*Permissions, error) {
	if err := d.loadPermissions(); err != nil {
		return nil, err
	}
	return &*d.permissions, nil
}

func listCmd(name string, recursive bool) ([]byte, error) {
	zfsArgs := []string{"get", "-Ht", "filesystem", "guid,canmount,mountpoint,encryptionroot,keystatus,mounted"}
	if recursive {
		zfsArgs = append(zfsArgs, "-r")
	}
	if name != "" {
		zfsArgs = append(zfsArgs, name)
	}

	return util.RunCommand(ZfsBin, nil, zfsArgs...)
}

func ParseListOutput(output []byte) ([]*Dataset, error) {
	idx := map[string]*Dataset{}

	sc := bufio.NewScanner(bytes.NewReader(output))
	for sc.Scan() {
		fields := strings.Split(sc.Text(), "\t")
		if len(fields) != 4 {
			return nil, errors.Errorf("invalid zfs list output: %q", sc.Text())
		}

		name := fields[0]
		key := fields[1]
		value := fields[2]

		dataset, ok := idx[name]
		if !ok {
			dataset = &Dataset{Name: name}
			idx[name] = dataset
		}
		switch key {
		case "guid":
			dataset.GUID = value
		case "canmount":
			dataset.CanMount = CanMount(value)
		case "mountpoint":
			dataset.MountPoint = value
		case "encryptionroot":
			dataset.EncryptionRoot = value
		case "keystatus":
			dataset.KeyStatus = KeyStatus(value)
		case "mounted":
			dataset.Mounted = value == "yes"
		}
	}

	return maps.Values(idx), nil
}

func List(name string, recursive bool) ([]*Dataset, error) {
	listOutput, err := listCmd(name, recursive)
	if err != nil {
		return nil, errors.WrapPrefix(err, "cannot obtain list of zfs datasets", 0)
	}

	datasets, err := ParseListOutput(listOutput)
	if err != nil {
		return nil, errors.WrapPrefix(err, "cannot parse zfs list output", 0)
	}
	return datasets, nil
}

type Mounter struct {
	preset  *mnt.Preset
	dataset *Dataset
}

func NewMounterFromPreset(p *mnt.Preset) (mnt.Mounter, error) {
	preset := &*p
	if preset.Path == "" {
		return nil, errors.New("preset path is empty")
	}
	m := &Mounter{
		preset: preset,
	}

	// find the dataset
	if err := m.refresh(); err != nil {
		return nil, err
	}

	return m, nil
}

func (m *Mounter) refresh() error {
	datasets, err := List(m.preset.Path, false)
	if err != nil {
		return errors.WrapPrefix(err, "cannot obtain list of zfs datasets", 0)
	}
	var newDataset *Dataset
	for _, d := range datasets {
		if d.Name == m.preset.Path {
			newDataset = d
		}
	}
	if newDataset == nil {
		return errors.Errorf("cannot find zfs dataset %q", m.preset.Path)
	}
	m.dataset = newDataset
	return nil
}

func (m *Mounter) loadKey() error {
	if m.dataset.KeyStatus != KeyStatusUnavailable {
		return nil
	}
	if m.dataset.Name != m.dataset.EncryptionRoot {
		return errors.Errorf("cannot load key for zfs dataset %q: not an encryption root", m.dataset.Name)
	}

	key, err := util.ReadCredentialFile(m.dataset.GUID, cfg.Cfg.CredentialStore)
	if err != nil {
		return errors.WrapPrefix(err, "cannot load zfs key", 0)
	}

	perm, err := m.dataset.Permissions()
	if err != nil {
		return errors.WrapPrefix(err, "failed to load zfs key", 0)
	}
	var run execFunc
	if perm.LoadKey {
		run = util.RunCommand
	} else {
		run = util.RunPrivilegedCommand
	}

	msg.Infof("zfs load-key %q", m.dataset.Name)
	_, err = run(ZfsBin, strings.NewReader(key), "load-key", m.dataset.Name)
	if err != nil {
		return errors.WrapPrefix(err, "failed to load zfs key", 0)
	}
	return nil
}

func (m *Mounter) mount() error {
	if m.dataset.Mounted {
		return nil
	}
	if !slices.Contains([]CanMount{CanMountNoAuto, CanMountOn}, m.dataset.CanMount) {
		return errors.Errorf("cannot mount zfs dataset %q: canmount is %q", m.dataset.Name, m.dataset.CanMount)
	}
	if m.dataset.KeyStatus != KeyStatusAvailable {
		return errors.Errorf("cannot mount zfs dataset %q: not unlocked", m.dataset.Name)
	}

	mountPoint := m.preset.MountPoint
	if mountPoint == "" {
		mountPoint = m.dataset.MountPoint
	}
	if !util.IsValidMountPoint(mountPoint) {
		return errors.Errorf(
			"cannot mount zfs dataset %q: invalid mount point %q",
			m.dataset.Name, mountPoint,
		)
	}

	// If systemd mount unit exists, use it.
	// In this case, if we do not use systemd for mounting, systemd will mess with the mounting process, and the zfs
	// dataset will get unmounted immediately after mounting. See https://github.com/openzfs/zfs/issues/11248
	if err := util.SystemdMount(mountPoint); err == nil {
		return nil
	} else if !util.ShouldSkipSdMount(err) {
		return errors.WrapPrefix(err, fmt.Sprintf("failed to mount zfs dataset %q", m.dataset.Name), 0)
	}

	// mount using zfs command
	mountArgs := []string{"mount"}
	if m.preset.MountPoint != "" {
		// user specified the mount point
		mountArgs = append(mountArgs, "-o", "mountpoint="+mountPoint)
	}
	mountArgs = append(mountArgs, m.dataset.Name)

	perm, err := m.dataset.Permissions()
	if err != nil {
		return errors.WrapPrefix(err, fmt.Sprintf("failed to mount zfs dataset %q", m.dataset.Name), 0)
	}
	var run execFunc
	if perm.Mount {
		run = util.RunCommand
	} else {
		run = util.RunPrivilegedCommand
	}

	_, err = run(ZfsBin, nil, mountArgs...)
	if err != nil {
		return errors.WrapPrefix(err, fmt.Sprintf("failed to mount zfs dataset %q", m.dataset.Name), 0)
	}
	return nil
}

func (m *Mounter) Mount() error {
	if err := m.loadKey(); err != nil {
		return err
	}
	if err := m.refresh(); err != nil {
		return err
	}
	if err := m.mount(); err != nil {
		return err
	}

	// check
	if err := m.refresh(); err != nil {
		return err
	}
	if !m.dataset.Mounted {
		return errors.Errorf("zfs dataset %q is not mounted", m.dataset.Name)
	}

	return nil
}

func (m *Mounter) unmount() error {
	if !m.dataset.Mounted {
		return nil
	}

	// try to unmount with systemd
	mp := m.dataset.MountPoint
	if util.IsValidMountPoint(mp) {
		if err := util.SystemdUnmount(mp); err == nil {
			return nil
		} else if !util.ShouldSkipSdMount(err) {
			return errors.WrapPrefix(err, fmt.Sprintf("failed to unmount zfs dataset %q", m.dataset.Name), 0)
		}
	}

	// try to unmount with zfs command.
	perm, err := m.dataset.Permissions()
	if err != nil {
		return errors.WrapPrefix(err, fmt.Sprintf("failed to unmount zfs dataset %q", m.dataset.Name), 0)
	}
	var run execFunc
	if perm.Mount {
		run = util.RunCommand
	} else {
		run = util.RunPrivilegedCommand
	}

	_, err = run(ZfsBin, nil, "unmount", "-u", m.dataset.Name)
	if err != nil {
		return errors.WrapPrefix(err, fmt.Sprintf("failed to unmount zfs dataset %q", m.dataset.Name), 0)
	}
	return nil
}

func (m *Mounter) unloadKey() error {
	if m.dataset.KeyStatus != KeyStatusAvailable || m.dataset.Name != m.dataset.EncryptionRoot {
		return nil
	}

	perm, err := m.dataset.Permissions()
	if err != nil {
		return errors.WrapPrefix(err, "failed to unload zfs key", 0)
	}
	var run execFunc
	if perm.LoadKey {
		run = util.RunCommand
	} else {
		run = util.RunPrivilegedCommand
	}

	msg.Infof("zfs unload-key %q", m.dataset.Name)
	_, err = run(ZfsBin, nil, "unload-key", m.dataset.Name)
	if err != nil {
		return errors.WrapPrefix(err, "failed to unload zfs key", 0)
	}
	return nil
}

func (m *Mounter) Unmount() error {
	if err := m.unmount(); err != nil {
		return errors.Wrap(err, 0)
	}

	// check
	if err := m.refresh(); err != nil {
		return errors.WrapPrefix(err, "failed to check for result of unmounting", 0)
	}
	if m.dataset.Mounted {
		return errors.Errorf("zfs dataset %q is still mounted", m.dataset.Name)
	}

	if err := m.unloadKey(); err != nil {
		return errors.Wrap(err, 0)
	}
	return nil
}

func match(s string) ([]*mnt.Preset, error) {
	datasets, err := List("", true)
	if err != nil {
		return nil, errors.Wrap(err, 0)
	}

	var partialMatch []*Dataset
	for _, d := range datasets {
		if d.Name == s {
			return []*mnt.Preset{{
				Name: s,
				Type: "zfs",
				Path: d.Name,
			}}, nil
		}
		if strings.HasSuffix(d.Name, "/"+s) {
			partialMatch = append(partialMatch, d)
		}
	}
	var ret []*mnt.Preset
	for _, d := range partialMatch {
		ret = append(ret, &mnt.Preset{
			Name: s,
			Type: "zfs",
			Path: d.Name,
		})
	}
	return ret, nil
}

func init() {
	mnt.RegisterMounter("zfs", NewMounterFromPreset)
	mnt.RegisterMatcher(match)
}