//go:build e2e

package web_console

import (
	"bytes"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/golang-jwt/jwt/v4"
	"github.com/gorilla/websocket"
	"github.com/pingcap-inc/labs.tidb.io/ssh-proxy/internal/config"
	"golang.org/x/crypto/ssh"
)

func TestE2E(t *testing.T) {
	jwtPub, jwtPriv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("ed25519.GenerateKey: %v", err)
	}
	_, jwtPrivBad, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("ed25519.GenerateKey(bad): %v", err)
	}

	tmpDir := t.TempDir()
	sshKeyPath, sshAllowedPub := writeTestSSHIdentityFile(t, tmpDir)

	t.Run("HappyPath", func(t *testing.T) {
		const (
			region = "test"
			user   = "testuser"
		)

		sshSrv := startTestSSHServer(t, user, sshAllowedPub)
		defer sshSrv.Close()

		httpAddr, closeHTTP := startTestHTTPServer(t, &config.Config{
			JwtPublicKey: jwtPub,
			SshIdentityFiles: map[string]string{
				region: sshKeyPath,
			},
		})
		defer closeHTTP()

		token := signJWT(t, jwtPriv, jwt.MapClaims{
			"region": region,
			"m_ip":   sshSrv.Host(),
			"m_port": sshSrv.Port(),
			"m_user": user,
			"exp":    time.Now().Add(2 * time.Minute).Unix(),
		})

		conn, resp, err := wsDial(httpAddr, token)
		if err != nil {
			if resp != nil {
				_ = resp.Body.Close()
			}
			t.Fatalf("wsDial: %v", err)
		}
		defer conn.Close()

		wsWriteText(t, conn, `1:{"high":24,"width":80}`)

		wantPty := ptySize{rows: 24, cols: 80}
		if got := sshSrv.WaitForPty(t, 5*time.Second); got != wantPty {
			t.Fatalf("pty size mismatch: got=%+v want=%+v", got, wantPty)
		}

		var output []byte
		output = wsReadBinaryUntil(t, conn, output, 5*time.Second, []byte("READY\n"), []byte("ERRREADY\n"))

		wsWriteBinary(t, conn, []byte("ping\n"))
		_ = wsReadBinaryUntil(t, conn, nil, 5*time.Second, []byte("pong\n"))

		wsWriteText(t, conn, `1:{"high":40,"width":100}`)
		wantResize := ptySize{rows: 40, cols: 100}
		if got := sshSrv.WaitForWindowChange(t, 5*time.Second); got != wantResize {
			t.Fatalf("window-change mismatch: got=%+v want=%+v", got, wantResize)
		}

		_ = conn.Close()
		waitForActiveSessions(t, 0, 5*time.Second)
	})

	t.Run("JWTInvalidSignature", func(t *testing.T) {
		httpAddr, closeHTTP := startTestHTTPServer(t, &config.Config{
			JwtPublicKey: jwtPub,
		})
		defer closeHTTP()

		token := signJWT(t, jwtPrivBad, jwt.MapClaims{
			"region": "test",
			"m_ip":   "127.0.0.1",
			"m_user": "root",
			"exp":    time.Now().Add(2 * time.Minute).Unix(),
		})

		_, resp, err := wsDial(httpAddr, token)
		if err == nil {
			t.Fatalf("expected handshake error")
		}
		if resp == nil {
			t.Fatalf("expected http response")
		}
		defer resp.Body.Close()
		if resp.StatusCode != http.StatusUnauthorized {
			t.Fatalf("unexpected status code: %d", resp.StatusCode)
		}
		waitForActiveSessions(t, 0, time.Second)
	})

	t.Run("JWTExpired", func(t *testing.T) {
		httpAddr, closeHTTP := startTestHTTPServer(t, &config.Config{
			JwtPublicKey: jwtPub,
		})
		defer closeHTTP()

		token := signJWT(t, jwtPriv, jwt.MapClaims{
			"region": "test",
			"m_ip":   "127.0.0.1",
			"m_user": "root",
			"exp":    time.Now().Add(-time.Minute).Unix(),
		})

		_, resp, err := wsDial(httpAddr, token)
		if err == nil {
			t.Fatalf("expected handshake error")
		}
		if resp == nil {
			t.Fatalf("expected http response")
		}
		defer resp.Body.Close()
		if resp.StatusCode != http.StatusUnauthorized {
			t.Fatalf("unexpected status code: %d", resp.StatusCode)
		}
		waitForActiveSessions(t, 0, time.Second)
	})

	t.Run("RegionMissingIdentityFile_NoUpgrade", func(t *testing.T) {
		httpAddr, closeHTTP := startTestHTTPServer(t, &config.Config{
			JwtPublicKey:     jwtPub,
			SshIdentityFiles: map[string]string{},
		})
		defer closeHTTP()

		token := signJWT(t, jwtPriv, jwt.MapClaims{
			"region": "missing",
			"m_ip":   "127.0.0.1",
			"m_user": "root",
			"exp":    time.Now().Add(2 * time.Minute).Unix(),
		})

		_, resp, err := wsDial(httpAddr, token)
		if err == nil {
			t.Fatalf("expected handshake error")
		}
		if resp != nil && resp.StatusCode == http.StatusSwitchingProtocols {
			_ = resp.Body.Close()
			t.Fatalf("unexpected websocket upgrade")
		}
		if resp != nil {
			_ = resp.Body.Close()
		}
		waitForActiveSessions(t, 0, time.Second)
	})

	t.Run("SSHDialFailure_SendsTextError", func(t *testing.T) {
		const (
			region = "test"
			user   = "testuser"
		)

		dummyPort, closeDummy := startOneShotTCPCloser(t)
		defer closeDummy()

		httpAddr, closeHTTP := startTestHTTPServer(t, &config.Config{
			JwtPublicKey: jwtPub,
			SshIdentityFiles: map[string]string{
				region: sshKeyPath,
			},
		})
		defer closeHTTP()

		token := signJWT(t, jwtPriv, jwt.MapClaims{
			"region": region,
			"m_ip":   "127.0.0.1",
			"m_port": dummyPort,
			"m_user": user,
			"exp":    time.Now().Add(2 * time.Minute).Unix(),
		})

		conn, resp, err := wsDial(httpAddr, token)
		if err != nil {
			if resp != nil {
				_ = resp.Body.Close()
			}
			t.Fatalf("wsDial: %v", err)
		}
		defer conn.Close()

		wsWriteText(t, conn, `1:{"high":24,"width":80}`)

		_ = conn.SetReadDeadline(time.Now().Add(5 * time.Second))
		msgType, payload, err := conn.ReadMessage()
		if err != nil {
			t.Fatalf("conn.ReadMessage: %v", err)
		}
		if msgType != websocket.TextMessage {
			t.Fatalf("expected TextMessage, got %d payload=%q", msgType, string(payload))
		}
		if !strings.Contains(string(payload), "failed to dial:") {
			t.Fatalf("unexpected error message: %q", string(payload))
		}

		_ = conn.Close()
		waitForActiveSessions(t, 0, 5*time.Second)
	})

	t.Run("HealthEndpoint", func(t *testing.T) {
		httpAddr, closeHTTP := startTestHTTPServer(t, &config.Config{})
		defer closeHTTP()

		client := &http.Client{Timeout: 3 * time.Second}
		resp, err := client.Get("http://" + httpAddr + "/health")
		if err != nil {
			t.Fatalf("GET /health: %v", err)
		}
		defer resp.Body.Close()
		if resp.StatusCode != http.StatusOK {
			t.Fatalf("unexpected status code: %d", resp.StatusCode)
		}
		var body struct {
			Message string `json:"message"`
		}
		if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
			t.Fatalf("decode response: %v", err)
		}
		if body.Message != "I'm alive" {
			t.Fatalf("unexpected message: %q", body.Message)
		}
	})
}

func startTestHTTPServer(t *testing.T, cfg *config.Config) (addr string, cleanup func()) {
	t.Helper()

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("net.Listen(http): %v", err)
	}

	srv := newHTTPServer(cfg)
	done := make(chan struct{})
	go func() {
		_ = srv.Serve(ln)
		close(done)
	}()

	cleanup = func() {
		_ = srv.Close()
		_ = ln.Close()
		<-done
	}
	return ln.Addr().String(), cleanup
}

func wsDial(httpAddr, jwtToken string) (*websocket.Conn, *http.Response, error) {
	wsURL := url.URL{
		Scheme:   "ws",
		Host:     httpAddr,
		Path:     "/ssh-server/web",
		RawQuery: "auth=" + url.QueryEscape(jwtToken),
	}
	dialer := websocket.Dialer{
		HandshakeTimeout: 5 * time.Second,
	}
	return dialer.Dial(wsURL.String(), nil)
}

func wsWriteText(t *testing.T, conn *websocket.Conn, msg string) {
	t.Helper()
	conn.SetWriteDeadline(time.Now().Add(2 * time.Second))
	if err := conn.WriteMessage(websocket.TextMessage, []byte(msg)); err != nil {
		t.Fatalf("conn.WriteMessage(text): %v", err)
	}
}

func wsWriteBinary(t *testing.T, conn *websocket.Conn, msg []byte) {
	t.Helper()
	conn.SetWriteDeadline(time.Now().Add(2 * time.Second))
	if err := conn.WriteMessage(websocket.BinaryMessage, msg); err != nil {
		t.Fatalf("conn.WriteMessage(binary): %v", err)
	}
}

func wsReadBinaryUntil(t *testing.T, conn *websocket.Conn, buf []byte, timeout time.Duration, wants ...[]byte) []byte {
	t.Helper()

	deadline := time.Now().Add(timeout)
	for {
		allFound := true
		for _, want := range wants {
			if !bytes.Contains(buf, want) {
				allFound = false
				break
			}
		}
		if allFound {
			return buf
		}
		if time.Now().After(deadline) {
			t.Fatalf("timeout waiting for %q (buffer=%q)", wants, string(buf))
		}

		_ = conn.SetReadDeadline(deadline)
		msgType, payload, err := conn.ReadMessage()
		if err != nil {
			t.Fatalf("conn.ReadMessage: %v", err)
		}
		if msgType != websocket.BinaryMessage {
			continue
		}
		buf = appendWindow(buf, payload, 1<<20)
	}
}

func appendWindow(dst, add []byte, max int) []byte {
	if max <= 0 {
		return dst
	}
	if len(add) >= max {
		dst = dst[:0]
		return append(dst, add[len(add)-max:]...)
	}
	if len(dst)+len(add) > max {
		drop := len(dst) + len(add) - max
		copy(dst, dst[drop:])
		dst = dst[:len(dst)-drop]
	}
	return append(dst, add...)
}

func waitForActiveSessions(t *testing.T, want int64, timeout time.Duration) {
	t.Helper()
	deadline := time.Now().Add(timeout)
	for {
		if GetActiveWebConsoleSessions() == want {
			return
		}
		if time.Now().After(deadline) {
			t.Fatalf("active session count mismatch: got=%d want=%d", GetActiveWebConsoleSessions(), want)
		}
		time.Sleep(10 * time.Millisecond)
	}
}

func signJWT(t *testing.T, priv ed25519.PrivateKey, claims jwt.Claims) string {
	t.Helper()
	token := jwt.NewWithClaims(jwt.SigningMethodEdDSA, claims)
	signed, err := token.SignedString(priv)
	if err != nil {
		t.Fatalf("token.SignedString: %v", err)
	}
	return signed
}

func writeTestSSHIdentityFile(t *testing.T, dir string) (path string, allowedPub ssh.PublicKey) {
	t.Helper()

	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("rsa.GenerateKey: %v", err)
	}
	signer, err := ssh.NewSignerFromKey(key)
	if err != nil {
		t.Fatalf("ssh.NewSignerFromKey: %v", err)
	}

	block := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(key),
	}
	pemBytes := pem.EncodeToMemory(block)
	if pemBytes == nil {
		t.Fatalf("pem.EncodeToMemory returned nil")
	}

	path = dir + "/id_rsa"
	if err := os.WriteFile(path, pemBytes, 0o600); err != nil {
		t.Fatalf("os.WriteFile: %v", err)
	}
	return path, signer.PublicKey()
}

type ptySize struct {
	rows uint32
	cols uint32
}

type testSSHServer struct {
	t            *testing.T
	ln           net.Listener
	expectedUser string
	allowedPub   ssh.PublicKey

	ptyCh    chan ptySize
	windowCh chan ptySize

	done       chan struct{}
	acceptDone chan struct{}
	connWg     sync.WaitGroup
}

func startTestSSHServer(t *testing.T, expectedUser string, allowedPub ssh.PublicKey) *testSSHServer {
	t.Helper()

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("net.Listen(ssh): %v", err)
	}

	hostKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		_ = ln.Close()
		t.Fatalf("rsa.GenerateKey(host): %v", err)
	}
	hostSigner, err := ssh.NewSignerFromKey(hostKey)
	if err != nil {
		_ = ln.Close()
		t.Fatalf("ssh.NewSignerFromKey(host): %v", err)
	}

	srv := &testSSHServer{
		t:            t,
		ln:           ln,
		expectedUser: expectedUser,
		allowedPub:   allowedPub,
		ptyCh:        make(chan ptySize, 1),
		windowCh:     make(chan ptySize, 1),
		done:         make(chan struct{}),
		acceptDone:   make(chan struct{}),
	}

	cfg := &ssh.ServerConfig{
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			if conn.User() != srv.expectedUser {
				return nil, fmt.Errorf("unexpected user: %q", conn.User())
			}
			if !bytes.Equal(key.Marshal(), srv.allowedPub.Marshal()) {
				return nil, fmt.Errorf("unexpected public key")
			}
			return nil, nil
		},
	}
	cfg.AddHostKey(hostSigner)

	go func() {
		defer close(srv.acceptDone)
		for {
			nc, err := ln.Accept()
			if err != nil {
				if errors.Is(err, net.ErrClosed) {
					return
				}
				select {
				case <-srv.done:
					return
				default:
					return
				}
			}
			srv.connWg.Add(1)
			go func(c net.Conn) {
				defer srv.connWg.Done()
				srv.handleConn(cfg, c)
			}(nc)
		}
	}()

	return srv
}

func (s *testSSHServer) Host() string {
	if tcp, ok := s.ln.Addr().(*net.TCPAddr); ok {
		return tcp.IP.String()
	}
	host, _, err := net.SplitHostPort(s.ln.Addr().String())
	if err != nil {
		return "127.0.0.1"
	}
	return host
}

func (s *testSSHServer) Port() int {
	if tcp, ok := s.ln.Addr().(*net.TCPAddr); ok {
		return tcp.Port
	}
	_, port, err := net.SplitHostPort(s.ln.Addr().String())
	if err != nil {
		s.t.Fatalf("net.SplitHostPort: %v", err)
	}
	p, err := net.LookupPort("tcp", port)
	if err != nil {
		s.t.Fatalf("net.LookupPort: %v", err)
	}
	return p
}

func (s *testSSHServer) Close() {
	close(s.done)
	_ = s.ln.Close()
	<-s.acceptDone
	s.connWg.Wait()
}

func (s *testSSHServer) WaitForPty(t *testing.T, timeout time.Duration) ptySize {
	t.Helper()
	select {
	case size := <-s.ptyCh:
		return size
	case <-time.After(timeout):
		t.Fatalf("timed out waiting for pty-req")
		return ptySize{}
	}
}

func (s *testSSHServer) WaitForWindowChange(t *testing.T, timeout time.Duration) ptySize {
	t.Helper()
	select {
	case size := <-s.windowCh:
		return size
	case <-time.After(timeout):
		t.Fatalf("timed out waiting for window-change")
		return ptySize{}
	}
}

func (s *testSSHServer) handleConn(cfg *ssh.ServerConfig, c net.Conn) {
	defer c.Close()

	conn, chans, reqs, err := ssh.NewServerConn(c, cfg)
	if err != nil {
		return
	}
	defer conn.Close()

	go ssh.DiscardRequests(reqs)

	for newChannel := range chans {
		if newChannel.ChannelType() != "session" {
			newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}

		channel, requests, err := newChannel.Accept()
		if err != nil {
			continue
		}
		s.handleSession(channel, requests)
		_ = channel.Close()
	}
}

type ptyReqMsg struct {
	Term     string
	Columns  uint32
	Rows     uint32
	Width    uint32
	Height   uint32
	Modelist string
}

type windowChangeMsg struct {
	Columns uint32
	Rows    uint32
	Width   uint32
	Height  uint32
}

func (s *testSSHServer) handleSession(channel ssh.Channel, requests <-chan *ssh.Request) {
	var shellOnce sync.Once

	for req := range requests {
		switch req.Type {
		case "pty-req":
			var msg ptyReqMsg
			if err := ssh.Unmarshal(req.Payload, &msg); err != nil {
				req.Reply(false, nil)
				continue
			}
			select {
			case s.ptyCh <- ptySize{rows: msg.Rows, cols: msg.Columns}:
			default:
			}
			req.Reply(true, nil)
		case "shell":
			req.Reply(true, nil)
			shellOnce.Do(func() {
				go runPseudoShell(channel)
			})
		case "window-change":
			var msg windowChangeMsg
			if err := ssh.Unmarshal(req.Payload, &msg); err != nil {
				if req.WantReply {
					req.Reply(false, nil)
				}
				continue
			}
			select {
			case s.windowCh <- ptySize{rows: msg.Rows, cols: msg.Columns}:
			default:
			}
			if req.WantReply {
				req.Reply(true, nil)
			}
		default:
			req.Reply(false, nil)
		}
	}
}

func runPseudoShell(ch ssh.Channel) {
	_, _ = io.WriteString(ch, "READY\n")
	_, _ = io.WriteString(ch.Stderr(), "ERRREADY\n")

	var buf [1024]byte
	var pending []byte
	for {
		n, err := ch.Read(buf[:])
		if n > 0 {
			pending = append(pending, buf[:n]...)
			for {
				pingIdx := bytes.Index(pending, []byte("ping\n"))
				exitIdx := bytes.Index(pending, []byte("exit\n"))
				if pingIdx == -1 && exitIdx == -1 {
					break
				}
				switch {
				case exitIdx != -1 && (pingIdx == -1 || exitIdx < pingIdx):
					_ = ch.Close()
					return
				case pingIdx != -1:
					_, _ = io.WriteString(ch, "pong\n")
					pending = pending[pingIdx+len("ping\n"):]
				}
			}
			if len(pending) > 4096 {
				pending = append([]byte(nil), pending[len(pending)-4096:]...)
			}
		}
		if err != nil {
			return
		}
	}
}

func startOneShotTCPCloser(t *testing.T) (port int, cleanup func()) {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("net.Listen(dummy): %v", err)
	}
	tcp := ln.Addr().(*net.TCPAddr)
	done := make(chan struct{})
	go func() {
		defer close(done)
		c, err := ln.Accept()
		if err != nil {
			return
		}
		_ = c.Close()
		_ = ln.Close()
	}()
	return tcp.Port, func() {
		_ = ln.Close()
		<-done
	}
}