// ABOUTME: CLI tool that mounts remote filesystems via sshfs based on
// ABOUTME: ~/.ssh/config entries, with optional editor launch on the mountpoint.

package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"io/fs"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"github.com/kevinburke/ssh_config"
	"github.com/moby/sys/mountinfo"
)

var (
	rxHostUser     = regexp.MustCompile(`^[A-Za-z0-9_][A-Za-z0-9._-]*$`)
	rxIdentityFile = regexp.MustCompile(`^[A-Za-z0-9/~][A-Za-z0-9._@/:+=~%-]*$`)
	rxRemoteDir    = regexp.MustCompile(`^[A-Za-z0-9/~][A-Za-z0-9._/~-]*$`)
)

var (
	eFlag = flag.Bool("e", false, "open mountpoint in your editor")
	vFlag = flag.Bool("v", false, "verbose: print resolved ssh_config fields")
	lFlag = flag.Bool("l", false, "list active mounts and exit")
	uFlag = flag.Bool("u", false, "unmount the given <Host> and exit")
	rDir  string
)

func init() {
	flag.StringVar(&rDir, "r", "", "remote directory to mount (default: remote home)")
	flag.StringVar(&rDir, "remote-dir", "", "remote directory to mount (default: remote home)")
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr,
			"usage: sshfsc [flags] <Host>\n\n"+
				"Mount a remote home directory via sshfs based on ~/.ssh/config entries.\n"+
				"Use -l to list active mounts; -u <Host> to unmount.\n\n"+
				"Flags:\n")
		flag.PrintDefaults()
	}
	flag.Parse()
	args := flag.Args()

	if *lFlag && *uFlag {
		fmt.Fprintln(os.Stderr, "-l and -u are mutually exclusive")
		os.Exit(2)
	}

	if *lFlag {
		if err := list_mounts(os.Stdout); err != nil {
			fmt.Fprintln(os.Stderr, "list_mounts() failed:", err)
			os.Exit(8)
		}
		return
	}

	if *uFlag {
		if len(args) == 0 {
			fmt.Fprintln(os.Stderr, "-u requires <Host> argument")
			os.Exit(2)
		}
		if err := unmount_sshfs(args[0]); err != nil {
			fmt.Fprintln(os.Stderr, "unmount_sshfs() failed:", err)
			os.Exit(9)
		}
		return
	}

	if len(args) == 0 {
		fmt.Fprintln(os.Stderr, "No hostname specified.")
		os.Exit(2)
	}

	alias := args[0]
	if err := validate_ssh_field("alias", alias, rxHostUser); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}

	hostname, err := lookup_ssh_field(args[0], "HostName")
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}
	user, err := lookup_ssh_field(args[0], "User")
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}
	port, err := lookup_ssh_field(args[0], "Port")
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}
	ifile, err := lookup_ssh_field(args[0], "IdentityFile")
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}

	if hostname == "" || user == "" || ifile == "" {
		var missing []string
		if hostname == "" {
			missing = append(missing, "HostName")
		}
		if user == "" {
			missing = append(missing, "User")
		}
		if ifile == "" {
			missing = append(missing, "IdentityFile")
		}
		fmt.Fprintf(os.Stderr, "ssh config %q missing required field(s): %s\n",
			args[0], strings.Join(missing, ", "))
		os.Exit(3)
	}

	if err := validate_ssh_field("HostName", hostname, rxHostUser); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}
	if err := validate_ssh_field("User", user, rxHostUser); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}
	if err := validate_ssh_field("IdentityFile", ifile, rxIdentityFile); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}
	if p, perr := strconv.Atoi(port); perr != nil || p < 1 || p > 65535 {
		fmt.Fprintf(os.Stderr, "ssh config Port %q is not a valid port number\n", port)
		os.Exit(3)
	}
	if rDir != "" {
		if err := validate_ssh_field("remote-dir", rDir, rxRemoteDir); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(3)
		}
	}

	mount, err := verify_mount_dir(alias)
	if err != nil {
		fmt.Fprintln(os.Stderr, "verify_mount_dir() failed:", err)
		os.Exit(4)
	}

	if *vFlag {
		fmt.Println("Hostname:	", hostname)
		fmt.Println("User:		", user)
		fmt.Println("Port:		", port)
		fmt.Println("Ifile:		", ifile)
		fmt.Println("Mount:		", mount)
		if rDir != "" {
			fmt.Println("Remote:		", rDir)
		}
		fmt.Println("---")
	} else {
		fmt.Println("Mount:	", mount)
	}

	chkmount, chkmount_err := is_mounted_at(mount)
	if chkmount_err != nil {
		fmt.Fprintf(os.Stderr, "is_mounted_at() failed with %s\n", chkmount_err)
		os.Exit(5)
	}
	if !chkmount {
		if err := mount_sshfs(alias, user, ifile, port, mount, rDir); err != nil {
			fmt.Fprintln(os.Stderr, "mount_sshfs() failed:", err)
			os.Exit(6)
		}
	} else {
		fmt.Println("!!! Already mounted")
	}
	// run_editor fires on -e regardless of mount state: re-invoking with -e
	// re-opens the editor on an already-mounted server.
	if err := run_editor(mount); err != nil {
		fmt.Fprintln(os.Stderr, "run_editor() failed:", err)
		os.Exit(7)
	}
}

func lookup_ssh_field(alias, field string) (string, error) {
	v, err := ssh_config.GetStrict(alias, field)
	if err != nil {
		return "", fmt.Errorf("parse ~/.ssh/config: %w", err)
	}
	return v, nil
}

func validate_ssh_field(name, value string, allow *regexp.Regexp) error {
	if !allow.MatchString(value) {
		return fmt.Errorf("ssh config %s %q contains disallowed characters", name, value)
	}
	return nil
}

func run_editor(mount string) error {
	if !*eFlag {
		return nil
	}
	cmd := exec.Command("subl", mount)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Start()
}

// mount_base resolves the per-user sshfs mount base under $XDG_RUNTIME_DIR.
// When create is true, the base directory is mkdir'd; otherwise a missing
// base bubbles up as fs.ErrNotExist so callers can decide how to react.
func mount_base(create bool) (string, error) {
	runtime := os.Getenv("XDG_RUNTIME_DIR")
	if runtime == "" {
		runtime = filepath.Join("/run/user", strconv.Itoa(os.Getuid()))
	}
	base := filepath.Join(runtime, "sshfs")
	if create {
		if err := os.MkdirAll(base, 0700); err != nil {
			return "", fmt.Errorf("create base %q: %w", base, err)
		}
	} else if _, err := os.Stat(base); err != nil {
		return "", err
	}
	realBase, err := filepath.EvalSymlinks(base)
	if err != nil {
		return "", fmt.Errorf("resolve base %q: %w", base, err)
	}
	return realBase, nil
}

// mount_path joins base + name with traversal and symlink guards.
// It does not create the mountpoint directory.
func mount_path(base, name string) (string, error) {
	mount := filepath.Clean(filepath.Join(base, name))
	if !strings.HasPrefix(mount, base+string(os.PathSeparator)) {
		return "", fmt.Errorf("name %q escapes mount base %q", name, base)
	}
	if info, err := os.Lstat(mount); err == nil && info.Mode()&os.ModeSymlink != 0 {
		return "", fmt.Errorf("mount path %q is a symlink", mount)
	}
	return mount, nil
}

func verify_mount_dir(name string) (string, error) {
	base, err := mount_base(true)
	if err != nil {
		return "", err
	}
	mount, err := mount_path(base, name)
	if err != nil {
		return "", err
	}
	// If the path is already a mountpoint (possibly stale), skip MkdirAll —
	// it would stat the target, which fails with EIO on stale FUSE mounts.
	if ok, _ := is_mounted_at(mount); ok {
		return mount, nil
	}
	if err := os.MkdirAll(mount, 0700); err != nil {
		return "", fmt.Errorf("create mount dir %q: %w", mount, err)
	}
	return mount, nil
}

// is_mounted_at reports whether path is a mountpoint per /proc/self/mountinfo.
// Unlike mountinfo.Mounted, this does not stat the path, so stale FUSE mounts
// (whose targets return EIO on stat) are still detected.
func is_mounted_at(path string) (bool, error) {
	mounts, err := mountinfo.GetMounts(mountinfo.SingleEntryFilter(path))
	if err != nil {
		return false, err
	}
	return len(mounts) > 0, nil
}

func list_mounts(out io.Writer) error {
	base, err := mount_base(false)
	if err != nil {
		if errors.Is(err, fs.ErrNotExist) {
			return nil
		}
		return err
	}
	mounts, err := mountinfo.GetMounts(mountinfo.PrefixFilter(base))
	if err != nil {
		return fmt.Errorf("read mountinfo: %w", err)
	}
	prefix := base + string(os.PathSeparator)
	for _, m := range mounts {
		rel := strings.TrimPrefix(m.Mountpoint, prefix)
		if rel == "" || strings.Contains(rel, string(os.PathSeparator)) {
			continue
		}
		fmt.Fprintln(out, rel)
	}
	return nil
}

func run_fusermount(flag, mount string) error {
	cmd := exec.Command("fusermount", flag, mount)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func unmount_sshfs(alias string) error {
	if err := validate_ssh_field("alias", alias, rxHostUser); err != nil {
		return err
	}
	base, err := mount_base(false)
	if err != nil {
		if errors.Is(err, fs.ErrNotExist) {
			return fmt.Errorf("not mounted: %s", alias)
		}
		return err
	}
	mount, err := mount_path(base, alias)
	if err != nil {
		return err
	}
	ok, err := is_mounted_at(mount)
	if err != nil {
		return fmt.Errorf("check mount %q: %w", mount, err)
	}
	if !ok {
		return fmt.Errorf("not mounted: %s", alias)
	}
	if err := run_fusermount("-u", mount); err != nil {
		// stale FUSE mounts often need lazy unmount; -u fails with
		// "Transport endpoint is not connected" on those.
		if err2 := run_fusermount("-uz", mount); err2 != nil {
			return fmt.Errorf("fusermount -u %q: %w (lazy retry: %v)", mount, err, err2)
		}
	}
	if err := os.Remove(mount); err != nil {
		fmt.Fprintf(os.Stderr, "warning: remove empty mount dir %q: %v\n", mount, err)
	}
	fmt.Println("Unmounted:", alias)
	return nil
}

func mount_sshfs(alias string, user string, ifile string, port string, mount string, remoteDir string) error {
	cmd := exec.Command("sshfs", "-p", port,
		"-o", "IdentityFile="+ifile,
		"-o", "idmap=user",
		"-o", "cache=yes",
		"-o", "auto_cache",
		"-o", "attr_timeout=60",
		"-o", "entry_timeout=60",
		"-o", "negative_timeout=20",
		"-o", "Ciphers=aes128-gcm@openssh.com",
		"-o", "Compression=no",
		"-o", "reconnect",
		"-o", "StrictHostKeyChecking=accept-new",
		"-o", "ServerAliveInterval=15",
		"-o", "ServerAliveCountMax=3",
		user+"@"+alias+":"+remoteDir, mount)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}