package integration

import (
	"encoding/json"
	"sort"
	"strings"
	"testing"

	"github.com/juanfont/headscale/integration/hsic"
	"github.com/oauth2-proxy/mockoidc"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestOIDCGroupsPersisted verifies that the `groups` claim from an OIDC
// provider is persisted into the users.groups column after the user logs in.
//
// The implementation under test:
//   - User.Groups column (TEXT, JSON-encoded []string) added by migration
//     202505141323 in hscontrol/db/db.go.
//   - User.SetGroups / User.GetGroups in hscontrol/types/users.go.
//   - FromClaim() calls SetGroups(claims.Groups) so login populates the column.
//   - OIDCClaims.Groups is FlexibleStringSlice so providers like JumpCloud
//     that return a single string instead of a one-element array also work.
//
// Verification is done by reading the SQLite database inside the headscale
// container directly, because the gRPC User message does not currently
// expose Groups. Adding groups to the gRPC API is a separate, larger change.
func TestOIDCGroupsPersisted(t *testing.T) {
	IntegrationSkip(t)

	// mockoidc serves logins in strict queue order, so keep NodesPerUser=1.
	spec := ScenarioSpec{
		NodesPerUser: 1,
		Users:        []string{"admin", "dev", "solo"},
		OIDCUsers: []mockoidc.MockUser{
			oidcMockUserWithGroups("admin", true, []string{"admins", "engineering"}),
			oidcMockUserWithGroups("dev", true, []string{"engineering"}),
			// User with empty groups — must round-trip as no Groups stored.
			oidcMockUserWithGroups("solo", true, nil),
		},
	}

	scenario, err := NewScenario(spec)
	require.NoError(t, err)
	defer scenario.ShutdownAssertNoPanics(t)

	oidcMap := map[string]string{
		"HEADSCALE_OIDC_ISSUER":             scenario.mockOIDC.Issuer(),
		"HEADSCALE_OIDC_CLIENT_ID":          scenario.mockOIDC.ClientID(),
		"CREDENTIALS_DIRECTORY_TEST":        "/tmp",
		"HEADSCALE_OIDC_CLIENT_SECRET_PATH": "${CREDENTIALS_DIRECTORY_TEST}/hs_client_oidc_secret",
		// Make sure the OIDC scope set includes "groups" so the IdP emits the claim.
		"HEADSCALE_OIDC_SCOPE": "openid,profile,email,groups",
	}

	err = scenario.CreateHeadscaleEnvWithLoginURL(
		nil,
		hsic.WithTestName("oidcgroups"),
		hsic.WithConfigEnv(oidcMap),
		hsic.WithFileInContainer("/tmp/hs_client_oidc_secret", []byte(scenario.mockOIDC.ClientSecret())),
	)
	requireNoErrHeadscaleEnv(t, err)

	// Drive the OIDC login flow for every client.
	_, err = scenario.ListTailscaleClients()
	requireNoErrListClients(t, err)
	err = scenario.WaitForTailscaleSync()
	requireNoErrSync(t, err)

	headscale, err := scenario.Headscale()
	require.NoError(t, err)

	// Query the SQLite database inside the headscale container for the groups
	// column. CLI/gRPC do not expose it yet; this is the authoritative store.
	const dbPath = "/tmp/integration_test_db.sqlite3"
	out, err := headscale.Execute([]string{
		"sqlite3", dbPath,
		"-cmd", ".mode tabs",
		"SELECT name, COALESCE(groups, '') FROM users WHERE provider = 'oidc' ORDER BY name;",
	})
	require.NoError(t, err, "querying users.groups from sqlite")

	got := parseGroupsRows(t, out)

	want := map[string][]string{
		"admin": {"admins", "engineering"},
		"dev":   {"engineering"},
		"solo":  nil,
	}

	for name, wantGroups := range want {
		gotGroups, ok := got[name]
		assert.True(t, ok, "user %q not present in users table", name)
		assert.ElementsMatch(t, wantGroups, gotGroups,
			"groups mismatch for user %q (raw rows: %q)", name, out)
	}
}

// parseGroupsRows parses the tab-separated output of:
//
//	SELECT name, COALESCE(groups, '') FROM users ...
//
// Returns a map of username -> decoded groups slice. An empty groups column
// (stored as "" by SetGroups when the input slice is empty) decodes to nil.
func parseGroupsRows(t *testing.T, raw string) map[string][]string {
	t.Helper()
	rows := map[string][]string{}
	for _, line := range strings.Split(strings.TrimSpace(raw), "\n") {
		if line == "" {
			continue
		}
		parts := strings.SplitN(line, "\t", 2)
		require.Len(t, parts, 2, "unexpected sqlite row format: %q", line)
		name, groupsJSON := parts[0], parts[1]

		if groupsJSON == "" {
			rows[name] = nil
			continue
		}
		var gs []string
		require.NoError(t, json.Unmarshal([]byte(groupsJSON), &gs),
			"groups column for %q is not valid JSON: %q", name, groupsJSON)
		sort.Strings(gs)
		rows[name] = gs
	}
	return rows
}

// oidcMockUserWithGroups extends [oidcMockUser] with a Groups claim.
// mockoidc populates the id_token / userinfo from this struct verbatim.
func oidcMockUserWithGroups(username string, emailVerified bool, groups []string) mockoidc.MockUser {
	u := oidcMockUser(username, emailVerified)
	u.Groups = groups
	return u
}