package runtimegrpc

import (
	"context"
	"errors"
	"net"
	"testing"
	"time"

	sessionsemantic "github.com/pingcap-inc/nexuscrm-contract/pkg/runtime/sessionv1"
	"github.com/pingcap-inc/nexuscrm-contract/wire-go/envelope"
	runtimev1 "github.com/pingcap-inc/nexuscrm-contract/wire-go/gen/runtime/v1"
	"github.com/pingcap-inc/nexuscrm-contract/wire-go/mapping"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/test/bufconn"
)

type sessionLifecycleTestServer struct {
	runtimev1.UnimplementedSessionLifecycleServiceServer
	t *testing.T
}

func (server *sessionLifecycleTestServer) CreateSession(ctx context.Context, request *runtimev1.CreateSessionRequest) (*runtimev1.CreateSessionResponse, error) {
	credentials, ok := RuntimeSessionCreateCredentials(ctx)
	if !ok || credentials.Identity != "user@example.test" || credentials.Secret != "password" || credentials.ClientSecret != "client-secret" {
		server.t.Fatalf("create credentials = %#v, %t", credentials, ok)
	}
	semanticRequest, err := mapping.FromWireCreateSessionRequest(ctx, request)
	if err != nil || semanticRequest.Method != sessionsemantic.AuthenticationMethodPassword {
		server.t.Fatalf("create request = %#v, %v", semanticRequest, err)
	}
	if err := SetIssuedSessionCredentials(ctx, SessionIssuedCredentials{AccessToken: "access-token", ControlProof: "control-proof"}); err != nil {
		server.t.Fatal(err)
	}
	return mapping.ToWireCreateSessionResponse(semanticRequest, sessionsemantic.CreateResponse{
		Session: sessionLifecycleTestView(), AccessTokenIssued: true, ControlProofIssued: true,
	})
}

func (server *sessionLifecycleTestServer) ReadSession(ctx context.Context, request *runtimev1.ReadSessionRequest) (*runtimev1.ReadSessionResponse, error) {
	credentials, ok := RuntimeSessionPresentedCredentials(ctx)
	if !ok || credentials.AccessToken != "access-token" || credentials.ControlProof != "control-proof" {
		server.t.Fatalf("read credentials = %#v, %t", credentials, ok)
	}
	semanticRequest, err := mapping.FromWireReadSessionRequest(ctx, request)
	if err != nil || !semanticRequest.AllowRenewal || !semanticRequest.UseControlProof {
		server.t.Fatalf("read request = %#v, %v", semanticRequest, err)
	}
	if err := SetIssuedSessionCredentials(ctx, SessionIssuedCredentials{AccessToken: "replacement-token"}); err != nil {
		server.t.Fatal(err)
	}
	return mapping.ToWireReadSessionResponse(semanticRequest, sessionsemantic.ReadResponse{
		Session: sessionLifecycleTestView(), ReplacementAccessTokenIssued: true,
		ControlProofExpiresAtUnixSecond: time.Now().Add(time.Hour).Unix(),
	})
}

func (server *sessionLifecycleTestServer) RevokeSession(ctx context.Context, request *runtimev1.RevokeSessionRequest) (*runtimev1.RevokeSessionResponse, error) {
	credentials, ok := RuntimeSessionPresentedCredentials(ctx)
	if !ok || credentials.AccessToken != "access-token" || credentials.ControlProof != "" {
		server.t.Fatalf("revoke credentials = %#v, %t", credentials, ok)
	}
	if _, err := mapping.FromWireRevokeSessionRequest(ctx, request); err != nil {
		server.t.Fatal(err)
	}
	return &runtimev1.RevokeSessionResponse{}, nil
}

func sessionLifecycleTestView() sessionsemantic.Session {
	return sessionsemantic.Session{Principal: sessionsemantic.Principal{
		PrincipalID: "user-1", DisplayName: "Test User", Email: "user@example.test", ProfileID: "profile-1",
	}, ExpiresAtUnixSecond: time.Now().Add(time.Hour).Unix()}
}

func TestSessionLifecycleClientKeepsCredentialsOutOfMessages(t *testing.T) {
	policy := mustIdentityPolicy(t, Config{
		Authority: func(ctx context.Context, _ string) (VerifiedCaller, error) {
			values, _ := metadata.FromIncomingContext(ctx)
			if got := values.Get("authorization"); len(got) != 1 || got[0] != "Bearer workload-secret" {
				t.Fatalf("workload authorization = %#v", got)
			}
			return "sfcrm", nil
		},
		AdmissionLimit: 2, AttemptCap: time.Second,
	}, time.Second)
	listener := bufconn.Listen(4 * envelope.MaxResponseWireBytes)
	server := grpc.NewServer(policy.ServerOptions()...)
	runtimev1.RegisterSessionLifecycleServiceServer(server, &sessionLifecycleTestServer{t: t})
	go func() { _ = server.Serve(listener) }()
	t.Cleanup(func() {
		server.Stop()
		_ = listener.Close()
	})

	dialCtx, cancelDial := context.WithTimeout(context.Background(), time.Second)
	connection, err := grpc.DialContext(dialCtx, "passthrough:///session-bufnet",
		grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return listener.Dial() }),
		grpc.WithTransportCredentials(insecure.NewCredentials()),
	)
	cancelDial()
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = connection.Close() })
	wireClient := runtimev1.NewSessionLifecycleServiceClient(connection)
	client, err := NewSessionLifecycleClient(wireClient.CreateSession, wireClient.ReadSession, wireClient.RevokeSession)
	if err != nil {
		t.Fatal(err)
	}
	baseCtx := metadata.AppendToOutgoingContext(context.Background(), "authorization", "Bearer workload-secret")
	ctx, cancel := context.WithTimeout(baseCtx, time.Second)
	defer cancel()
	createResponse, issued, err := client.CreateSession(ctx, SessionCreateCredentials{
		Identity: "user@example.test", Secret: "password", ClientSecret: "client-secret",
	}, sessionsemantic.CreateRequest{
		AppInstallationID: "installation-1", Method: sessionsemantic.AuthenticationMethodPassword,
		ApplicationClientID: "client-1", IssueControlProof: true,
	})
	if err != nil || !createResponse.AccessTokenIssued || issued.AccessToken != "access-token" || issued.ControlProof != "control-proof" {
		t.Fatalf("create response=%#v issued=%#v err=%v", createResponse, issued, err)
	}
	readResponse, replacement, err := client.ReadSession(ctx, SessionPresentedCredentials{
		AccessToken: "access-token", ControlProof: "control-proof",
	}, sessionsemantic.ReadRequest{AppInstallationID: "installation-1", AllowRenewal: true, UseControlProof: true})
	if err != nil || !readResponse.ReplacementAccessTokenIssued || replacement.AccessToken != "replacement-token" || replacement.ControlProof != "" {
		t.Fatalf("read response=%#v replacement=%#v err=%v", readResponse, replacement, err)
	}
	if err := client.RevokeSession(ctx, SessionPresentedCredentials{AccessToken: "access-token"}, sessionsemantic.RevokeRequest{AppInstallationID: "installation-1"}); err != nil {
		t.Fatal(err)
	}
}

func TestSessionLifecycleClientRejectsCredentialProfileBeforeRPC(t *testing.T) {
	called := false
	client, err := NewSessionLifecycleClient(
		func(context.Context, *runtimev1.CreateSessionRequest, ...grpc.CallOption) (*runtimev1.CreateSessionResponse, error) {
			called = true
			return nil, nil
		},
		func(context.Context, *runtimev1.ReadSessionRequest, ...grpc.CallOption) (*runtimev1.ReadSessionResponse, error) {
			called = true
			return nil, nil
		},
		func(context.Context, *runtimev1.RevokeSessionRequest, ...grpc.CallOption) (*runtimev1.RevokeSessionResponse, error) {
			called = true
			return nil, nil
		},
	)
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	_, _, err = client.CreateSession(ctx, SessionCreateCredentials{Secret: "id-token", Identity: "must-not-be-present"}, sessionsemantic.CreateRequest{
		AppInstallationID: "installation-1", Method: sessionsemantic.AuthenticationMethodFederatedIDToken, FederatedProviderID: "google",
	})
	if err == nil || called {
		t.Fatalf("invalid federated credentials err=%v called=%t", err, called)
	}
	_, _, err = client.ReadSession(ctx, SessionPresentedCredentials{AccessToken: "token"}, sessionsemantic.ReadRequest{
		AppInstallationID: "installation-1", AllowRenewal: true, UseControlProof: true,
	})
	if err == nil || called {
		t.Fatalf("missing control proof err=%v called=%t", err, called)
	}
}

func TestIssuedSessionCredentialsRejectMetadataMismatch(t *testing.T) {
	tests := []struct {
		name                string
		metadata            metadata.MD
		requireAccessToken  bool
		requireControlProof bool
		wantError           bool
	}{
		{name: "valid access", metadata: metadata.Pairs(IssuedSessionAccessTokenMetadataKey, "token"), requireAccessToken: true},
		{name: "valid pair", metadata: metadata.Pairs(IssuedSessionAccessTokenMetadataKey, "token", IssuedSessionControlProofMetadataKey, "proof"), requireAccessToken: true, requireControlProof: true},
		{name: "missing access", metadata: metadata.MD{}, requireAccessToken: true, wantError: true},
		{name: "duplicate access", metadata: metadata.Pairs(IssuedSessionAccessTokenMetadataKey, "one", IssuedSessionAccessTokenMetadataKey, "two"), requireAccessToken: true, wantError: true},
		{name: "unexpected access", metadata: metadata.Pairs(IssuedSessionAccessTokenMetadataKey, "token"), wantError: true},
		{name: "missing proof", metadata: metadata.Pairs(IssuedSessionAccessTokenMetadataKey, "token"), requireAccessToken: true, requireControlProof: true, wantError: true},
		{name: "unexpected proof", metadata: metadata.Pairs(IssuedSessionAccessTokenMetadataKey, "token", IssuedSessionControlProofMetadataKey, "proof"), requireAccessToken: true, wantError: true},
		{name: "noncanonical access", metadata: metadata.Pairs(IssuedSessionAccessTokenMetadataKey, " token"), requireAccessToken: true, wantError: true},
		{name: "noncanonical proof", metadata: metadata.Pairs(IssuedSessionAccessTokenMetadataKey, "token", IssuedSessionControlProofMetadataKey, "proof "), requireAccessToken: true, requireControlProof: true, wantError: true},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			issued, err := issuedSessionCredentials(test.metadata, test.requireAccessToken, test.requireControlProof)
			if (err != nil) != test.wantError {
				t.Fatalf("issued=%#v err=%v, want error=%t", issued, err, test.wantError)
			}
			if err != nil {
				var failure *sessionsemantic.Failure
				if !errors.As(err, &failure) || failure.Category != sessionsemantic.FailureInternal {
					t.Fatalf("mismatch error = %#v", err)
				}
			}
		})
	}
}

func TestIssuedSessionCredentialsAreOwnedBySessionResponseMethods(t *testing.T) {
	tests := []struct {
		name        string
		method      string
		credentials SessionIssuedCredentials
	}{
		{
			name: "untrusted context", credentials: SessionIssuedCredentials{AccessToken: "token"},
		},
		{
			name: "foreign runtime method", method: runtimev1.RecordsService_ReadSelectedRoot_FullMethodName,
			credentials: SessionIssuedCredentials{AccessToken: "token"},
		},
		{
			name: "revoke cannot issue", method: runtimev1.SessionLifecycleService_RevokeSession_FullMethodName,
			credentials: SessionIssuedCredentials{AccessToken: "token"},
		},
		{
			name: "read cannot rotate proof", method: runtimev1.SessionLifecycleService_ReadSession_FullMethodName,
			credentials: SessionIssuedCredentials{AccessToken: "token", ControlProof: "proof"},
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			ctx := context.WithValue(context.Background(), methodContextKey{}, test.method)
			err := SetIssuedSessionCredentials(ctx, test.credentials)
			var failure *sessionsemantic.Failure
			if !errors.As(err, &failure) || failure.Category != sessionsemantic.FailureInternal {
				t.Fatalf("error=%#v", err)
			}
		})
	}
}