package kvm2

import (
	"context"
	"errors"
	"io/fs"
	"net"
	"os"
	"strings"
	"testing"
	"time"

	"golang.org/x/crypto/ssh/knownhosts"

	"gitea.home.hrajfrisbee.cz/kacerr/egress-proxies-operator/internal/provider"
)

// sshFixture is one in-process server with a client key it accepts and a
// known_hosts file naming its host key.
type sshFixture struct {
	server *testSSHServer
	socket string
	cfg    provider.KVMConfig
	host   provider.KVMHost
}

func newSSHFixture(t *testing.T) *sshFixture {
	t.Helper()
	dir := t.TempDir()
	socket := startEchoSocket(t)
	keyPath, pub := writeKeyFile(t, dir, "id_ed25519")
	server := startSSHServer(t, pub, socket)
	return &sshFixture{
		server: server,
		socket: socket,
		cfg: provider.KVMConfig{
			SSHKeyFile:     keyPath,
			KnownHostsFile: writeKnownHosts(t, dir, "known_hosts", server.addr, server.hostKey),
		},
		host: provider.KVMHost{Name: "beelink", Address: server.addr, User: "virt", Socket: socket},
	}
}

const dialTestTimeout = 5 * time.Second

func TestDial_forwardsUnixSocket(t *testing.T) {
	t.Parallel()
	fx := newSSHFixture(t)
	for round := range 2 {
		sock, err := dialSSHSocket(context.Background(), fx.host, fx.cfg, dialTestTimeout)
		if err != nil {
			t.Fatalf("round %d: dial error = %v", round, err)
		}
		msg := []byte("ping over libvirt-sock")
		if _, err := sock.Write(msg); err != nil {
			t.Fatalf("round %d: write: %v", round, err)
		}
		buf := make([]byte, len(msg))
		_ = sock.SetReadDeadline(time.Now().Add(dialTestTimeout))
		if _, err := readFull(sock, buf); err != nil {
			t.Fatalf("round %d: read: %v", round, err)
		}
		if string(buf) != string(msg) {
			t.Errorf("round %d: echoed %q, want %q", round, buf, msg)
		}
		if err := sock.Close(); err != nil {
			t.Errorf("round %d: close error = %v", round, err)
		}
		if err := sock.Close(); err != nil {
			t.Errorf("round %d: second close error = %v, want nil", round, err)
		}
	}
	for _, p := range fx.server.requestedPaths() {
		if p != fx.socket {
			t.Errorf("server saw channel for %q, want %q", p, fx.socket)
		}
	}
}

func readFull(sock *sshSocket, buf []byte) (int, error) {
	n := 0
	for n < len(buf) {
		m, err := sock.Read(buf[n:])
		n += m
		if err != nil {
			return n, err
		}
	}
	return n, nil
}

func TestDial_unknownHostKeyRejected(t *testing.T) {
	t.Parallel()
	fx := newSSHFixture(t)
	fx.cfg.KnownHostsFile = writeKnownHosts(t, t.TempDir(), "known_hosts", "other.example:22", fx.server.hostKey)
	_, err := dialSSHSocket(context.Background(), fx.host, fx.cfg, dialTestTimeout)
	assertDialFailure(t, err, stageSSH, provider.ErrPermanent)
	var keyErr *knownhosts.KeyError
	if !errors.As(err, &keyErr) {
		t.Fatalf("error = %v, want *knownhosts.KeyError", err)
	}
	if len(keyErr.Want) != 0 {
		t.Errorf("KeyError.Want = %v, want empty (host unknown)", keyErr.Want)
	}
}

func TestDial_mismatchedHostKeyRejected(t *testing.T) {
	t.Parallel()
	fx := newSSHFixture(t)
	fx.cfg.KnownHostsFile = writeKnownHosts(t, t.TempDir(), "known_hosts", fx.server.addr, otherPublicKey(t))
	_, err := dialSSHSocket(context.Background(), fx.host, fx.cfg, dialTestTimeout)
	assertDialFailure(t, err, stageSSH, provider.ErrPermanent)
	var keyErr *knownhosts.KeyError
	if !errors.As(err, &keyErr) {
		t.Fatalf("error = %v, want *knownhosts.KeyError", err)
	}
	if len(keyErr.Want) == 0 {
		t.Error("KeyError.Want empty, want the mismatching known key")
	}
}

func TestDial_insecureIgnoreHostKeyAccepts(t *testing.T) {
	t.Parallel()
	fx := newSSHFixture(t)
	fx.cfg.KnownHostsFile = ""
	fx.cfg.InsecureIgnoreHostKey = true
	sock, err := dialSSHSocket(context.Background(), fx.host, fx.cfg, dialTestTimeout)
	if err != nil {
		t.Fatalf("dial error = %v, want nil with host keys ignored", err)
	}
	_ = sock.Close()
}

func TestDial_missingKeyFileIsTransient(t *testing.T) {
	t.Parallel()
	fx := newSSHFixture(t)
	fx.cfg.SSHKeyFile = fx.cfg.SSHKeyFile + ".missing"
	_, err := dialSSHSocket(context.Background(), fx.host, fx.cfg, dialTestTimeout)
	assertDialFailure(t, err, stageCredentials, provider.ErrTransient)
	if !errors.Is(err, fs.ErrNotExist) {
		t.Errorf("error = %v, want fs.ErrNotExist", err)
	}
	if !strings.Contains(err.Error(), fx.cfg.SSHKeyFile) {
		t.Errorf("error = %q, want the path in it", err)
	}
}

func TestDial_missingKnownHostsIsTransient(t *testing.T) {
	t.Parallel()
	fx := newSSHFixture(t)
	fx.cfg.KnownHostsFile = fx.cfg.KnownHostsFile + ".missing"
	_, err := dialSSHSocket(context.Background(), fx.host, fx.cfg, dialTestTimeout)
	assertDialFailure(t, err, stageCredentials, provider.ErrTransient)
	if !errors.Is(err, fs.ErrNotExist) {
		t.Errorf("error = %v, want fs.ErrNotExist", err)
	}
}

func TestDial_unparsableKeyIsPermanent(t *testing.T) {
	t.Parallel()
	fx := newSSHFixture(t)
	fx.cfg.SSHKeyFile = writeFile(t, t.TempDir(), "id_bad", "this is not a key\n")
	_, err := dialSSHSocket(context.Background(), fx.host, fx.cfg, dialTestTimeout)
	assertDialFailure(t, err, stageCredentials, provider.ErrPermanent)
	if !strings.Contains(err.Error(), "parsing ssh key file "+fx.cfg.SSHKeyFile) {
		t.Errorf("error = %q, want 'parsing ssh key file <path>'", err)
	}
}

func TestDial_unparsableKnownHostsIsPermanent(t *testing.T) {
	t.Parallel()
	fx := newSSHFixture(t)
	fx.cfg.KnownHostsFile = writeFile(t, t.TempDir(), "known_hosts", "host.example ssh-ed25519 not-base64!\n")
	_, err := dialSSHSocket(context.Background(), fx.host, fx.cfg, dialTestTimeout)
	assertDialFailure(t, err, stageCredentials, provider.ErrPermanent)
	if !strings.Contains(err.Error(), "parsing known_hosts file "+fx.cfg.KnownHostsFile) {
		t.Errorf("error = %q, want 'parsing known_hosts file <path>'", err)
	}
}

func TestDial_wrongKeyIsUnauthorized(t *testing.T) {
	t.Parallel()
	fx := newSSHFixture(t)
	fx.cfg.SSHKeyFile, _ = writeKeyFile(t, t.TempDir(), "id_other")
	_, err := dialSSHSocket(context.Background(), fx.host, fx.cfg, dialTestTimeout)
	assertDialFailure(t, err, stageSSH, provider.ErrPermanent)
	if !strings.Contains(err.Error(), sshAuthFailed) {
		t.Errorf("error = %q, want %q in it", err, sshAuthFailed)
	}
}

func TestDial_handshakeTimeout(t *testing.T) {
	t.Parallel()
	fx := newSSHFixture(t)
	fx.host.Address = silentListener(t)
	const budget = 200 * time.Millisecond
	start := time.Now()
	_, err := dialSSHSocket(context.Background(), fx.host, fx.cfg, budget)
	elapsed := time.Since(start)
	assertDialFailure(t, err, stageSSH, provider.ErrTransient)
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Errorf("error = %v, want context.DeadlineExceeded", err)
	}
	if elapsed > budget+time.Second {
		t.Errorf("dial took %v, want about the %v budget", elapsed, budget)
	}
}

func TestDial_cancelledContext(t *testing.T) {
	t.Parallel()
	fx := newSSHFixture(t)
	fx.host.Address = silentListener(t)
	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		time.Sleep(20 * time.Millisecond)
		cancel()
	}()
	start := time.Now()
	_, err := dialSSHSocket(ctx, fx.host, fx.cfg, dialTestTimeout)
	if time.Since(start) > time.Second {
		t.Errorf("dial took %v after cancel", time.Since(start))
	}
	assertDialFailure(t, err, stageSSH, provider.ErrTransient)
	if !errors.Is(err, context.Canceled) {
		t.Errorf("error = %v, want context.Canceled", err)
	}
}

func TestDial_refusedIsTransient(t *testing.T) {
	t.Parallel()
	fx := newSSHFixture(t)
	fx.host.Address = closedPort(t)
	_, err := dialSSHSocket(context.Background(), fx.host, fx.cfg, dialTestTimeout)
	assertDialFailure(t, err, stageTCP, provider.ErrTransient)
	var opErr *net.OpError
	if !errors.As(err, &opErr) {
		t.Errorf("error = %v, want *net.OpError", err)
	}
}

func TestDial_socketRefusedIsSocketStage(t *testing.T) {
	t.Parallel()
	fx := newSSHFixture(t)
	fx.host.Socket = "/run/nonexistent/libvirt-sock"
	_, err := dialSSHSocket(context.Background(), fx.host, fx.cfg, dialTestTimeout)
	assertDialFailure(t, err, stageSocket, provider.ErrTransient)
	if !strings.Contains(err.Error(), fx.host.Socket) {
		t.Errorf("error = %q, want the socket path in it", err)
	}
}

func TestSSH_keyMaterialNeverInErrorsOrLogs(t *testing.T) {
	t.Parallel()
	const keySentinel = "SENTINEL-KEY-4f9c2b"
	const hostsSentinel = "SENTINEL-HOSTS-8a1d7e"
	fx := newSSHFixture(t)
	dir := t.TempDir()
	tests := []struct {
		name string
		cfg  func(c *provider.KVMConfig)
	}{
		{
			name: "bad key file",
			cfg: func(c *provider.KVMConfig) {
				c.SSHKeyFile = writeFile(t, dir, "id_bad", "-----BEGIN NOT A KEY-----\n"+keySentinel+"\n-----END NOT A KEY-----\n")
			},
		},
		{
			name: "bad known_hosts file",
			cfg: func(c *provider.KVMConfig) {
				c.KnownHostsFile = writeFile(t, dir, "known_hosts_bad", "beelink ssh-ed25519 "+hostsSentinel+"\n")
			},
		},
		{
			name: "good key, host unknown",
			cfg: func(c *provider.KVMConfig) {
				c.KnownHostsFile = writeFile(t, dir, "known_hosts_other", knownhosts.Line([]string{"other.example:22"}, fx.server.hostKey)+"\n")
			},
		},
	}
	keyBytes := readFileT(t, fx.cfg.SSHKeyFile)
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			cfg := fx.cfg
			tc.cfg(&cfg)
			_, err := dialSSHSocket(context.Background(), fx.host, cfg, dialTestTimeout)
			if err == nil {
				t.Fatal("dial succeeded, want an error")
			}
			msg := err.Error()
			for _, s := range []string{keySentinel, hostsSentinel} {
				if strings.Contains(msg, s) {
					t.Errorf("error %q contains sentinel %q", msg, s)
				}
			}
			if strings.Contains(msg, keyBytes) || strings.Contains(msg, "PRIVATE KEY") {
				t.Errorf("error %q contains key material", msg)
			}
		})
	}
}

func readFileT(t *testing.T, path string) string {
	t.Helper()
	b, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("reading %s: %v", path, err)
	}
	return strings.TrimSpace(string(b))
}

func assertDialFailure(t *testing.T, err error, wantStage dialStage, wantClass error) {
	t.Helper()
	if err == nil {
		t.Fatal("dial succeeded, want an error")
	}
	stage, ok := stageOf(err)
	if !ok {
		t.Fatalf("error %v is not a *dialError", err)
	}
	if stage != wantStage {
		t.Errorf("stage = %q, want %q (error: %v)", stage, wantStage, err)
	}
	if got := classify("get", err); !errors.Is(got, wantClass) {
		t.Errorf("classify() = %v, want %v (error: %v)", got, wantClass, err)
	}
}

func TestDialError_unwrapsToCause(t *testing.T) {
	t.Parallel()
	cause := &knownhosts.KeyError{}
	err := &dialError{stage: stageSSH, host: "beelink", err: cause}
	if !errors.Is(err, cause) {
		t.Error("errors.Is(dialError, cause) = false")
	}
	if !strings.Contains(err.Error(), "beelink") || !strings.Contains(err.Error(), "(ssh)") {
		t.Errorf("Error() = %q, want the host and the stage", err)
	}
	if _, ok := stageOf(errors.New("plain")); ok {
		t.Error("stageOf(plain error) reported a stage")
	}
}