package web

import (
	"crypto/ed25519"
	"crypto/rand"
	"crypto/rsa"
	"encoding/json"
	"encoding/pem"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/labstack/echo/v5"
	"github.com/shellhub-io/shellhub/pkg/api/authorizer"
	"github.com/shellhub-io/shellhub/pkg/api/identity"
	"github.com/shellhub-io/shellhub/pkg/api/identity/identitytest"
	"github.com/shellhub-io/shellhub/pkg/models"
	"github.com/shellhub-io/shellhub/server/api/pkg/echo/handlers"
	"github.com/shellhub-io/shellhub/server/api/pkg/gateway"
	"github.com/shellhub-io/shellhub/server/api/pkg/origin"
	servicemocks "github.com/shellhub-io/shellhub/server/api/services/mocks"
	"github.com/shellhub-io/shellhub/server/api/store"
	"github.com/shellhub-io/shellhub/server/ssh/pkg/webhandoff"
	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	"golang.org/x/crypto/ssh"
	"golang.org/x/net/websocket"
)

func writeHostKey(t *testing.T) string {
	t.Helper()

	_, key, err := ed25519.GenerateKey(rand.Reader)
	require.NoError(t, err)

	block, err := ssh.MarshalPrivateKey(key, "")
	require.NoError(t, err)

	path := filepath.Join(t.TempDir(), "ssh.key")
	require.NoError(t, os.WriteFile(path, pem.EncodeToMemory(block), 0o600))

	return path
}

// TestExitLogLevel verifies that exitLogLevel returns logrus.WarnLevel for
// expected/banner-derived errors and logrus.ErrorLevel for genuine server faults.
func TestExitLogLevel(t *testing.T) {
	tests := []struct {
		description string
		err         error
		expected    logrus.Level
	}{
		{
			description: "ErrConnect uses Warn",
			err:         ErrConnect,
			expected:    logrus.WarnLevel,
		},
		{
			description: "ErrAccessDenied uses Warn",
			err:         ErrAccessDenied,
			expected:    logrus.WarnLevel,
		},
		{
			description: "ErrInvalidSSHID uses Warn",
			err:         ErrInvalidSSHID,
			expected:    logrus.WarnLevel,
		},
		{
			description: "ErrAuthentication uses Warn",
			err:         ErrAuthentication,
			expected:    logrus.WarnLevel,
		},
		{
			description: "ErrGetAuth uses Warn",
			err:         ErrGetAuth,
			expected:    logrus.WarnLevel,
		},
		{
			description: "ErrFindDevice uses Warn",
			err:         ErrFindDevice,
			expected:    logrus.WarnLevel,
		},
		{
			description: "ErrForbiddenPublicKey uses Warn",
			err:         ErrForbiddenPublicKey,
			expected:    logrus.WarnLevel,
		},
		{
			description: "ErrBridgeCredentialsNotFound uses Warn",
			err:         ErrBridgeCredentialsNotFound,
			expected:    logrus.WarnLevel,
		},
		{
			description: "ErrWebSocketGetToken uses Warn",
			err:         ErrWebSocketGetToken,
			expected:    logrus.WarnLevel,
		},
		{
			description: "ErrWebSocketGetDimensions uses Warn",
			err:         ErrWebSocketGetDimensions,
			expected:    logrus.WarnLevel,
		},
		{
			description: "ErrWebSocketGetIP uses Warn",
			err:         ErrWebSocketGetIP,
			expected:    logrus.WarnLevel,
		},
		{
			description: "ErrSession uses Error",
			err:         ErrSession,
			expected:    logrus.ErrorLevel,
		},
		{
			description: "ErrPty uses Error",
			err:         ErrPty,
			expected:    logrus.ErrorLevel,
		},
		{
			description: "ErrShell uses Error",
			err:         ErrShell,
			expected:    logrus.ErrorLevel,
		},
		{
			description: "unrecognised error uses Error",
			err:         ErrPublicKey,
			expected:    logrus.ErrorLevel,
		},
	}

	for _, test := range tests {
		t.Run(test.description, func(t *testing.T) {
			assert.Equal(t, test.expected, exitLogLevel(test.err))
		})
	}
}

func TestNewSSHServerBridge_CredentialsNotFound(t *testing.T) {
	e := echo.New()

	require.NoError(t, NewSSHServerBridge(e, nil, nil, webhandoff.NewStore(), &Config{HostKeyFile: writeHostKey(t)}))

	server := httptest.NewServer(e)
	defer server.Close()

	assert.NotPanics(t, func() {
		assert.Contains(t, openBridge(t, server.URL, "nonexistent"), ErrBridgeCredentialsNotFound.Error())
	}, "handler must not panic when credentials are not found")
}

const tenantID = "00000000-0000-4000-0000-000000000000"

func deviceInNamespace(inScope bool) (*models.Device, error) {
	if !inScope {
		return nil, store.ErrNoDocuments
	}

	return &models.Device{UID: "device-uid", TenantID: tenantID}, nil //nolint:exhaustruct // the handler only checks that the device resolves
}

func TestWebSessionRouteRequiresTheConnectPermission(t *testing.T) {
	tests := []struct {
		description   string
		role          string
		deviceInScope bool
		expectedCode  int
	}{
		{
			description:  "observer is refused",
			role:         authorizer.RoleObserver.String(),
			expectedCode: http.StatusForbidden,
		},
		{
			description:  "an identity carrying no role is refused",
			role:         "",
			expectedCode: http.StatusForbidden,
		},
		{
			description:   "operator is served a token for a device in their namespace",
			role:          authorizer.RoleOperator.String(),
			deviceInScope: true,
			expectedCode:  http.StatusOK,
		},
		{
			description:  "a device outside the caller's namespace is refused",
			role:         authorizer.RoleOperator.String(),
			expectedCode: http.StatusForbidden,
		},
	}

	for _, test := range tests {
		t.Run(test.description, func(t *testing.T) {
			service := servicemocks.NewMockService(t)
			if test.role == authorizer.RoleOperator.String() {
				service.On("GetDevice", mock.Anything, mock.Anything, models.UID("device-uid")).
					Return(deviceInNamespace(test.deviceInScope)).Once()
			}

			e := echo.New()
			e.HTTPErrorHandler = handlers.NewErrors(nil)
			e.Use(gateway.WithContext())

			require.NoError(t, NewSSHServerBridge(e, nil, service, webhandoff.NewStore(), &Config{HostKeyFile: writeHostKey(t)}))

			server := httptest.NewServer(identitytest.Serve(e, identity.Identity{
				TenantID: tenantID,
				Role:     authorizer.RoleFromString(test.role),
			}))
			defer server.Close()

			body := strings.NewReader(`{"device":"device-uid","username":"root","password":"secret"}`)

			req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, server.URL+WebSessionRoute, body)
			require.NoError(t, err)

			req.Header.Set("Content-Type", "application/json")

			res, err := http.DefaultClient.Do(req)
			require.NoError(t, err)

			defer res.Body.Close() //nolint:errcheck

			assert.Equal(t, test.expectedCode, res.StatusCode)

			if test.expectedCode != http.StatusOK {
				var refusal struct {
					Message string `json:"message"`
					Error   string `json:"error"`
				}

				require.NoError(t, json.NewDecoder(res.Body).Decode(&refusal))
				assert.True(t, refusal.Message != "" || refusal.Error != "",
					"a refusal must name what failed, through the error handler's message or this route's own error")

				return
			}

			var success struct {
				Token string `json:"token"`
			}

			require.NoError(t, json.NewDecoder(res.Body).Decode(&success))
			assert.NotEmpty(t, success.Token)
		})
	}
}

func TestNewSSHServerBridge_TokenOpensOneBridge(t *testing.T) {
	service := servicemocks.NewMockService(t)
	service.On("GetDevice", mock.Anything, mock.Anything, models.UID("device-uid")).
		Return(deviceInNamespace(true)).Once()
	service.On("GetDevice", mock.Anything, mock.Anything, models.UID("device-uid")).
		Return(nil, store.ErrNoDocuments).Once()

	server := startBridge(t, service)

	status, reply := postSession(t, server.URL, `{"device":"device-uid","username":"root","fingerprint":"fingerprint"}`)
	require.Equal(t, http.StatusOK, status)

	assert.Contains(t, openBridge(t, server.URL, reply.Token), ErrGetAuth.Error())
	assert.Contains(t, openBridge(t, server.URL, reply.Token), ErrBridgeCredentialsNotFound.Error())
}

func TestNewSSHServerBridge_FailedOpenSpendsTheToken(t *testing.T) {
	tests := []struct {
		description string
		query       string
		resolved    bool
		expected    error
	}{
		{
			description: "an invalid terminal size",
			query:       "cols=wide&rows=24",
			resolved:    true,
			expected:    ErrWebSocketGetDimensions,
		},
		{
			description: "a client origin the router did not resolve",
			query:       "cols=80&rows=24",
			resolved:    false,
			expected:    ErrWebSocketGetIP,
		},
	}

	for _, test := range tests {
		t.Run(test.description, func(t *testing.T) {
			service := servicemocks.NewMockService(t)
			service.On("GetDevice", mock.Anything, mock.Anything, models.UID("device-uid")).
				Return(deviceInNamespace(true)).Once()

			server := startBridgeResolving(t, service, test.resolved)

			status, reply := postSession(t, server.URL, `{"device":"device-uid","username":"root","fingerprint":"fingerprint"}`)
			require.Equal(t, http.StatusOK, status)

			assert.Contains(t, openBridgeWith(t, server.URL, "token="+reply.Token+"&"+test.query), test.expected.Error())
			assert.Contains(t, openBridge(t, server.URL, reply.Token), ErrBridgeCredentialsNotFound.Error())
		})
	}
}

func TestWebSessionRouteRefusesAPasswordItCannotEncrypt(t *testing.T) {
	service := servicemocks.NewMockService(t)
	service.On("GetDevice", mock.Anything, mock.Anything, models.UID("device-uid")).
		Return(deviceInNamespace(true)).Once()

	server := startBridge(t, service)

	status, reply := postSession(t, server.URL, `{"device":"device-uid","username":"root","password":"`+strings.Repeat("a", 191)+`"}`)
	assert.Equal(t, http.StatusBadRequest, status)
	assert.Contains(t, reply.Error, ErrCredentialsEncryptPassword.Error())
	assert.Contains(t, reply.Error, rsa.ErrMessageTooLong.Error())
}

func startBridge(t *testing.T, service *servicemocks.MockService) *httptest.Server {
	t.Helper()

	return startBridgeResolving(t, service, true)
}

func startBridgeResolving(t *testing.T, service *servicemocks.MockService, resolved bool) *httptest.Server {
	t.Helper()

	e := echo.New()
	e.HTTPErrorHandler = handlers.NewErrors(nil)
	e.Use(gateway.WithContext())
	if resolved {
		e.Pre(func(next echo.HandlerFunc) echo.HandlerFunc {
			return func(c *echo.Context) error {
				req := c.Request()
				c.SetRequest(req.WithContext(origin.NewContext(req.Context(), origin.Loopback().Resolve(req))))

				return next(c)
			}
		})
	}

	require.NoError(t, NewSSHServerBridge(e, nil, service, webhandoff.NewStore(), &Config{HostKeyFile: writeHostKey(t)}))

	server := httptest.NewServer(identitytest.Serve(e, identity.Identity{TenantID: tenantID, Role: authorizer.RoleOperator}))
	t.Cleanup(server.Close)

	return server
}

type sessionReply struct {
	Token string `json:"token"`
	Error string `json:"error"`
}

func postSession(t *testing.T, serverURL, body string) (int, sessionReply) {
	t.Helper()

	req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, serverURL+WebSessionRoute, strings.NewReader(body))
	require.NoError(t, err)

	req.Header.Set("Content-Type", "application/json")

	res, err := http.DefaultClient.Do(req)
	require.NoError(t, err)

	defer res.Body.Close() //nolint:errcheck // the body is decoded before the close runs, and a failed close cannot change it

	var reply sessionReply
	require.NoError(t, json.NewDecoder(res.Body).Decode(&reply))

	return res.StatusCode, reply
}

func openBridge(t *testing.T, serverURL, token string) string {
	t.Helper()

	return openBridgeWith(t, serverURL, "token="+token+"&cols=80&rows=24")
}

func openBridgeWith(t *testing.T, serverURL, query string) string {
	t.Helper()

	cfg, err := websocket.NewConfig("ws"+strings.TrimPrefix(serverURL, "http")+WebsocketSSHBridgeRoute+"?"+query, serverURL)
	require.NoError(t, err)

	conn, err := websocket.DialConfig(cfg)
	require.NoError(t, err)

	defer conn.Close() //nolint:errcheck // the bridge has already answered, and a failed close does not change it

	var raw []byte
	require.NoError(t, websocket.Message.Receive(conn, &raw))

	var msg Message
	require.NoError(t, json.Unmarshal(raw, &msg))
	require.Equal(t, messageKindError, msg.Kind)

	data, ok := msg.Data.(string)
	require.True(t, ok)

	return data
}

func TestNewSSHServerBridge_MissingHostKey(t *testing.T) {
	err := NewSSHServerBridge(echo.New(), nil, nil, webhandoff.NewStore(), &Config{HostKeyFile: filepath.Join(t.TempDir(), "absent.key")})

	assert.ErrorIs(t, err, ErrBridgeReadHostKey)
}