package main

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/bramvdbogaerde/go-scp"
	"github.com/pkg/sftp"
	"github.com/shellhub-io/shellhub/pkg/api/requests"
	"github.com/shellhub-io/shellhub/pkg/models"
	"github.com/shellhub-io/shellhub/tests/environment"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/crypto/ssh"
)

func listSessions(ctx context.Context, compose *environment.DockerCompose) ([]models.Session, error) {
	sessions := []models.Session{}

	resp, err := compose.R(ctx).SetResult(&sessions).Get("/api/sessions")
	if err != nil {
		return nil, err
	}

	if resp.StatusCode() != 200 {
		return nil, fmt.Errorf("the list answered %d, and the strict validator names the mismatch: %s", resp.StatusCode(), resp.String())
	}

	return sessions, nil
}

func currentSessions(t *testing.T, ctx context.Context, compose *environment.DockerCompose) []models.Session {
	t.Helper()

	var sessions []models.Session

	require.EventuallyWithT(t, func(tt *assert.CollectT) {
		var err error

		sessions, err = listSessions(ctx, compose)
		assert.NoError(tt, err)
	}, 30*time.Second, 1*time.Second)

	return sessions
}

func sessionAfter(t *testing.T, ctx context.Context, compose *environment.DockerCompose, before []models.Session, action func()) *models.Session {
	t.Helper()

	known := make(map[string]struct{}, len(before))
	for _, session := range before {
		known[session.UID] = struct{}{}
	}

	action()

	opened := &models.Session{} //nolint:exhaustruct // filled from the API response

	require.EventuallyWithT(t, func(tt *assert.CollectT) {
		sessions, err := listSessions(ctx, compose)
		if !assert.NoError(tt, err) {
			return
		}

		for _, session := range sessions {
			if _, seen := known[session.UID]; !seen {
				*opened = session

				return
			}
		}

		assert.Fail(tt, "the action opened no session the list did not already carry")
	}, 30*time.Second, 1*time.Second)

	return opened
}

func sessionDetail(t *testing.T, ctx context.Context, compose *environment.DockerCompose, uid string) models.Session {
	t.Helper()

	return awaitSession(t, ctx, compose, uid, 30*time.Second, func(tt *assert.CollectT, session models.Session) {
		assert.NotEmpty(tt, session.Events.Items, "the timeline is written after the channel opens")
	})
}

func itemOfType(events []models.SessionEvent, eventType models.SessionEventType) *models.SessionEvent {
	for i, event := range events {
		if event.Type == eventType {
			return &events[i]
		}
	}

	return nil
}

func payload[T any](t *testing.T, event *models.SessionEvent, key string) T {
	t.Helper()

	require.NotNil(t, event, "no event of the wanted type is in the timeline")

	data, ok := event.Data.(map[string]any)
	require.True(t, ok, "the event payload is not an object: %T", event.Data)

	value, ok := data[key].(T)
	require.True(t, ok, "the payload carries no %T under %q: %v", value, key, data)

	return value
}

func commandsBySeat(events []models.SessionEvent) map[int]string {
	commands := make(map[int]string)

	for _, event := range events {
		if event.Type != models.SessionEventTypeExec {
			continue
		}

		if data, ok := event.Data.(map[string]any); ok {
			commands[event.Seat], _ = data["command"].(string)
		}
	}

	return commands
}

func TestSessionDetailSaysWhatTheSessionDid(t *testing.T) {
	t.Parallel()

	ctx := context.Background()

	compose := newSSHEnvironment(t, ctx, legacyNamespace)
	_, device := startAcceptedAgent(t, ctx, compose)
	signer := registerDeviceKey(t, ctx, compose)

	cases := []struct {
		name   string
		open   func(t *testing.T, conn *ssh.Client)
		expect func(t *testing.T, session models.Session)
	}{
		{
			name: "an interactive shell names the terminal the client asked for",
			open: func(t *testing.T, conn *ssh.Client) {
				t.Helper()

				sess, err := conn.NewSession()
				require.NoError(t, err)

				require.NoError(t, sess.RequestPty("xterm-256color", 30, 141, ssh.TerminalModes{ssh.ECHO: 1}))
				require.NoError(t, sess.Shell())
			},
			expect: func(t *testing.T, session models.Session) {
				t.Helper()

				assert.Equal(t, models.SessionEventTypePtyRequest, session.Events.First)
				assert.Equal(t, "xterm-256color",
					payload[string](t, itemOfType(session.Events.Items, models.SessionEventTypePtyRequest), "term"))
			},
		},
		{
			name: "a single command carries the command it ran",
			open: func(t *testing.T, conn *ssh.Client) {
				t.Helper()

				sess, err := conn.NewSession()
				require.NoError(t, err)

				_, err = sess.CombinedOutput("uptime")
				require.NoError(t, err)
			},
			expect: func(t *testing.T, session models.Session) {
				t.Helper()

				assert.Equal(t, models.SessionEventTypeExec, session.Events.First)
				assert.Equal(t, "uptime",
					payload[string](t, itemOfType(session.Events.Items, models.SessionEventTypeExec), "command"))
			},
		},
		{
			name: "an SFTP transfer opens with the subsystem",
			open: func(t *testing.T, conn *ssh.Client) {
				t.Helper()

				client, err := sftp.NewClient(conn)
				require.NoError(t, err)

				file, err := client.OpenFile("/tmp/sent-by-sftp", os.O_WRONLY|os.O_CREATE|os.O_TRUNC)
				require.NoError(t, err)

				_, err = file.Write([]byte("sent"))
				require.NoError(t, err)

				require.NoError(t, file.Close())
				require.NoError(t, client.Close())
			},
			expect: func(t *testing.T, session models.Session) {
				t.Helper()

				assert.Equal(t, models.SessionEventTypeSubsystem, session.Events.First)
			},
		},
		{
			name: "a port forward records the destination it reached",
			open: func(t *testing.T, conn *ssh.Client) {
				t.Helper()

				forward := struct {
					DestAddr   string
					DestPort   uint32
					OriginAddr string
					OriginPort uint32
				}{DestAddr: "127.0.0.1", DestPort: 80, OriginAddr: "127.0.0.1", OriginPort: 40000}

				ch, requests, err := conn.OpenChannel("direct-tcpip", ssh.Marshal(forward))
				require.NoError(t, err)
				go ssh.DiscardRequests(requests)

				_, err = ch.Write([]byte("GET / HTTP/1.0\r\n\r\n"))
				require.NoError(t, err)

				status, err := bufio.NewReader(ch).ReadString('\n')
				require.NoError(t, err)
				require.True(t, strings.HasPrefix(status, "HTTP/1."), status)

				_ = ch.Close()
			},
			expect: func(t *testing.T, session models.Session) {
				t.Helper()

				forward := itemOfType(session.Events.Items, models.SessionEventTypeDirectTCPIP)
				assert.Equal(t, "127.0.0.1", payload[string](t, forward, "dest_addr"))
				assert.InDelta(t, 80, payload[float64](t, forward, "dest_port"), 0)
				assert.Contains(t, session.Events.Types, string(models.SessionEventTypeDirectTCPIP))
				assert.Empty(t, session.Events.First, "a forward opens no session channel")
			},
		},
		{
			name: "a legacy scp is told from a plain command by its own payload",
			open: func(t *testing.T, conn *ssh.Client) {
				t.Helper()

				client, err := scp.NewClientBySSH(conn)
				require.NoError(t, err)

				file := bytes.NewBuffer(make([]byte, 64))
				require.NoError(t, client.CopyFilePassThru(ctx, file, "/tmp/sent-by-scp", "0644", io.LimitReader))

				client.Close()
			},
			expect: func(t *testing.T, session models.Session) {
				t.Helper()

				assert.Equal(t, models.SessionEventTypeExec, session.Events.First)
				assert.Contains(t,
					payload[string](t, itemOfType(session.Events.Items, models.SessionEventTypeExec), "command"),
					"scp",
					"only the payload separates a legacy transfer from an ordinary command")
			},
		},
		{
			name: "a resized terminal records its new size",
			open: func(t *testing.T, conn *ssh.Client) {
				t.Helper()

				sess, err := conn.NewSession()
				require.NoError(t, err)

				require.NoError(t, sess.RequestPty("xterm", 24, 80, ssh.TerminalModes{ssh.ECHO: 0}))

				startResizeReporter(t, sess)(50, 132)
			},
			expect: func(t *testing.T, session models.Session) {
				t.Helper()

				resize := itemOfType(session.Events.Items, models.SessionEventTypeWindowChange)
				assert.InDelta(t, 132, payload[float64](t, resize, "columns"), 0)
				assert.InDelta(t, 50, payload[float64](t, resize, "rows"), 0)
				assert.Contains(t, session.Events.Types, string(models.SessionEventTypeWindowChange))
			},
		},
		{
			name: "a command records the status it exited with",
			open: func(t *testing.T, conn *ssh.Client) {
				t.Helper()

				sess, err := conn.NewSession()
				require.NoError(t, err)

				var exit *ssh.ExitError

				require.ErrorAs(t, sess.Run("exit 3"), &exit)
				require.Equal(t, 3, exit.ExitStatus())
			},
			expect: func(t *testing.T, session models.Session) {
				t.Helper()

				assert.InDelta(t, 3,
					payload[float64](t, itemOfType(session.Events.Items, models.SessionEventTypeExitStatus), "status"), 0)
				assert.Contains(t, session.Events.Types, string(models.SessionEventTypeExitStatus))
			},
		},
		{
			name: "channels on one connection take a seat each",
			open: func(t *testing.T, conn *ssh.Client) {
				t.Helper()

				for _, command := range []string{"echo first", "echo second"} {
					sess, err := conn.NewSession()
					require.NoError(t, err)

					_, err = sess.CombinedOutput(command)
					require.NoError(t, err)
				}
			},
			expect: func(t *testing.T, session models.Session) {
				t.Helper()

				assert.ElementsMatch(t, []int{0, 1}, session.Events.Seats)
				assert.Equal(t, map[int]string{0: "echo first", 1: "echo second"}, commandsBySeat(session.Events.Items),
					"each command is recorded against the seat its channel took")
			},
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			before := currentSessions(t, ctx, compose)

			conn := dialDevice(t, ctx, compose, device, signer)
			t.Cleanup(func() { _ = conn.Close() })

			opened := sessionAfter(t, ctx, compose, before, func() { tc.open(t, conn) })

			assert.Empty(t, opened.Events.Items, "the list carries no timeline")
			assert.Nil(t, opened.Principal, "a legacy namespace authorizes no principal")

			require.NoError(t, conn.Close())
			requireSessionActive(t, ctx, compose, opened.UID, false)

			session := sessionDetail(t, ctx, compose, opened.UID)
			tc.expect(t, session)

			for i := 1; i < len(session.Events.Items); i++ {
				assert.False(t, session.Events.Items[i].Timestamp.Before(session.Events.Items[i-1].Timestamp),
					"the timeline is oldest first")
			}

			for _, item := range session.Events.Items {
				assert.NotEqual(t, models.SessionEventTypePtyOutput, item.Type,
					"terminal output is the recording, not a timeline entry")
			}
		})
	}
}

func TestSessionPrincipalNamesWhoOpenedIt(t *testing.T) {
	t.Parallel()

	ctx := context.Background()

	compose := newSSHEnvironment(t, ctx, newNamespace)
	_, device := startAcceptedAgent(t, ctx, compose)

	openedWith := func(t *testing.T, signer ssh.Signer) *models.Session {
		t.Helper()

		before := currentSessions(t, ctx, compose)

		conn := dialClient(t, ctx, compose.SSHAddress(), &ssh.ClientConfig{ //nolint:exhaustruct // the remaining fields keep their defaults
			User:            deviceSSHID(device),
			Auth:            []ssh.AuthMethod{ssh.PublicKeys(signer)},
			HostKeyCallback: ssh.InsecureIgnoreHostKey(), //nolint:gosec // the test stack's host key is ephemeral
		})
		t.Cleanup(func() { _ = conn.Close() })

		return sessionAfter(t, ctx, compose, before, func() {
			sess, err := conn.NewSession()
			require.NoError(t, err)

			_, err = sess.CombinedOutput("uptime")
			require.NoError(t, err)
		})
	}

	t.Run("a person", func(t *testing.T) {
		owner := compose.AuthUser(t, ShellHubUsername, ShellHubPassword)

		signer, data := newSigner(t)
		compose.EnrollIdentity(t, "person", data)

		opened := openedWith(t, signer)

		require.NotNil(t, opened.Principal, "a new namespace's session is opened by someone")
		assert.Equal(t, models.Principal{Kind: models.PrincipalUser, ID: owner.ID}, *opened.Principal)
	})

	t.Run("an API key", func(t *testing.T) {
		key := compose.CreateAPIKey(t, &requests.CreateAPIKey{Name: "automation", ExpiresAt: -1})

		signer, data := newSigner(t)

		resp, err := compose.R(ctx).
			SetBody(map[string]any{"name": "automation", "data": data}).
			Post("/api/namespaces/api-key/" + key.Name + "/ssh-identities")
		require.NoError(t, err)
		require.Equal(t, 200, resp.StatusCode(), resp.String())

		compose.CreateAccessPolicy(t, &requests.AccessPolicyCreate{ //nolint:exhaustruct // the remaining fields keep their defaults
			Name:    "automation",
			Subject: requests.AccessPolicySubject{Type: string(models.PolicySubjectAPIKey), Value: key.ID},
			Filter:  requests.AccessPolicyFilter{All: true},
			Logins:  []string{ShellHubAgentUsername},
			Action:  string(models.PolicyActionAllow),
		})

		opened := openedWith(t, signer)

		require.NotNil(t, opened.Principal, "a new namespace's session is opened by someone")
		assert.Equal(t, models.Principal{Kind: models.PrincipalAPIKey, ID: key.ID}, *opened.Principal)
	})
}