// Copyright 2025 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package commands

import (
	"bytes"
	"context"
	"crypto"
	"crypto/rand"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"golang.org/x/crypto/ed25519"

	"github.com/openpubkey/openpubkey/client"
	"github.com/openpubkey/openpubkey/jose"
	"github.com/openpubkey/openpubkey/pktoken"
	"github.com/openpubkey/openpubkey/providers"
	"github.com/openpubkey/openpubkey/util"
	"github.com/openpubkey/opkssh/commands/config"
	"github.com/openpubkey/opkssh/sshcert"
	"github.com/spf13/afero"
	"github.com/stretchr/testify/require"
	"golang.org/x/crypto/ssh"
)

const providerAlias1 = "op1"
const providerIssuer1 = "https://example.com/tokens-1/"
const providerArg1 = providerIssuer1 + ",client-id1234,,"
const providerStr1 = providerAlias1 + "," + providerArg1

const providerAlias2 = "op2"
const providerIssuer2 = "https://auth.issuer/tokens-2/"
const providerArg2 = providerIssuer2 + ",client-id5678,,"
const providerStr2 = providerAlias2 + "," + providerArg2

const providerAlias3 = "op3"
const providerIssuer3 = "https://openidprovider.openidconnect/tokens-3/"
const providerArg3 = providerIssuer3 + ",client-id91011,,"
const providerStr3 = providerAlias3 + "," + providerArg3

const allProvidersStr = providerStr1 + ";" + providerStr2 + ";" + providerStr3

func Mocks(t *testing.T, keyType KeyType, extraClaims ...map[string]any) (*pktoken.PKToken, crypto.Signer, providers.OpenIdProvider) {
	var err error
	var alg jose.KeyAlgorithm
	var signer crypto.Signer

	switch keyType {
	case ECDSA:
		alg = jose.ES256
		signer, err = util.GenKeyPair(alg)
	case ED25519:
		alg = jose.EdDSA
		_, signer, err = ed25519.GenerateKey(rand.Reader)
	}
	require.NoError(t, err)

	providerOpts := providers.DefaultMockProviderOpts()
	op, _, idtTemplate, err := providers.NewMockProvider(providerOpts)
	require.NoError(t, err)

	// Default: include email claim
	if len(extraClaims) > 0 {
		idtTemplate.ExtraClaims = extraClaims[0]
	} else {
		mockEmail := "arthur.aardvark@example.com"
		idtTemplate.ExtraClaims = map[string]any{
			"email": mockEmail,
		}
	}

	client, err := client.New(op, client.WithSigner(signer, alg))
	require.NoError(t, err)

	pkt, err := client.Auth(context.Background())
	require.NoError(t, err)
	return pkt, signer, op
}

func TestLoginCmd(t *testing.T) {
	logDir := "./logs"
	logPath := filepath.Join(logDir, "opkssh.log")

	defaultConfig, err := config.NewClientConfig(config.DefaultClientConfig)
	require.NoError(t, err, "Failed to get default client config")

	_, _, mockOp := Mocks(t, ECDSA)
	configWithAccessToken := &config.ClientConfig{
		Providers: []config.ProviderConfig{
			{
				AliasList:       []string{"mockOp"},
				Issuer:          mockOp.Issuer(),
				SendAccessToken: true,
			},
		},
		DefaultProvider: "mockOp",
	}

	tests := []struct {
		name            string
		envVars         map[string]string
		loginCmd        LoginCmd
		ClientConfig    *config.ClientConfig
		wantAccessToken bool
		wantError       bool
		errorString     string
	}{
		{
			name:    "Good path with no vars",
			envVars: map[string]string{},
			loginCmd: LoginCmd{
				Verbosity:       2,
				PrintIdTokenArg: true,
				LogDirArg:       logDir,
				Config:          defaultConfig,
			},
			wantError: false,
		},
		{
			name:    "Good path (load config)",
			envVars: map[string]string{},
			loginCmd: LoginCmd{
				Verbosity:       2,
				PrintIdTokenArg: true,
				LogDirArg:       logDir,
			},
			wantError: false,
		},
		{
			name:    "Good path PrintKey",
			envVars: map[string]string{},
			loginCmd: LoginCmd{
				Verbosity:   0,
				PrintKeyArg: true,
				LogDirArg:   logDir,
			},
			wantError: false,
		},
		{
			name:    "Good path PrincipalsArg set to nil",
			envVars: map[string]string{},
			loginCmd: LoginCmd{
				Verbosity:     0,
				PrintKeyArg:   true,
				LogDirArg:     logDir,
				PrincipalsArg: nil,
			},
			wantError: false,
		},
		{
			name:    "Good path PrincipalsArg set to empty",
			envVars: map[string]string{},
			loginCmd: LoginCmd{
				Verbosity:     0,
				PrintKeyArg:   true,
				LogDirArg:     logDir,
				PrincipalsArg: []string{},
			},
			wantError: false,
		},
		{
			name:    "Good path PrincipalsArg set to specific principals",
			envVars: map[string]string{},
			loginCmd: LoginCmd{
				Verbosity:     0,
				PrintKeyArg:   true,
				LogDirArg:     logDir,
				PrincipalsArg: []string{"alice", "bob", "root"},
			},
			wantError: false,
		},
		{
			name:    "Good path InspectCert",
			envVars: map[string]string{},
			loginCmd: LoginCmd{
				Verbosity:      0,
				InspectCertArg: true,
				LogDirArg:      logDir,
			},
			wantError: false,
		},
		{
			name:    "Good path with SendAccessToken set in arg and config",
			envVars: map[string]string{},
			loginCmd: LoginCmd{
				Verbosity:          2,
				LogDirArg:          logDir,
				Config:             configWithAccessToken,
				SendAccessTokenArg: true,
			},
			wantAccessToken: true,
			wantError:       false,
		},
		{
			name:    "Good path with SendAccessToken set in config but not in arg",
			envVars: map[string]string{},
			loginCmd: LoginCmd{
				Verbosity:          2,
				LogDirArg:          logDir,
				Config:             configWithAccessToken,
				SendAccessTokenArg: false,
			},
			wantAccessToken: true,
			wantError:       false,
		},
		{
			name:    "Good path with SendAccessToken Arg (issuer not found in config)",
			envVars: map[string]string{},
			loginCmd: LoginCmd{
				Verbosity:          2,
				LogDirArg:          logDir,
				Config:             defaultConfig,
				SendAccessTokenArg: true,
			},
			wantAccessToken: true,
			wantError:       false,
		},
	}
	keyTypes := [...]KeyType{ECDSA, ED25519}

	for _, tt := range tests {
		for _, keyType := range keyTypes {
			t.Run(fmt.Sprintf("%s %s", tt.name, keyType.String()), func(t *testing.T) {
				for k, v := range tt.envVars {
					err := os.Setenv(k, v)
					require.NoError(t, err, "Failed to set env var")
					defer func(key string) {
						_ = os.Unsetenv(key)
					}(k)
				}

				_, _, mockOp := Mocks(t, keyType)
				mockFs := afero.NewMemMapFs()

				tt.loginCmd.overrideProvider = &mockOp
				tt.loginCmd.Fs = mockFs

				// Allows us to capture non-logged CLI output
				cliOutputBuffer := &bytes.Buffer{}
				tt.loginCmd.OutWriter = cliOutputBuffer

				err = tt.loginCmd.Run(context.Background())
				if tt.wantError {
					require.Error(t, err, "Expected error but got none")
					if tt.errorString != "" {
						require.ErrorContains(t, err, tt.errorString, "Got a wrong error message")
					}
				} else {
					require.NoError(t, err, "Unexpected error")

					var pubKeyBytes []byte

					if tt.loginCmd.PrintKeyArg {
						got := cliOutputBuffer.String()
						gotLines := strings.Split(strings.TrimSpace(got), "\n")
						require.GreaterOrEqual(t, len(gotLines), 2, "expected at least 2 lines in output")
						require.Contains(t, gotLines[0], "cert-v01@openssh.com AAAA")
						require.Contains(t, gotLines[1], "-----BEGIN OPENSSH PRIVATE KEY-----")
						pubKeyBytes = []byte(gotLines[0])
					} else if tt.loginCmd.InspectCertArg {
						got := cliOutputBuffer.String()
						require.Contains(t, got, "--- SSH Certificate Information ---")
						require.Contains(t, got, "--- PKToken Structure ---")

						homePath, err := os.UserHomeDir()
						require.NoError(t, err)
						// KeyTypeArg defaults to ECDSA so keys are written to id_ecdsa path
						sshPubPath := filepath.Join(homePath, ".ssh", "id_ecdsa-cert.pub")
						pubKeyBytes, err = afero.ReadFile(mockFs, sshPubPath)
						require.NoError(t, err)
					} else {
						homePath, err := os.UserHomeDir()
						require.NoError(t, err)

						sshPath := filepath.Join(homePath, ".ssh", "id_ecdsa")
						secKeyBytes, err := afero.ReadFile(mockFs, sshPath)
						require.NoError(t, err)
						require.NotNil(t, secKeyBytes)
						require.Contains(t, string(secKeyBytes), "-----BEGIN OPENSSH PRIVATE KEY-----")

						logBytes, err := afero.ReadFile(mockFs, logPath)
						require.NoError(t, err)
						require.NotNil(t, logBytes)
						require.Contains(t, string(logBytes), "running login command with args:")

						sshPubPath := filepath.Join(homePath, ".ssh", "id_ecdsa-cert.pub")
						pubKeyBytes, err = afero.ReadFile(mockFs, sshPubPath)
						require.NoError(t, err)
					}
					certSmug, err := sshcert.NewFromAuthorizedKey("fake-cert-type", string(pubKeyBytes))
					require.NoError(t, err)

					if tt.loginCmd.PrincipalsArg == nil {
						// We use the wildcard principal when PrincipalsArg is nil
						require.Equal(t, []string{"opkssh-wildcard"}, certSmug.SshCert.ValidPrincipals)
					} else if len(tt.loginCmd.PrincipalsArg) == 0 {
						// A list of length 0 gets deserialized as nil in an sshCert
						require.Nil(t, certSmug.SshCert.ValidPrincipals)
					} else {
						require.Equal(t, tt.loginCmd.PrincipalsArg, certSmug.SshCert.ValidPrincipals)
					}

					accToken := certSmug.GetAccessToken()
					if tt.wantAccessToken {
						require.NotEmpty(t, accToken, "expected access token to be set in SSH cert")
					} else {
						require.Empty(t, accToken, "expected access token to not be set in SSH cert")
					}
				}
			})
		}
	}
}

func TestDetermineProvider(t *testing.T) {
	tests := []struct {
		name              string
		envVars           map[string]string
		providerArg       string
		providerAlias     string
		remoteRedirectURI string
		wantIssuer        string
		wantChooser       string
		wantError         bool
		errorString       string
	}{
		{
			name:          "Good path with env vars",
			envVars:       map[string]string{"OPKSSH_DEFAULT": providerAlias1, "OPKSSH_PROVIDERS": providerStr1},
			providerArg:   "",
			providerAlias: "",
			wantIssuer:    providerIssuer1,
			wantError:     false,
		},
		{
			name:          "Good path with env vars and provider arg (provider arg takes precedence)",
			envVars:       map[string]string{"OPKSSH_DEFAULT": providerAlias1, "OPKSSH_PROVIDERS": providerStr1},
			providerArg:   providerArg2,
			providerAlias: "",
			wantIssuer:    providerIssuer2,
			wantError:     false,
		},
		{
			name:          "Good path with env vars and no alias",
			envVars:       map[string]string{"OPKSSH_DEFAULT": providerAlias1, "OPKSSH_PROVIDERS": providerStr1},
			providerArg:   "",
			providerAlias: "",
			wantIssuer:    providerIssuer1,
			wantError:     false,
		},
		{
			name:          "Good path with env vars single provider and no default",
			envVars:       map[string]string{"OPKSSH_DEFAULT": "", "OPKSSH_PROVIDERS": providerStr1},
			providerArg:   "",
			providerAlias: "",
			wantIssuer:    "",
			wantError:     false,
			errorString:   "",
			wantChooser:   `[{"ClientSecret":"","Scopes":["openid profile email"],"PromptType":"consent","AccessType":"offline","RedirectURIs":["http://localhost:3000/login-callback","http://localhost:10001/login-callback","http://localhost:11110/login-callback"],"RemoteRedirectURI":"","GQSign":false,"DeviceFlow":false,"OpenBrowser":false,"HttpClient":null,"IssuedAtOffset":60000000000,"CallbackHTML":"You may now close this window","ExtraURLParamOpts":{}}]`,
		},
		{
			name:          "Good path with env vars many providers and no default",
			envVars:       map[string]string{"OPKSSH_DEFAULT": "", "OPKSSH_PROVIDERS": allProvidersStr},
			providerArg:   "",
			providerAlias: "",
			wantIssuer:    "",
			wantError:     false,
			wantChooser:   `[{"ClientSecret":"","Scopes":["openid profile email"],"PromptType":"consent","AccessType":"offline","RedirectURIs":["http://localhost:3000/login-callback","http://localhost:10001/login-callback","http://localhost:11110/login-callback"],"RemoteRedirectURI":"","GQSign":false,"DeviceFlow":false,"OpenBrowser":false,"HttpClient":null,"IssuedAtOffset":60000000000,"CallbackHTML":"You may now close this window","ExtraURLParamOpts":{}},{"ClientSecret":"","Scopes":["openid profile email"],"PromptType":"consent","AccessType":"offline","RedirectURIs":["http://localhost:3000/login-callback","http://localhost:10001/login-callback","http://localhost:11110/login-callback"],"RemoteRedirectURI":"","GQSign":false,"DeviceFlow":false,"OpenBrowser":false,"HttpClient":null,"IssuedAtOffset":60000000000,"CallbackHTML":"You may now close this window","ExtraURLParamOpts":{}},{"ClientSecret":"","Scopes":["openid profile email"],"PromptType":"consent","AccessType":"offline","RedirectURIs":["http://localhost:3000/login-callback","http://localhost:10001/login-callback","http://localhost:11110/login-callback"],"RemoteRedirectURI":"","GQSign":false,"DeviceFlow":false,"OpenBrowser":false,"HttpClient":null,"IssuedAtOffset":60000000000,"CallbackHTML":"You may now close this window","ExtraURLParamOpts":{}}]`,
		},
		{
			name:          "Good path with env vars many providers and providerAlias",
			envVars:       map[string]string{"OPKSSH_DEFAULT": "", "OPKSSH_PROVIDERS": allProvidersStr},
			providerArg:   "",
			providerAlias: providerAlias2,
			wantIssuer:    providerIssuer2,
			wantError:     false,
		},
		{
			name:          "Good path with env vars many providers and providerAlias",
			envVars:       map[string]string{"OPKSSH_DEFAULT": providerAlias3, "OPKSSH_PROVIDERS": allProvidersStr},
			providerArg:   "",
			providerAlias: "",
			wantIssuer:    providerIssuer3,
			wantError:     false,
		},
		{
			name:              "Good path remoteRedirectURI set (no default)",
			envVars:           map[string]string{"OPKSSH_DEFAULT": "", "OPKSSH_PROVIDERS": allProvidersStr},
			providerArg:       "",
			providerAlias:     "",
			remoteRedirectURI: "https://example.com/login_callback",
			wantChooser:       `[{"ClientSecret":"","Scopes":["openid profile email"],"PromptType":"consent","AccessType":"offline","RedirectURIs":["http://localhost:3000/login-callback","http://localhost:10001/login-callback","http://localhost:11110/login-callback"],"RemoteRedirectURI":"https://example.com/login_callback","GQSign":false,"DeviceFlow":false,"OpenBrowser":false,"HttpClient":null,"IssuedAtOffset":60000000000,"CallbackHTML":"You may now close this window","ExtraURLParamOpts":{}},{"ClientSecret":"","Scopes":["openid profile email"],"PromptType":"consent","AccessType":"offline","RedirectURIs":["http://localhost:3000/login-callback","http://localhost:10001/login-callback","http://localhost:11110/login-callback"],"RemoteRedirectURI":"https://example.com/login_callback","GQSign":false,"DeviceFlow":false,"OpenBrowser":false,"HttpClient":null,"IssuedAtOffset":60000000000,"CallbackHTML":"You may now close this window","ExtraURLParamOpts":{}},{"ClientSecret":"","Scopes":["openid profile email"],"PromptType":"consent","AccessType":"offline","RedirectURIs":["http://localhost:3000/login-callback","http://localhost:10001/login-callback","http://localhost:11110/login-callback"],"RemoteRedirectURI":"https://example.com/login_callback","GQSign":false,"DeviceFlow":false,"OpenBrowser":false,"HttpClient":null,"IssuedAtOffset":60000000000,"CallbackHTML":"You may now close this window","ExtraURLParamOpts":{}}]`,
			wantError:         false,
		},
		{
			name:              "Good path remoteRedirectURI set (with default)",
			envVars:           map[string]string{"OPKSSH_DEFAULT": providerAlias3, "OPKSSH_PROVIDERS": allProvidersStr},
			providerArg:       "",
			providerAlias:     "",
			remoteRedirectURI: "https://example.com/login_callback",
			wantIssuer:        providerIssuer3,
			wantError:         false,
		},
		{
			name:              "Good path remoteRedirectURI set (when provider arg specified)",
			envVars:           map[string]string{"OPKSSH_DEFAULT": providerAlias3, "OPKSSH_PROVIDERS": allProvidersStr},
			providerArg:       providerArg2,
			providerAlias:     "",
			remoteRedirectURI: "https://example.com/login_callback",
			wantIssuer:        providerIssuer2,
			wantError:         false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			for k, v := range tt.envVars {
				err := os.Setenv(k, v)
				require.NoError(t, err, "Failed to set env var")
				defer func(key string) {
					_ = os.Unsetenv(key)
				}(k)
			}

			defaultConfig, err := config.NewClientConfig(config.DefaultClientConfig)
			require.NoError(t, err, "Failed to get default client config")

			loginCmd := LoginCmd{
				DisableBrowserOpenArg: true,
				ProviderArg:           tt.providerArg,
				ProviderAliasArg:      tt.providerAlias,
				PrintIdTokenArg:       true,
				RemoteRedirectURI:     tt.remoteRedirectURI,
				Config:                defaultConfig,
			}

			provider, chooser, err := loginCmd.determineProvider()
			if tt.wantError {
				require.Error(t, err, "Expected error but got none")
				if tt.errorString != "" {
					require.ErrorContains(t, err, tt.errorString, "Got a wrong error message")
				}
			} else {
				require.NoError(t, err, "Unexpected error")
				require.True(t, provider != nil || chooser != nil, "Provider or chooser should never both be nil")
				require.False(t, provider != nil && chooser != nil, "Provider or chooser should never both be non-nil")

				if tt.wantIssuer != "" {
					require.NotNil(t, provider)
				}

				if tt.wantChooser != "" {
					require.NotNil(t, chooser)
				}

				if provider != nil {
					require.Equal(t, provider.Issuer(), tt.wantIssuer)

					if tt.remoteRedirectURI != "" {
						// This only covers the case where a single provider is selected.
						// We handle the chooser case by matching against the expected JSON.
						unwrappedOp, ok := provider.(*providers.StandardOp)
						require.True(t, ok, "Expected provider to be of type StandardOp")
						require.Equal(t, tt.remoteRedirectURI, unwrappedOp.RemoteRedirectURI)
					}
				} else {
					require.NotNil(t, chooser.OpList, "Chooser OpList should not be nil")
					jsonBytes, err := json.Marshal(chooser.OpList)
					require.NoError(t, err)
					require.Equal(t, tt.wantChooser, string(jsonBytes))
				}

			}
		})
	}
}

func TestNewLogin(t *testing.T) {
	autoRefresh := false
	configPathArg := filepath.Join("..", "default-client-config.yml")
	createConfig := false
	configureArg := false
	logDir := "./testdata"
	sendAccessTokenArg := false
	disableBrowserOpenArg := true
	printIdTokenArg := false
	providerArg := ""
	keyPathArg := ""
	providerAlias := ""
	keyAsOutputArg := false
	keyTypeArg := ECDSA
	remoteRedirectURIArg := ""
	var principalsDesired []string = nil

	loginCmd := NewLogin(autoRefresh, configPathArg, createConfig, configureArg, logDir,
		sendAccessTokenArg, disableBrowserOpenArg, printIdTokenArg, providerArg, keyAsOutputArg, keyPathArg, providerAlias, keyTypeArg, remoteRedirectURIArg, false, principalsDesired)
	require.NotNil(t, loginCmd)
}

func TestCreateSSHCert(t *testing.T) {
	tests := []struct {
		name    string
		keyType KeyType
	}{
		{
			name:    "ECDSA Certificate",
			keyType: ECDSA,
		},
		{
			name:    "ED25519 Certificate",
			keyType: ED25519,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {

			pkt, signer, _ := Mocks(t, tt.keyType)
			principals := []string{"guest", "dev"}

			sshCertBytes, signKeyBytes, err := createSSHCert(pkt, signer, principals)
			require.NoError(t, err)
			require.NotNil(t, sshCertBytes)
			require.NotNil(t, signKeyBytes)

			// Simple smoke test to verify we can parse the cert
			certPubkey, _, _, _, err := ssh.ParseAuthorizedKey([]byte("certType" + " " + string(sshCertBytes)))
			require.NoError(t, err)
			require.NotNil(t, certPubkey)
		})
	}
}

func TestIdentityString(t *testing.T) {
	t.Run("with email claim", func(t *testing.T) {
		pkt, _, _ := Mocks(t, ECDSA)
		idString, err := IdentityString(*pkt)
		require.NoError(t, err)
		expIdString := "Email, sub, issuer, audience: \narthur.aardvark@example.com me https://accounts.example.com test_client_id"
		require.Equal(t, expIdString, idString)
	})

	t.Run("without email claim", func(t *testing.T) {
		// Create a mock without email claim by passing empty ExtraClaims
		pkt, _, _ := Mocks(t, ECDSA, map[string]any{})

		idString, err := IdentityString(*pkt)
		require.NoError(t, err)
		require.Contains(t, idString, "WARNING: Email claim is missing from ID token")
		require.Contains(t, idString, "Policies based on email will not work")
		require.Contains(t, idString, "Sub, issuer, audience:")
		require.Contains(t, idString, "me")                           // subject
		require.Contains(t, idString, "https://accounts.example.com") // issuer
		require.Contains(t, idString, "test_client_id")               // audience
	})
}

func TestPrettyPrintIdToken(t *testing.T) {
	pkt, _, _ := Mocks(t, ECDSA)
	iss, err := pkt.Issuer()
	require.NoError(t, err)

	pktStr, err := PrettyIdToken(*pkt)
	require.NoError(t, err)
	require.NotNil(t, pktStr)
	require.Contains(t, pktStr, iss)
}
