package server

import (
	"context"
	"errors"
	"fmt"
	"testing"

	grpc_ctxtags "github.com/grpc-ecosystem/go-grpc-middleware/tags"
	"github.com/oklog/ulid/v2"
	"github.com/stretchr/testify/require"
	"go.uber.org/goleak"

	openfgav1 "github.com/openfga/api/proto/openfga/v1"

	"github.com/openfga/openfga/cmd/util"
	"github.com/openfga/openfga/internal/condition"
	"github.com/openfga/openfga/internal/graph"
	"github.com/openfga/openfga/pkg/featureflags"
	"github.com/openfga/openfga/pkg/server/commands"
	"github.com/openfga/openfga/pkg/server/config"
	"github.com/openfga/openfga/pkg/testutils"
	"github.com/openfga/openfga/pkg/tuple"
)

func TestBatchCheckUsesTypesystemModel(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	_, ds, _ := util.MustBootstrapDatastore(t, "memory")

	s := MustNewServerWithOpts(WithDatastore(ds))
	t.Cleanup(s.Close)

	createStoreResp, err := s.CreateStore(context.Background(), &openfgav1.CreateStoreRequest{
		Name: "openfga-test",
	})
	require.NoError(t, err)

	storeID := createStoreResp.GetId()

	model := testutils.MustTransformDSLToProtoWithID(`
		model
			schema 1.1

		type user

		type document
			relations
				define viewer: [user]
	`)

	writeAuthModelResp, err := s.WriteAuthorizationModel(context.Background(), &openfgav1.WriteAuthorizationModelRequest{
		StoreId:         storeID,
		SchemaVersion:   model.GetSchemaVersion(),
		TypeDefinitions: model.GetTypeDefinitions(),
	})
	require.NoError(t, err)

	modelID := writeAuthModelResp.GetAuthorizationModelId()

	_, err = s.Write(context.Background(), &openfgav1.WriteRequest{
		StoreId: storeID,
		Writes: &openfgav1.WriteRequestWrites{
			TupleKeys: []*openfgav1.TupleKey{
				tuple.NewTupleKey("document:1", "viewer", "user:anne"),
			},
		},
	})
	require.NoError(t, err)

	// no model ID in this request
	batchCheckRequest := &openfgav1.BatchCheckRequest{
		StoreId: storeID,
		Checks: []*openfgav1.BatchCheckItem{
			{
				TupleKey: &openfgav1.CheckRequestTupleKey{
					User:     "user:anne",
					Relation: "viewer",
					Object:   "document:1",
				},
				CorrelationId: "id1",
			},
			{
				TupleKey: &openfgav1.CheckRequestTupleKey{
					User:     "user:bob",
					Relation: "viewer",
					Object:   "document:1",
				},
				CorrelationId: "id2",
			},
		},
	}
	firstResponse, err := s.BatchCheck(context.Background(), batchCheckRequest)
	require.NoError(t, err)

	// Now add the auth model ID to the request and send it again
	batchCheckRequest.AuthorizationModelId = modelID
	secondResponse, err := s.BatchCheck(context.Background(), batchCheckRequest)
	require.NoError(t, err)

	firstResult := firstResponse.GetResult()
	secondResult := secondResponse.GetResult()

	// Both responses should be identical
	for k, v := range firstResult {
		require.Equal(t, v, secondResult[k])
	}
}

func TestBatchCheckValidatesInboundRequest(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	scenarios := map[string]struct {
		request       *openfgav1.BatchCheckRequest
		errorContains string
	}{
		`test_invalid_correlation_id`: {
			request: &openfgav1.BatchCheckRequest{
				Checks: []*openfgav1.BatchCheckItem{
					{
						TupleKey: &openfgav1.CheckRequestTupleKey{
							User:     "user:bob",
							Relation: "viewer",
							Object:   "document:1",
						},
						CorrelationId: "", // this is invalid and should fail
					},
				},
			},
			errorContains: "invalid BatchCheckItem.CorrelationId",
		},
		`test_empty_checks`: {
			request: &openfgav1.BatchCheckRequest{
				Checks: []*openfgav1.BatchCheckItem{},
			},
			errorContains: "invalid BatchCheckRequest.Checks",
		},
		`test_empty_tuple_key`: {
			request: &openfgav1.BatchCheckRequest{
				Checks: []*openfgav1.BatchCheckItem{
					{
						TupleKey: nil,
					},
				},
			},
			errorContains: "invalid BatchCheckItem.TupleKey",
		},
	}

	_, ds, _ := util.MustBootstrapDatastore(t, "memory")
	s := MustNewServerWithOpts(WithDatastore(ds))
	t.Cleanup(s.Close)

	for name, scenario := range scenarios {
		t.Run(name, func(t *testing.T) {
			batchReq := scenario.request
			batchReq.StoreId = ulid.Make().String()
			_, err := s.BatchCheck(context.Background(), batchReq)
			require.ErrorContains(t, err, scenario.errorContains)
		})
	}
}

func TestBatchCheckFailsIfTooManyChecks(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	numChecks := config.DefaultMaxChecksPerBatchCheck + 1
	_, ds, _ := util.MustBootstrapDatastore(t, "memory")
	s := MustNewServerWithOpts(WithDatastore(ds))
	t.Cleanup(s.Close)

	createStoreResp, err := s.CreateStore(context.Background(), &openfgav1.CreateStoreRequest{
		Name: "openfga-test",
	})
	require.NoError(t, err)

	storeID := createStoreResp.GetId()

	model := testutils.MustTransformDSLToProtoWithID(`
		model
			schema 1.1

		type user

		type document
			relations
				define viewer: [user]
	`)

	_, err = s.WriteAuthorizationModel(context.Background(), &openfgav1.WriteAuthorizationModelRequest{
		StoreId:         storeID,
		SchemaVersion:   model.GetSchemaVersion(),
		TypeDefinitions: model.GetTypeDefinitions(),
	})
	require.NoError(t, err)

	checks := make([]*openfgav1.BatchCheckItem, numChecks)
	for i := 0; i < numChecks; i++ {
		checks[i] = &openfgav1.BatchCheckItem{
			TupleKey: &openfgav1.CheckRequestTupleKey{
				Object:   "doc:doc1",
				Relation: "viewer",
				User:     "user:justin",
			},
			CorrelationId: fmt.Sprintf("abc%d", i),
		}
	}

	request := &openfgav1.BatchCheckRequest{
		StoreId: storeID,
		Checks:  checks,
	}

	_, err = s.BatchCheck(context.Background(), request)
	msg := fmt.Sprintf(
		"batchCheck received %d checks, the maximum allowed is %d",
		numChecks,
		config.DefaultMaxChecksPerBatchCheck,
	)
	require.ErrorContains(t, err, msg)
}
func TestTransformCheckCommandErrorToBatchCheckError(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	errMsg := "oh_no"
	scenarios := map[string]struct {
		inputError     error
		expectedOutput *openfgav1.CheckError
	}{
		`test_invalid_relation_error`: {
			inputError: &commands.InvalidRelationError{Cause: errors.New(errMsg)},
			expectedOutput: &openfgav1.CheckError{
				Code:    &openfgav1.CheckError_InputError{InputError: openfgav1.ErrorCode_validation_error},
				Message: "invalid relation: " + errMsg,
			},
		},
		`test_invalid_tuple_error`: {
			inputError: &commands.InvalidTupleError{Cause: errors.New(errMsg)},
			expectedOutput: &openfgav1.CheckError{
				Code:    &openfgav1.CheckError_InputError{InputError: openfgav1.ErrorCode_invalid_tuple},
				Message: "invalid tuple: " + errMsg,
			},
		},
		`test_resolution_depth_error`: {
			inputError: graph.ErrResolutionDepthExceeded,
			expectedOutput: &openfgav1.CheckError{
				Code:    &openfgav1.CheckError_InputError{InputError: openfgav1.ErrorCode_authorization_model_resolution_too_complex},
				Message: graph.ErrResolutionDepthExceeded.Error(),
			},
		},
		`test_condition_error`: {
			inputError: condition.ErrEvaluationFailed,
			expectedOutput: &openfgav1.CheckError{
				Code:    &openfgav1.CheckError_InputError{InputError: openfgav1.ErrorCode_validation_error},
				Message: condition.ErrEvaluationFailed.Error(),
			},
		},
		`test_throttled_error`: {
			inputError: &commands.ThrottledError{Cause: errors.New(errMsg)},
			expectedOutput: &openfgav1.CheckError{
				Code:    &openfgav1.CheckError_InputError{InputError: openfgav1.ErrorCode_validation_error},
				Message: "throttled: " + errMsg,
			},
		},
		`test_deadline_exceeded`: {
			inputError: context.DeadlineExceeded,
			expectedOutput: &openfgav1.CheckError{
				Code:    &openfgav1.CheckError_InternalError{InternalError: openfgav1.InternalErrorCode_deadline_exceeded},
				Message: context.DeadlineExceeded.Error(),
			},
		},
		`test_generic_error`: {
			inputError: errors.New(errMsg),
			expectedOutput: &openfgav1.CheckError{
				Code:    &openfgav1.CheckError_InternalError{InternalError: openfgav1.InternalErrorCode_internal_error},
				Message: errMsg,
			},
		},
	}

	for name, scenario := range scenarios {
		t.Run(name, func(t *testing.T) {
			result := transformCheckCommandErrorToBatchCheckError(scenario.inputError)
			require.Equal(t, scenario.expectedOutput, result)
		})
	}
}

func TestTransformCheckResultToProto(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	// Create two fake outcomes, one happy path and one error
	// These two make 100% coverage
	happyOutcome := &commands.BatchCheckOutcome{Allowed: true}
	sadOutcome := &commands.BatchCheckOutcome{Err: graph.ErrResolutionDepthExceeded}

	// format the expected final output of the transform function
	expectedHappyResult := &openfgav1.BatchCheckSingleResult{
		CheckResult: &openfgav1.BatchCheckSingleResult_Allowed{Allowed: true},
	}

	expectedSadResult := &openfgav1.BatchCheckSingleResult{
		CheckResult: &openfgav1.BatchCheckSingleResult_Error{
			Error: &openfgav1.CheckError{
				Code:    &openfgav1.CheckError_InputError{InputError: openfgav1.ErrorCode_authorization_model_resolution_too_complex},
				Message: graph.ErrResolutionDepthExceeded.Error(),
			},
		},
	}

	t.Run("test_transform_check_result_to_proto", func(t *testing.T) {
		resultOne := transformCheckResultToProto(happyOutcome)
		resultTwo := transformCheckResultToProto(sadOutcome)
		require.Equal(t, expectedHappyResult, resultOne)
		require.Equal(t, expectedSadResult, resultTwo)
	})
}

// TestBatchCheckV2UsedWhenFlagEnabled verifies that when ExperimentalWeightedGraphCheck is
// enabled, BatchCheck routes through v2 and the v2_check_count tag is set on the context.
func TestBatchCheckV2UsedWhenFlagEnabled(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	_, ds, _ := util.MustBootstrapDatastore(t, "memory")

	s := MustNewServerWithOpts(
		WithDatastore(ds),
		WithFeatureFlagClient(featureflags.NewDefaultClient([]string{config.ExperimentalWeightedGraphCheck})),
	)
	t.Cleanup(s.Close)

	ctx := context.Background()

	createStoreResp, err := s.CreateStore(ctx, &openfgav1.CreateStoreRequest{Name: "test"})
	require.NoError(t, err)
	storeID := createStoreResp.GetId()

	model := testutils.MustTransformDSLToProtoWithID(`
		model
			schema 1.1
		type user
		type document
			relations
				define viewer: [user]
	`)
	writeModelResp, err := s.WriteAuthorizationModel(ctx, &openfgav1.WriteAuthorizationModelRequest{
		StoreId:         storeID,
		SchemaVersion:   model.GetSchemaVersion(),
		TypeDefinitions: model.GetTypeDefinitions(),
	})
	require.NoError(t, err)
	modelID := writeModelResp.GetAuthorizationModelId()

	_, err = s.Write(ctx, &openfgav1.WriteRequest{
		StoreId:              storeID,
		AuthorizationModelId: modelID,
		Writes: &openfgav1.WriteRequestWrites{
			TupleKeys: []*openfgav1.TupleKey{
				tuple.NewTupleKey("document:1", "viewer", "user:alice"),
			},
		},
	})
	require.NoError(t, err)

	ctx = grpc_ctxtags.SetInContext(ctx, grpc_ctxtags.NewTags())

	resp, err := s.BatchCheck(ctx, &openfgav1.BatchCheckRequest{
		StoreId:              storeID,
		AuthorizationModelId: modelID,
		Checks: []*openfgav1.BatchCheckItem{
			{
				TupleKey:      &openfgav1.CheckRequestTupleKey{Object: "document:1", Relation: "viewer", User: "user:alice"},
				CorrelationId: "id1",
			},
			{
				TupleKey:      &openfgav1.CheckRequestTupleKey{Object: "document:1", Relation: "viewer", User: "user:bob"},
				CorrelationId: "id2",
			},
		},
	})
	require.NoError(t, err)
	require.True(t, resp.GetResult()["id1"].GetAllowed())
	require.False(t, resp.GetResult()["id2"].GetAllowed())

	tags := grpc_ctxtags.Extract(ctx).Values()
	v2Count, ok := tags["v2_check_count"]
	require.True(t, ok, "v2_check_count should be set in context tags")
	require.EqualValues(t, uint32(2), v2Count)

	v2Fallback, ok := tags["v2_fallback_count"]
	require.True(t, ok, "v2_fallback_count should be set in context tags")
	require.EqualValues(t, uint32(0), v2Fallback)
}

// TestBatchCheckV2FallbackCountWhenGraphParseFails verifies that when ExperimentalWeightedGraphCheck
// is enabled but the model fails weighted graph construction (ErrInvalidModel), the v2_fallback_count
// tag equals the number of checks in the batch (all fell back to v1).
func TestBatchCheckV2FallbackCountWhenGraphParseFails(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	// An intersection whose branches resolve to different terminal types triggers ErrInvalidModel
	// in the weighted graph builder, causing a batch-level fallback to v1.
	modelDSL := `
		model
			schema 1.1
		type user
		type bot
		type document
			relations
				define user_access: [user]
				define bot_access: [bot]
				define viewer: user_access and bot_access
	`

	_, ds, _ := util.MustBootstrapDatastore(t, "memory")

	s := MustNewServerWithOpts(
		WithDatastore(ds),
		WithFeatureFlagClient(featureflags.NewDefaultClient([]string{config.ExperimentalWeightedGraphCheck})),
	)
	t.Cleanup(s.Close)

	ctx := context.Background()

	createStoreResp, err := s.CreateStore(ctx, &openfgav1.CreateStoreRequest{Name: "test"})
	require.NoError(t, err)
	storeID := createStoreResp.GetId()

	model := testutils.MustTransformDSLToProtoWithID(modelDSL)
	writeModelResp, err := s.WriteAuthorizationModel(ctx, &openfgav1.WriteAuthorizationModelRequest{
		StoreId:         storeID,
		SchemaVersion:   model.GetSchemaVersion(),
		TypeDefinitions: model.GetTypeDefinitions(),
	})
	require.NoError(t, err)
	modelID := writeModelResp.GetAuthorizationModelId()

	ctx = grpc_ctxtags.SetInContext(ctx, grpc_ctxtags.NewTags())

	_, err = s.BatchCheck(ctx, &openfgav1.BatchCheckRequest{
		StoreId:              storeID,
		AuthorizationModelId: modelID,
		Checks: []*openfgav1.BatchCheckItem{
			{
				TupleKey:      &openfgav1.CheckRequestTupleKey{Object: "document:1", Relation: "viewer", User: "user:alice"},
				CorrelationId: "id1",
			},
			{
				TupleKey:      &openfgav1.CheckRequestTupleKey{Object: "document:1", Relation: "viewer", User: "user:bob"},
				CorrelationId: "id2",
			},
		},
	})
	require.NoError(t, err)

	tags := grpc_ctxtags.Extract(ctx).Values()

	v2Count, ok := tags["v2_check_count"]
	require.True(t, ok, "v2_check_count should be set in context tags")
	require.EqualValues(t, uint32(0), v2Count)

	v2Fallback, ok := tags["v2_fallback_count"]
	require.True(t, ok, "v2_fallback_count should be set in context tags")
	require.EqualValues(t, uint32(2), v2Fallback)
}

// transformCheckResultToProto takes ~100-200ns per BatchCheckOutcome, or .0001 - .0002ms per.
// At smaller batch sizes that's fine, but if sizes increase into the thousands the
// transform might be better in its own concurrent routine rather than as post-processing.
func BenchmarkTransformCheckResultToProto(b *testing.B) {
	outcomes := map[commands.CorrelationID]*commands.BatchCheckOutcome{
		"abc123": {
			Allowed: true,
		},
		"def456": {
			Err: graph.ErrResolutionDepthExceeded,
		},
	}

	b.Run("benchmark_transform_check_result_to_proto", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			for _, v := range outcomes {
				transformCheckResultToProto(v)
			}
		}
	})
}
