package main

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"net/http"
	"sync"
	"testing"
	"time"

	"github.com/shellhub-io/shellhub/pkg/api/authorizer"
	"github.com/shellhub-io/shellhub/pkg/api/jwttoken"
	"github.com/shellhub-io/shellhub/pkg/api/requests"
	"github.com/shellhub-io/shellhub/pkg/devicekey"
	"github.com/shellhub-io/shellhub/pkg/models"
	"github.com/shellhub-io/shellhub/tests/environment"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/testcontainers/testcontainers-go"
)

const (
	otherNamespaceName     = "otherspace"
	otherNamespace         = "00000000-0000-4000-0000-000000000001"
	fastTunnelPingInterval = 2 * time.Second
	tunnelBeatTimeout      = 15 * time.Second
	staleBeatCycles        = 3
	staleBeatWindow        = 2500 * time.Millisecond
)

// TestDeviceAuth drives the endpoint an agent enrolls through with requests built by hand, so each
// case controls exactly which part of the device's identity changes. The provisioning-key paths
// are covered by [TestProvisioningKeyEnrollment].
func TestDeviceAuth(t *testing.T) {
	t.Parallel()

	compose := newSSHEnvironment(t, t.Context(), legacyNamespace)
	compose.NewLegacyNamespace(t, ShellHubUsername, otherNamespaceName, otherNamespace)

	t.Run("a tenant id earns a token signed for the device", func(t *testing.T) {
		res := authDevice(t, compose, newDeviceAuthRequest(t, "token", "02:00:00:00:01:00"))

		assert.NotEmpty(t, res.UID)
		assert.Equal(t, ShellHubNamespace, res.TenantID)
		assert.Equal(t, ShellHubNamespaceName, res.Namespace)
		assert.Equal(t, models.DeviceStatusPending, res.Status)

		claims, err := jwttoken.ClaimsFromBearerToken(compose.APIPublicKey(t), res.Token)
		require.NoError(t, err)
		assert.Equal(t, &authorizer.DeviceClaims{UID: res.UID, TenantID: ShellHubNamespace, KeyVerified: true}, claims)
	})

	t.Run("authenticating alone leaves the device offline", func(t *testing.T) {
		req := newDeviceAuthRequest(t, "offline", "02:00:00:00:01:f0")

		res := authDevice(t, compose, req)
		compose.UpdateDeviceStatus(t, res.UID, environment.DeviceActionAccept)
		authDevice(t, compose, req)

		device := requireDevice(t, compose, res.UID)
		assert.False(t, device.Online, "a device that never opened a tunnel reads online")
		assert.True(t, device.LastSeen.IsZero(), "last_seen %s was written by an authentication", device.LastSeen)
	})

	t.Run("the uid is derived from the identity", func(t *testing.T) {
		otherPublicKey := newDevicePublicKey(t)

		cases := []struct {
			description string
			change      func(req *requests.DeviceAuth)
			compare     assert.ComparisonAssertionFunc
		}{
			{
				description: "authenticating again keeps the uid",
				change:      func(*requests.DeviceAuth) {},
				compare:     assert.Equal,
			},
			{
				description: "a new hostname keeps the uid",
				change: func(req *requests.DeviceAuth) {
					req.Hostname = "renamed"
				},
				compare: assert.Equal,
			},
			{
				description: "a new mac address changes the uid",
				change: func(req *requests.DeviceAuth) {
					req.Identity = &requests.DeviceIdentity{MAC: "02:00:00:00:03:ff"}
				},
				compare: assert.NotEqual,
			},
			{
				description: "a new public key changes the uid",
				change: func(req *requests.DeviceAuth) {
					req.PublicKey = otherPublicKey
				},
				compare: assert.NotEqual,
			},
			{
				description: "another namespace changes the uid",
				change: func(req *requests.DeviceAuth) {
					req.TenantID = otherNamespace
				},
				compare: assert.NotEqual,
			},
		}

		for i, tc := range cases {
			t.Run(tc.description, func(t *testing.T) {
				req := newDeviceAuthRequest(t, fmt.Sprintf("uid-%d", i), fmt.Sprintf("02:00:00:00:03:%02x", i))

				first := authDevice(t, compose, req)

				tc.change(&req)
				second := authDevice(t, compose, req)

				tc.compare(t, first.UID, second.UID)
			})
		}
	})
}

// TestDeviceAuthAgent runs a real agent, for what the server does with an agent beyond a single
// request: the hostname the agent reports, the tunnel it keeps open, and the enrollment it repeats
// after the device is removed.
func TestDeviceAuthAgent(t *testing.T) {
	t.Parallel()

	t.Run("an agent's hostname has its dots replaced by underscores", func(t *testing.T) {
		compose := newSSHEnvironment(t, t.Context(), legacyNamespace)

		startAgent(t, t.Context(), compose, NewAgentContainerWithHostname("e2e.agent.example"))

		device := compose.AwaitDeviceWithStatus(t, models.DeviceStatusPending)
		assert.Equal(t, "e2e_agent_example", device.Name)
	})

	t.Run("the tunnel ping keeps advancing last_seen", func(t *testing.T) {
		compose := newFastTunnelPingEnvironment(t, t.Context(), legacyNamespace)

		_, device := startAcceptedAgent(t, t.Context(), compose)

		seen := device.LastSeen
		for range 2 {
			seen = awaitLastSeenPast(t, compose, device.UID, seen)
		}
	})

	t.Run("an agent stopped soon after it connects stays offline", func(t *testing.T) {
		compose := newSSHEnvironment(t, t.Context(), legacyNamespace)

		agent, device := startAcceptedAgent(t, t.Context(), compose)
		require.NoError(t, agent.Stop(t.Context(), nil))
		compose.AwaitDeviceOffline(t, device.UID)

		for cycle := range staleBeatCycles {
			require.NoError(t, agent.Start(t.Context()))
			awaitAgentConnections(t, agent, cycle+2)

			require.NoError(t, agent.Stop(t.Context(), nil))
			compose.AwaitDeviceOffline(t, device.UID)

			require.Never(t, func() bool {
				current, _, err := compose.GetDevice(t.Context(), device.UID)

				return err == nil && current.Online
			}, staleBeatWindow, 100*time.Millisecond, "cycle %d: the device came back online after its agent stopped", cycle)
		}
	})

	t.Run("a removed device re-registers as pending under the same uid", func(t *testing.T) {
		compose := newSSHEnvironment(t, t.Context(), legacyNamespace)

		agent, device := startAcceptedAgent(t, t.Context(), compose)
		require.NoError(t, agent.Stop(t.Context(), nil))

		compose.DeleteDevice(t, device.UID)

		removed := compose.AwaitDeviceWithStatus(t, models.DeviceStatusRemoved)
		require.Equal(t, device.UID, removed.UID)

		require.NoError(t, agent.Start(t.Context()))

		pending := compose.AwaitDeviceWithStatus(t, models.DeviceStatusPending)
		assert.Equal(t, device.UID, pending.UID)
	})

	t.Run("a removed device's running agent re-registers as pending without a restart", func(t *testing.T) {
		compose := newSSHEnvironment(t, t.Context(), legacyNamespace)

		_, device := startAcceptedAgent(t, t.Context(), compose)

		compose.DeleteDevice(t, device.UID)

		pending := compose.AwaitDeviceWithStatus(t, models.DeviceStatusPending)
		assert.Equal(t, device.UID, pending.UID)
	})
}

func newFastTunnelPingEnvironment(t *testing.T, ctx context.Context, model namespaceModel) *environment.DockerCompose {
	t.Helper()

	configurator := environment.New(t, run).WithEnv("SHELLHUB_TUNNEL_PING_INTERVAL", fastTunnelPingInterval.String())

	return newConfiguredSSHEnvironment(t, ctx, configurator, model)
}

func awaitLastSeenPast(t *testing.T, compose *environment.DockerCompose, uid string, seen time.Time) time.Time {
	t.Helper()

	var current *models.Device

	require.EventuallyWithT(t, func(tt *assert.CollectT) {
		device, resp, err := compose.GetDevice(t.Context(), uid)
		if !assert.NoError(tt, err) {
			return
		}

		if !assert.Equal(tt, http.StatusOK, resp.StatusCode(), resp.String()) {
			return
		}

		assert.True(tt, device.LastSeen.After(seen), "last_seen %s has not moved past %s", device.LastSeen, seen)

		current = device
	}, 5*fastTunnelPingInterval, fastTunnelPingInterval/4)

	return current.LastSeen
}

func awaitAgentConnections(t *testing.T, agent testcontainers.Container, connections int) {
	t.Helper()

	require.EventuallyWithT(t, func(tt *assert.CollectT) {
		logs, err := agentLogs(t.Context(), agent)
		if !assert.NoError(tt, err) {
			return
		}

		assert.Len(tt, agentConnectedLine.FindAll(logs, -1), connections, "the agent has not connected to the server")
	}, 30*time.Second, 100*time.Millisecond)
}

func newDeviceAuthRequest(t *testing.T, hostname, mac string) requests.DeviceAuth {
	t.Helper()

	return requests.DeviceAuth{
		Info: &requests.DeviceInfo{
			ID:         "e2e",
			PrettyName: "e2e",
			Version:    "e2e",
			Arch:       "amd64",
			Platform:   "native",
		},
		Hostname:  hostname,
		Identity:  &requests.DeviceIdentity{MAC: mac},
		PublicKey: newDevicePublicKey(t),
		TenantID:  ShellHubNamespace,
	}
}

var devicePrivateKeys sync.Map

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

	key, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)

	publicKey := string(pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: x509.MarshalPKCS1PublicKey(&key.PublicKey),
	}))
	devicePrivateKeys.Store(publicKey, key)

	return publicKey
}

func signDeviceAuth(ctx context.Context, compose *environment.DockerCompose, req requests.DeviceAuth) (requests.DeviceAuth, error) {
	if req.Challenge != "" || req.Signature != "" {
		return req, nil
	}

	key, ok := devicePrivateKeys.Load(req.PublicKey)
	if !ok {
		return req, errors.New("the request's public key was not made by newDevicePublicKey, so nothing can sign it")
	}

	challenge, err := requestDeviceAuthChallenge(ctx, compose)
	if err != nil {
		return req, err
	}

	return signDeviceAuthWith(key.(*rsa.PrivateKey), req, challenge) //nolint:forcetypeassert // devicePrivateKeys holds nothing but device keys
}

func requestDeviceAuthChallenge(ctx context.Context, compose *environment.DockerCompose) (string, error) {
	challenge := new(models.DeviceAuthChallenge)

	resp, err := compose.Anonymous(ctx).SetResult(challenge).Post("/api/devices/auth/challenge")
	if err != nil {
		return "", err
	}

	if resp.StatusCode() != http.StatusOK {
		return "", fmt.Errorf("issuing a device auth challenge: %s", resp.String())
	}

	return challenge.Challenge, nil
}

func signDeviceAuthWith(key *rsa.PrivateKey, req requests.DeviceAuth, challenge string) (requests.DeviceAuth, error) {
	signature, err := devicekey.Sign(key, devicekey.Statement{
		Challenge:       challenge,
		TenantID:        req.TenantID,
		ProvisioningKey: req.ProvisioningKey,
		PublicKey:       req.PublicKey,
	})
	if err != nil {
		return req, err
	}

	req.Challenge, req.Signature = challenge, signature

	return req, nil
}

func authDevice(t *testing.T, compose *environment.DockerCompose, req requests.DeviceAuth) *models.DeviceAuthResponse {
	t.Helper()

	req, err := signDeviceAuth(t.Context(), compose, req)
	require.NoError(t, err)

	res := new(models.DeviceAuthResponse)

	resp, err := compose.Anonymous(t.Context()).SetBody(req).SetResult(res).Post("/api/devices/auth")
	require.NoError(t, err)
	require.Equal(t, http.StatusOK, resp.StatusCode(), resp.String())

	return res
}