package kvm

import (
	"bytes"
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"encoding/pem"
	"errors"
	"io"
	"net"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	libvirt "github.com/digitalocean/go-libvirt"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/knownhosts"

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

// --- hostPool over a fake dialer ---

type fakeClient struct {
	done   chan struct{}
	closed atomic.Bool
}

func (c *fakeClient) close() {
	if c.closed.CompareAndSwap(false, true) {
		close(c.done)
	}
}

type fakeDialer struct {
	mu      sync.Mutex
	dials   int
	err     error
	delay   time.Duration
	clients []*fakeClient
}

func (d *fakeDialer) dial(ctx context.Context, _ provider.KVMHost) (*hostClient, error) {
	if d.delay > 0 {
		select {
		case <-time.After(d.delay):
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	d.dials++
	if d.err != nil {
		return nil, d.err
	}
	fc := &fakeClient{done: make(chan struct{})}
	d.clients = append(d.clients, fc)
	return &hostClient{api: nil, done: fc.done, close: fc.close}, nil
}

func (d *fakeDialer) dialCount() int {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.dials
}

var testHost = provider.KVMHost{Name: "nas1", Address: "nas1:22", User: "virt", Socket: defaultSocket, Pool: "default", Bridge: "br0"}

func noop(libvirtAPI) error { return nil }

func TestHostPool_call_reusesConnection(t *testing.T) {
	t.Parallel()
	d := &fakeDialer{}
	p := newHostPool(d.dial)
	for range 3 {
		if err := p.call(context.Background(), testHost, noop); err != nil {
			t.Fatalf("call() error = %v", err)
		}
	}
	if d.dialCount() != 1 {
		t.Errorf("dials = %d, want 1", d.dialCount())
	}
}

func TestHostPool_call_redialsAfterDisconnect(t *testing.T) {
	t.Parallel()
	d := &fakeDialer{}
	p := newHostPool(d.dial)
	if err := p.call(context.Background(), testHost, noop); err != nil {
		t.Fatalf("call() error = %v", err)
	}
	d.clients[0].close()
	if err := p.call(context.Background(), testHost, noop); err != nil {
		t.Fatalf("second call() error = %v", err)
	}
	if d.dialCount() != 2 {
		t.Errorf("dials = %d, want 2 after the first connection dropped", d.dialCount())
	}
}

func TestHostPool_call_transportErrorDropsConnection(t *testing.T) {
	t.Parallel()
	d := &fakeDialer{}
	p := newHostPool(d.dial)
	err := p.call(context.Background(), testHost, func(libvirtAPI) error { return io.ErrUnexpectedEOF })
	if !errors.Is(err, io.ErrUnexpectedEOF) {
		t.Fatalf("call() error = %v, want the transport error passed through", err)
	}
	if !d.clients[0].closed.Load() {
		t.Error("connection not closed after a transport error")
	}
	if err := p.call(context.Background(), testHost, noop); err != nil {
		t.Fatalf("call() after transport error = %v", err)
	}
	if d.dialCount() != 2 {
		t.Errorf("dials = %d, want 2", d.dialCount())
	}
}

func TestHostPool_call_libvirtErrorKeepsConnection(t *testing.T) {
	t.Parallel()
	d := &fakeDialer{}
	p := newHostPool(d.dial)
	err := p.call(context.Background(), testHost, func(libvirtAPI) error { return lvErr(libvirt.ErrNoDomain) })
	if !hasCode(err, libvirt.ErrNoDomain) {
		t.Fatalf("call() error = %v, want the libvirt error passed through", err)
	}
	if err := p.call(context.Background(), testHost, noop); err != nil {
		t.Fatalf("call() error = %v", err)
	}
	if d.dialCount() != 1 {
		t.Errorf("dials = %d, want 1: a daemon-reported error must not cost the connection", d.dialCount())
	}
	if d.clients[0].closed.Load() {
		t.Error("connection closed after a libvirt error")
	}
}

func TestHostPool_call_dialErrorPropagates(t *testing.T) {
	t.Parallel()
	d := &fakeDialer{err: errors.New("ssh: handshake failed")}
	p := newHostPool(d.dial)
	err := p.call(context.Background(), testHost, noop)
	if err == nil || !strings.Contains(err.Error(), "nas1") || !strings.Contains(err.Error(), "handshake failed") {
		t.Errorf("call() error = %v, want the dial error naming the host", err)
	}
}

func TestHostPool_call_cancelAbandonsBlockedCall(t *testing.T) {
	t.Parallel()
	d := &fakeDialer{}
	p := newHostPool(d.dial)
	ctx, cancel := context.WithCancel(context.Background())
	released := make(chan struct{})
	go func() {
		time.Sleep(20 * time.Millisecond)
		cancel()
	}()
	err := p.call(ctx, testHost, func(libvirtAPI) error {
		// Simulates an RPC stuck on a hung hypervisor: it only returns once
		// the pool closes the connection out from under it.
		<-d.clients[0].done
		close(released)
		return io.EOF
	})
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("call() error = %v, want context.Canceled", err)
	}
	select {
	case <-released:
	case <-time.After(time.Second):
		t.Fatal("blocked call was not released by closing the connection")
	}
	if err := p.call(context.Background(), testHost, noop); err != nil {
		t.Fatalf("call() after cancel = %v", err)
	}
	if d.dialCount() != 2 {
		t.Errorf("dials = %d, want 2", d.dialCount())
	}
}

func TestHostPool_call_concurrentCallersShareOneDial(t *testing.T) {
	t.Parallel()
	d := &fakeDialer{delay: 30 * time.Millisecond}
	p := newHostPool(d.dial)
	var wg sync.WaitGroup
	for range 8 {
		wg.Go(func() {
			if err := p.call(context.Background(), testHost, noop); err != nil {
				t.Errorf("call() error = %v", err)
			}
		})
	}
	wg.Wait()
	if d.dialCount() != 1 {
		t.Errorf("dials = %d, want 1", d.dialCount())
	}
}

func TestHostPool_call_hostsAreIndependent(t *testing.T) {
	t.Parallel()
	d := &fakeDialer{}
	p := newHostPool(d.dial)
	other := testHost
	other.Name = "lab2"
	for _, h := range []provider.KVMHost{testHost, other, testHost, other} {
		if err := p.call(context.Background(), h, noop); err != nil {
			t.Fatalf("call(%s) error = %v", h.Name, err)
		}
	}
	if d.dialCount() != 2 {
		t.Errorf("dials = %d, want one per host", d.dialCount())
	}
}

// --- SSH dial against an in-process server ---

func newKeyPair(t *testing.T) (ssh.Signer, []byte) {
	t.Helper()
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	signer, err := ssh.NewSignerFromKey(priv)
	if err != nil {
		t.Fatal(err)
	}
	block, err := ssh.MarshalPrivateKey(priv, "")
	if err != nil {
		t.Fatal(err)
	}
	return signer, pem.EncodeToMemory(block)
}

// testSSHServer accepts one public key and serves
// direct-streamlocal@openssh.com channels by dialing the requested unix
// path, which is exactly what OpenSSH does for ssh.Client.Dial("unix").
type testSSHServer struct {
	addr    string
	hostKey ssh.PublicKey

	mu    sync.Mutex
	paths []string
}

func startSSHServer(t *testing.T, clientKey ssh.PublicKey, answerGlobalRequests bool) *testSSHServer {
	t.Helper()
	hostSigner, _ := newKeyPair(t)
	cfg := &ssh.ServerConfig{
		PublicKeyCallback: func(_ ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			if bytes.Equal(key.Marshal(), clientKey.Marshal()) {
				return nil, nil
			}
			return nil, errors.New("unknown public key")
		},
	}
	cfg.AddHostKey(hostSigner)
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { ln.Close() })
	s := &testSSHServer{addr: ln.Addr().String(), hostKey: hostSigner.PublicKey()}
	go func() {
		for {
			nc, err := ln.Accept()
			if err != nil {
				return
			}
			go s.serve(nc, cfg, answerGlobalRequests)
		}
	}()
	return s
}

func (s *testSSHServer) serve(nc net.Conn, cfg *ssh.ServerConfig, answerGlobalRequests bool) {
	_, chans, reqs, err := ssh.NewServerConn(nc, cfg)
	if err != nil {
		return
	}
	if answerGlobalRequests {
		go ssh.DiscardRequests(reqs)
	} else {
		go func() {
			for range reqs { //nolint:revive // deliberately never replies, to stall keepalives
			}
		}()
	}
	for newCh := range chans {
		if newCh.ChannelType() != "direct-streamlocal@openssh.com" {
			_ = newCh.Reject(ssh.UnknownChannelType, "test server only forwards unix sockets")
			continue
		}
		var payload struct {
			Path      string
			Reserved0 string
			Reserved1 uint32
		}
		if err := ssh.Unmarshal(newCh.ExtraData(), &payload); err != nil {
			_ = newCh.Reject(ssh.ConnectionFailed, err.Error())
			continue
		}
		s.mu.Lock()
		s.paths = append(s.paths, payload.Path)
		s.mu.Unlock()
		target, err := net.Dial("unix", payload.Path)
		if err != nil {
			_ = newCh.Reject(ssh.ConnectionFailed, err.Error())
			continue
		}
		ch, chReqs, err := newCh.Accept()
		if err != nil {
			target.Close()
			continue
		}
		go ssh.DiscardRequests(chReqs)
		go func() { _, _ = io.Copy(ch, target); ch.Close() }()
		go func() { _, _ = io.Copy(target, ch); target.Close() }()
	}
}

func (s *testSSHServer) requestedPaths() []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]string(nil), s.paths...)
}

// startEchoSocket serves a unix socket that echoes everything back, standing
// in for libvirt-sock.
func startEchoSocket(t *testing.T) string {
	t.Helper()
	path := filepath.Join(t.TempDir(), "libvirt-sock")
	ln, err := net.Listen("unix", path)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { ln.Close() })
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			go func() { _, _ = io.Copy(c, c); c.Close() }()
		}
	}()
	return path
}

type sshFixture struct {
	server  *testSSHServer
	cfg     provider.KVMConfig
	host    provider.KVMHost
	keyFile string
}

func newSSHFixture(t *testing.T, answerGlobalRequests bool) *sshFixture {
	t.Helper()
	clientSigner, clientPEM := newKeyPair(t)
	server := startSSHServer(t, clientSigner.PublicKey(), answerGlobalRequests)
	dir := t.TempDir()
	keyFile := filepath.Join(dir, "id_ed25519")
	if err := os.WriteFile(keyFile, clientPEM, 0o600); err != nil {
		t.Fatal(err)
	}
	khFile := filepath.Join(dir, "known_hosts")
	line := knownhosts.Line([]string{server.addr}, server.hostKey)
	if err := os.WriteFile(khFile, []byte(line+"\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	return &sshFixture{
		server:  server,
		keyFile: keyFile,
		cfg:     provider.KVMConfig{SSHKeyFile: keyFile, KnownHostsFile: khFile},
		host:    provider.KVMHost{Name: "nas1", Address: server.addr, User: "virt", Socket: startEchoSocket(t)},
	}
}

func TestSSHAuth_dialSSH_forwardsUnixSocket(t *testing.T) {
	t.Parallel()
	f := newSSHFixture(t, true)
	auth, err := newSSHAuth(f.cfg)
	if err != nil {
		t.Fatalf("newSSHAuth() error = %v", err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	client, err := auth.dialSSH(ctx, f.host)
	if err != nil {
		t.Fatalf("dialSSH() error = %v", err)
	}
	defer client.Close()

	conn, err := unixForward{client: client, path: f.host.Socket}.Dial()
	if err != nil {
		t.Fatalf("unixForward.Dial() error = %v", err)
	}
	defer conn.Close()
	if _, err := conn.Write([]byte("ping")); err != nil {
		t.Fatal(err)
	}
	buf := make([]byte, 4)
	if _, err := io.ReadFull(conn, buf); err != nil || string(buf) != "ping" {
		t.Fatalf("echo through the forward = %q, %v; want ping", buf, err)
	}
	if paths := f.server.requestedPaths(); len(paths) != 1 || paths[0] != f.host.Socket {
		t.Errorf("server saw forward requests %v, want [%s]", paths, f.host.Socket)
	}
}

func TestSSHAuth_dialSSH_rejectsUnknownHostKey(t *testing.T) {
	t.Parallel()
	f := newSSHFixture(t, true)
	impostor, _ := newKeyPair(t)
	line := knownhosts.Line([]string{f.server.addr}, impostor.PublicKey())
	if err := os.WriteFile(f.cfg.KnownHostsFile, []byte(line+"\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	auth, err := newSSHAuth(f.cfg)
	if err != nil {
		t.Fatalf("newSSHAuth() error = %v", err)
	}
	client, err := auth.dialSSH(context.Background(), f.host)
	if err == nil {
		client.Close()
		t.Fatal("dialSSH() succeeded against a host whose key is not in known_hosts")
	}
	var keyErr *knownhosts.KeyError
	if !errors.As(err, &keyErr) {
		t.Errorf("dialSSH() error = %v, want a knownhosts.KeyError", err)
	}
}

func TestSSHAuth_dialSSH_insecureIgnoreHostKey(t *testing.T) {
	t.Parallel()
	f := newSSHFixture(t, true)
	f.cfg.KnownHostsFile = ""
	f.cfg.InsecureIgnoreHostKey = true
	auth, err := newSSHAuth(f.cfg)
	if err != nil {
		t.Fatalf("newSSHAuth() error = %v", err)
	}
	client, err := auth.dialSSH(context.Background(), f.host)
	if err != nil {
		t.Fatalf("dialSSH() error = %v", err)
	}
	client.Close()
}

func TestSSHAuth_dialSSH_rejectsWrongClientKey(t *testing.T) {
	t.Parallel()
	f := newSSHFixture(t, true)
	_, otherPEM := newKeyPair(t)
	if err := os.WriteFile(f.keyFile, otherPEM, 0o600); err != nil {
		t.Fatal(err)
	}
	auth, err := newSSHAuth(f.cfg)
	if err != nil {
		t.Fatalf("newSSHAuth() error = %v", err)
	}
	client, err := auth.dialSSH(context.Background(), f.host)
	if err == nil {
		client.Close()
		t.Fatal("dialSSH() succeeded with a key the server does not know")
	}
	if !strings.Contains(err.Error(), "ssh handshake") {
		t.Errorf("dialSSH() error = %v, want an auth failure from the handshake", err)
	}
}

func TestSSHAuth_dialSSH_cancelledContext(t *testing.T) {
	t.Parallel()
	f := newSSHFixture(t, true)
	auth, err := newSSHAuth(f.cfg)
	if err != nil {
		t.Fatalf("newSSHAuth() error = %v", err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if client, err := auth.dialSSH(ctx, f.host); err == nil {
		client.Close()
		t.Fatal("dialSSH() succeeded with an already-cancelled context")
	}
}

func TestNewSSHAuth_missingFiles(t *testing.T) {
	t.Parallel()
	_, keyPEM := newKeyPair(t)
	keyFile := filepath.Join(t.TempDir(), "key")
	if err := os.WriteFile(keyFile, keyPEM, 0o600); err != nil {
		t.Fatal(err)
	}
	tests := []struct {
		name    string
		cfg     provider.KVMConfig
		wantSub string
	}{
		{"key file missing", provider.KVMConfig{SSHKeyFile: "/nonexistent/key", InsecureIgnoreHostKey: true}, "sshKeyFile"},
		{"key file garbage", provider.KVMConfig{SSHKeyFile: keyFile + ".bad", InsecureIgnoreHostKey: true}, "sshKeyFile"},
		{"known_hosts missing", provider.KVMConfig{SSHKeyFile: keyFile, KnownHostsFile: "/nonexistent/kh"}, "knownHostsFile"},
	}
	if err := os.WriteFile(keyFile+".bad", []byte("not a key"), 0o600); err != nil {
		t.Fatal(err)
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			_, err := newSSHAuth(tc.cfg)
			if err == nil || !strings.Contains(err.Error(), tc.wantSub) {
				t.Errorf("newSSHAuth() error = %v, want one mentioning %s", err, tc.wantSub)
			}
		})
	}
}

func TestKeepalive_closesClientWhenServerStalls(t *testing.T) {
	t.Parallel()
	f := newSSHFixture(t, false)
	auth, err := newSSHAuth(f.cfg)
	if err != nil {
		t.Fatalf("newSSHAuth() error = %v", err)
	}
	client, err := auth.dialSSH(context.Background(), f.host)
	if err != nil {
		t.Fatalf("dialSSH() error = %v", err)
	}
	defer client.Close()
	stop := make(chan struct{})
	defer close(stop)
	go keepalive(client, 20*time.Millisecond, stop)

	closed := make(chan struct{})
	go func() { _ = client.Wait(); close(closed) }()
	select {
	case <-closed:
	case <-time.After(2 * time.Second):
		t.Fatal("client still open after keepalives went unanswered")
	}
}

func TestKeepalive_keepsHealthyClientOpen(t *testing.T) {
	t.Parallel()
	f := newSSHFixture(t, true)
	auth, err := newSSHAuth(f.cfg)
	if err != nil {
		t.Fatalf("newSSHAuth() error = %v", err)
	}
	client, err := auth.dialSSH(context.Background(), f.host)
	if err != nil {
		t.Fatalf("dialSSH() error = %v", err)
	}
	defer client.Close()
	stop := make(chan struct{})
	defer close(stop)
	go keepalive(client, 10*time.Millisecond, stop)
	time.Sleep(60 * time.Millisecond)
	fwd := unixForward{client: client, path: f.host.Socket}
	conn, err := fwd.Dial()
	if err != nil {
		t.Fatalf("forward after several keepalives failed: %v", err)
	}
	conn.Close()
}