package remote

import (
	"bytes"
	"context"
	"embed"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"runtime"
	"slices"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/klauspost/compress/zstd"
	"github.com/samber/lo"
	"tobi/internal/session"
	"tobi/internal/target"
)

//go:embed payloads/*
var embeddedPayloads embed.FS

var hostMu sync.Mutex

func hostMapPath() string {
	base, _ := os.UserConfigDir()
	if base == "" {
		return ""
	}
	return filepath.Join(base, "tobi", "hosts.json")
}

func loadHosts() map[string][]string {
	data, _ := os.ReadFile(hostMapPath())
	var m map[string][]string
	if json.Unmarshal(data, &m) == nil && m != nil {
		return m
	}
	var old map[string]string
	_ = json.Unmarshal(data, &old)
	m = make(map[string][]string, len(old))
	for host, dest := range old {
		if dest != "" {
			m[host] = []string{dest}
		}
	}
	return m
}

func remembered(host string) []string {
	hostMu.Lock()
	defer hostMu.Unlock()
	return loadHosts()[host]
}

func peerEnv() string {
	u := os.Getenv("USER")
	if u == "" {
		if cur, err := user.Current(); err == nil {
			u = cur.Username
		}
	}
	host := session.ShortHostname()
	if u == "" || host == "" || host == "localhost" {
		return ""
	}
	return "TOBI_PEER=" + u + "@" + host + " "
}

func RememberHost(dest string) {
	t := target.Parse(dest)
	if t.User == "" || t.Host == "" {
		return
	}
	hostMu.Lock()
	defer hostMu.Unlock()
	m := loadHosts()
	if len(m[t.Host]) > 0 && m[t.Host][0] == dest {
		return
	}
	m[t.Host] = append([]string{dest}, slices.DeleteFunc(m[t.Host], func(d string) bool { return d == dest })...)
	m[t.Host] = m[t.Host][:min(len(m[t.Host]), 8)]
	saveHosts(m)
}

func saveHosts(m map[string][]string) {
	p := hostMapPath()
	data, err := json.Marshal(m)
	if p == "" || err != nil || os.MkdirAll(filepath.Dir(p), 0700) != nil {
		return
	}
	_ = os.WriteFile(p, data, 0600)
}

func hostDest(host string) string {
	if target.Parse(host).User != "" {
		RememberHost(host)
		return host
	}
	if d := remembered(host); len(d) > 0 {
		return d[0]
	}
	return host
}

func KnownDests(host string) []string {
	return lo.Uniq(append([]string{host}, slices.DeleteFunc(remembered(host), func(d string) bool { return d == "" || d == host })...))
}

func ProbeDests(host string) []string {
	d := KnownDests(host)
	return lo.Ternary(len(d) > 1, d[1:], []string{host})
}

func QueryHost(ctx context.Context, host string) []session.Session {
	t := target.Parse(host)
	dests := ProbeDests(t.Host)
	if t.User != "" {
		dests = lo.Uniq(append([]string{host}, dests...))
	}
	var found []session.Session
	for _, dest := range dests {
		if list, err := QueryRemote(ctx, dest); err == nil {
			found = append(found, list...)
		}
	}
	return lo.UniqBy(found, func(s session.Session) string {
		return strings.ToLower(target.Parse(s.Host).Host + "." + s.ID)
	})
}

const (
	uploadMarker  = "__TOBI_UPLOADED__"
	spawnMarker   = "__TOBI_SPAWNED__"
	killMarker    = "__TOBI_KILLED__"
	remotePathEnv = `PATH="$HOME/.local/bin:$PATH:$HOME/.nix-profile/bin:/opt/homebrew/bin:/usr/local/bin:$HOME/go/bin"`
	uploadCmd     = "mkdir -p ~/.local/bin && cat > ~/.local/bin/tobi.tmp && chmod 755 ~/.local/bin/tobi.tmp && mv -f ~/.local/bin/tobi.tmp ~/.local/bin/tobi"
)

func rpcCommand(op, id string) string {
	cmd := remotePathEnv + " " + peerEnv() + "TOBI_RPC=" + op
	if id != "" {
		cmd += " TOBI_ID=" + id
	}
	return cmd + " tobi"
}

var archAlias = map[string]string{"x86_64": "amd64", "aarch64": "arm64"}

var (
	remoteInstalled sync.Map
	// the declared release; bump it (or append "+n" for a same-day
	// re-release) to ship, remotes refuse older builds
	Version = "2026.10.21"
)

// verTuple parses a declared version - "2026.09.21" or "2026.09.21+2", the
// suffix counting same-day re-releases; undated builds sort before dated ones
func verTuple(v string) (string, int) {
	if len(v) >= 10 && v[4] == '.' && v[7] == '.' {
		if len(v) == 10 {
			return v, 0
		}
		if v[10] == '+' {
			if n, err := strconv.Atoi(v[11:]); err == nil && n >= 0 {
				return v[:10], n
			}
		}
	}
	return "", 0
}

type sshOpts struct {
	cmd           string
	batch, master bool
	timeout       string
}

func sshCommand(ctx context.Context, dest string, opts sshOpts) *exec.Cmd {
	bin, err := exec.LookPath("tssh")
	if err != nil {
		bin = "ssh"
	}
	args := []string{"-T", "-o", "ConnectTimeout=" + opts.timeout}
	if opts.batch {
		args = append(args, "-o", "BatchMode=yes")
	}
	if home, err := os.UserHomeDir(); err == nil && home != "" {
		ctl := filepath.Join(home, ".ssh", "tobi-control-%r@%h:%p")
		args = append(args, "-o", "ControlMaster="+lo.Ternary(opts.master, "auto", "no"),
			"-o", "ControlPath="+ctl, "-o", "ControlPersist=60")
	}
	c := exec.CommandContext(ctx, bin, append(args, dest, opts.cmd)...)
	c.WaitDelay = 500 * time.Millisecond
	if opts.batch {
		// no controlling terminal: tssh ignores BatchMode and prompts for
		// unknown host keys on /dev/tty, eating the picker's keystrokes.
		// the new session is also a process group, so a timeout takes the
		// ProxyCommand down with ssh instead of orphaning it
		c.SysProcAttr = &syscall.SysProcAttr{Setsid: true}
		c.Cancel = func() error { return syscall.Kill(-c.Process.Pid, syscall.SIGKILL) }
	}
	return c
}

var ErrUnreachable = errors.New("unreachable")

func sshRun(ctx context.Context, dest string, in io.Reader, opts sshOpts) (string, error) {
	c := sshCommand(ctx, dest, opts)
	c.Stdin = in
	var out, errBuf bytes.Buffer
	c.Stdout, c.Stderr = &out, &errBuf
	if err := c.Run(); err != nil {
		var ee *exec.ExitError
		if (errors.As(err, &ee) && ee.ExitCode() == 255) || ctx.Err() != nil {
			err = fmt.Errorf("%w: %w", ErrUnreachable, err)
		}
		if msg := strings.TrimSpace(errBuf.String()); msg != "" {
			return out.String(), fmt.Errorf("%w: %s", err, msg)
		}
		return out.String(), err
	}
	return out.String(), nil
}

// sshConfirmed runs opts.cmd and wants marker in its output: a remote that
// exits clean without printing it did not do the work
func sshConfirmed(ctx context.Context, dest string, in io.Reader, opts sshOpts, marker, missing string) error {
	out, err := sshRun(ctx, dest, in, opts)
	if err == nil && !strings.Contains(out, marker) {
		return errors.New(missing)
	}
	return err
}

func getBinaryPayload(goos, arch string) ([]byte, error) {
	if data, err := embeddedPayloads.ReadFile(fmt.Sprintf("payloads/tobi-%s-%s.zst", goos, arch)); err == nil && len(data) > 0 {
		if zr, err := zstd.NewReader(bytes.NewReader(data)); err == nil {
			defer zr.Close()
			return io.ReadAll(zr)
		}
	}
	if goos == runtime.GOOS && arch == runtime.GOARCH {
		return os.ReadFile(session.Executable())
	}
	return nil, fmt.Errorf("no binary available for %s/%s (build the carrier slices: nix build or scripts/build-universal.sh)", goos, arch)
}

func EnsureRemote(ctx context.Context, host string) error {
	dest := hostDest(host)
	if _, ok := remoteInstalled.Load(dest); ok {
		return nil
	}
	// the PATH prefix is a per-command assignment: it must wrap the version
	// rpc too, or the check 127s on hosts where tobi is not in default PATH
	check := remotePathEnv + " " + peerEnv() + `command -v tobi >/dev/null 2>&1 && ` + remotePathEnv + ` TOBI_RPC=version tobi`
	if out, err := sshRun(ctx, dest, nil, sshOpts{cmd: check, batch: true, master: true, timeout: "3"}); err == nil {
		remoteVer := strings.TrimSpace(out)
		ldate, lseq := verTuple(Version)
		rdate, rseq := verTuple(remoteVer)
		newer := rdate != "" && (rdate > ldate || (rdate == ldate && rseq > lseq))
		if ldate == "" || rdate != "" && (rdate > ldate || rdate == ldate && rseq >= lseq) {
			if ldate == "" && rdate != "" {
				fmt.Fprintf(os.Stderr, "not downgrading %s: remote tobi %s is versioned, the local build is not\n", dest, remoteVer)
			} else if newer {
				fmt.Fprintf(os.Stderr, "not downgrading %s: remote tobi %s is newer than local %s\n", dest, remoteVer, Version)
			}
			remoteInstalled.Store(dest, true)
			return nil
		}
	}
	out, err := sshRun(ctx, dest, nil, sshOpts{cmd: "uname -s && uname -m", batch: true, timeout: "3"})
	if err != nil {
		return fmt.Errorf("detect %s: %w", dest, err)
	}
	osName, arch, ok := strings.Cut(strings.TrimSpace(out), "\n")
	if !ok {
		return fmt.Errorf("unexpected uname from %s: %q", dest, out)
	}
	osName, arch = strings.ToLower(osName), strings.ToLower(arch)
	if a := archAlias[arch]; a != "" {
		arch = a
	}
	bin, err := getBinaryPayload(osName, arch)
	if err != nil {
		return fmt.Errorf("prepare tobi for %s (%s/%s): %w", dest, osName, arch, err)
	}
	upload := sshOpts{cmd: uploadCmd + " && echo " + uploadMarker, batch: true, timeout: "5"}
	if err := sshConfirmed(ctx, dest, bytes.NewReader(bin), upload, uploadMarker, "upload did not complete"); err != nil {
		return fmt.Errorf("upload tobi to %s: %w", dest, err)
	}
	remoteInstalled.Store(dest, true)
	return nil
}

func QueryRemote(ctx context.Context, host string) ([]session.Session, error) {
	dest := hostDest(host)
	out, err := sshRun(ctx, dest, nil, sshOpts{cmd: rpcCommand("list", ""), batch: true, timeout: "10"})
	if err != nil {
		return nil, err
	}
	var list []session.Session
	if err := json.Unmarshal([]byte(strings.TrimSpace(out)), &list); err != nil {
		return nil, err
	}
	for i := range list {
		list[i].Host = dest
	}
	return list, nil
}

func AttachRemote(host, id string) (io.ReadWriteCloser, error) {
	ctx := context.Background()
	dest := hostDest(host)
	if err := EnsureRemote(ctx, dest); err != nil {
		return nil, err
	}
	c := sshCommand(ctx, dest, sshOpts{cmd: rpcCommand("attach", id), batch: true, timeout: "2"})
	in, err := c.StdinPipe()
	if err != nil {
		return nil, err
	}
	out, err := c.StdoutPipe()
	if err != nil {
		return nil, err
	}
	c.Stderr = nil // unreachable peers are expected; their stderr is noise
	if err := c.Start(); err != nil {
		return nil, err
	}
	return &Stream{Reader: out, WriteCloser: in, cmd: c}, nil
}

type Stream struct {
	io.Reader
	io.WriteCloser
	cmd  *exec.Cmd
	once sync.Once
}

func (s *Stream) Close() error {
	s.once.Do(func() {
		_ = s.WriteCloser.Close()
		_ = s.cmd.Wait()
	})
	return nil
}

type tsPeer struct {
	Online                bool
	OS, DNSName, HostName string
}

type Peer struct {
	Name   string
	Online bool
}

func DiscoverPeers(ctx context.Context) []Peer {
	c := exec.CommandContext(ctx, "tailscale", "status", "--json")
	c.WaitDelay = 500 * time.Millisecond
	out, err := c.Output()
	if err != nil {
		return nil
	}
	var st struct{ Peer map[string]tsPeer }
	if json.Unmarshal(out, &st) != nil {
		return nil
	}
	self := strings.ToLower(session.ShortHostname())
	peers := lo.FilterMap(lo.Values(st.Peer), func(p tsPeer, _ int) (Peer, bool) {
		if p.OS != "linux" && p.OS != "macOS" && p.OS != "darwin" {
			return Peer{}, false
		}
		name, _, _ := strings.Cut(p.DNSName, ".")
		if name == "" {
			name = p.HostName
		}
		return Peer{Name: name, Online: p.Online}, name != "" && !strings.EqualFold(name, self)
	})
	slices.SortFunc(peers, func(a, b Peer) int { return strings.Compare(a.Name, b.Name) })
	return peers
}

func RunRPC(op, id string) error {
	if t := target.Parse(os.Getenv("TOBI_PEER")); t.User != "" && t.Host != "" && !strings.EqualFold(t.Host, session.ShortHostname()) {
		RememberHost(t.Dest())
	}
	switch op {
	case "list":
		return json.NewEncoder(os.Stdout).Encode(session.ListLocal())
	case "spawn":
		return runSpawnRPC(os.Stdin)
	case "version":
		fmt.Println(Version)
		return nil
	case "kill":
		if err := session.Kill(id); err != nil {
			return err
		}
		fmt.Println(killMarker)
		return nil
	}
	conn, err := net.Dial("unix", session.SocketPath(id))
	if err != nil {
		return err
	}
	defer conn.Close()
	done := make(chan struct{}, 2)
	pump := func(dst io.Writer, src io.Reader) {
		_, _ = io.Copy(dst, src)
		_ = conn.Close()
		done <- struct{}{}
	}
	go pump(conn, os.Stdin)
	go pump(os.Stdout, conn)
	<-done
	return nil
}

type spawnRequest struct {
	ID   string   `json:"id"`
	Argv []string `json:"argv"`
}

func validateSpawnRequest(req spawnRequest) error {
	if !session.ValidateID(req.ID) {
		return fmt.Errorf("invalid session id %q", req.ID)
	}
	if len(req.Argv) == 0 || len(req.Argv) > 128 {
		return fmt.Errorf("invalid argv length %d", len(req.Argv))
	}
	return nil
}

func runSpawnRPC(r io.Reader) error {
	var req spawnRequest
	if err := json.NewDecoder(io.LimitReader(r, 64<<10)).Decode(&req); err != nil {
		return fmt.Errorf("spawn request: %w", err)
	}
	if err := validateSpawnRequest(req); err != nil {
		return err
	}
	return session.EnsureOwner(req.ID, req.Argv)
}

// an older remote tobi takes "kill" for an attach and exits clean, so only the marker proves it
func KillRemote(ctx context.Context, host, id string) error {
	dest := hostDest(host)
	kill := sshOpts{cmd: rpcCommand("kill", id), batch: true, timeout: "5", master: true}
	if err := sshConfirmed(ctx, dest, nil, kill, killMarker, "remote tobi is too old to kill sessions"); err != nil {
		return fmt.Errorf("kill %s.%s: %w", dest, id, err)
	}
	return nil
}

func SpawnRemote(host, id string, cmdArgs []string) error {
	req := spawnRequest{ID: id, Argv: cmdArgs}
	if err := validateSpawnRequest(req); err != nil {
		return err
	}
	dest := hostDest(host)
	if err := EnsureRemote(context.Background(), dest); err != nil {
		return err
	}
	body, _ := json.Marshal(req) // strings only, it cannot fail
	spawn := sshOpts{cmd: rpcCommand("spawn", "") + " && echo " + spawnMarker, timeout: "5", master: true}
	if err := sshConfirmed(context.Background(), dest, bytes.NewReader(append(body, '\n')), spawn, spawnMarker, "did not confirm"); err != nil {
		return fmt.Errorf("spawn on %s: %w", dest, err)
	}
	return nil
}