package application

import (
	"context"
	"errors"
	"google.golang.org/protobuf/proto"
	"testing"

	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"
	"google.golang.org/grpc"

	gatewayapi "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/api/agent-gateway/gateway/v1"
	gatewayv1 "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/entity/agent-gateway/gateway/v1"
	sessionv1 "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/entity/agent-gateway/session/v1"
	gatewayerrors "github.com/tidbcloud/aisupport/packages/sdk-go/infra/agentgateway/errors"
)

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

	client := NewMockGatewayServiceClient(gomock.NewController(t))
	application := NewApplication(client, ApplicationConfig{
		InstanceID:       "ticket-manager",
		WorkingDirectory: "/workspace",
	})
	require.NoError(t, application.RegisterApplicationTaskEventHandler("dispatch.triage", func(context.Context, *gatewayv1.TaskEvent) error {
		return nil
	}))
	sessionKey, err := EncodeApplicationSessionKey("dispatch.triage", "execution/42")
	require.NoError(t, err)
	request := &gatewayapi.ProcessMessageRequest{
		Message: &gatewayv1.ApplicationMessage{
			MessageId:   "message-1",
			SessionKey:  sessionKey,
			SenderEmail: "user@example.com",
			Content:     "triage this ticket",
		},
		Context: &gatewayapi.MessageContext{AgentId: "codex", AgentModel: "model-override"},
	}
	want := &gatewayapi.ProcessMessageResponse{
		Disposition: gatewayapi.ProcessMessageDisposition_PROCESS_MESSAGE_DISPOSITION_ACCEPTED,
		Task:        &sessionv1.Task{TaskId: "task-1", SessionId: "session-1"},
	}
	ctx := &testContext{Context: context.Background()}
	client.EXPECT().ProcessMessage(gomock.Any(), gomock.Any()).DoAndReturn(
		func(callCtx context.Context, req *gatewayapi.ProcessMessageRequest, _ ...grpc.CallOption) (*gatewayapi.ProcessMessageResponse, error) {
			require.Same(t, ctx, callCtx)
			require.Equal(t, "ticket-manager", req.GetContext().GetInstanceId())
			require.Equal(t, "/workspace", req.GetContext().GetWorkingDirectory())
			require.Equal(t, request.GetMessage(), req.GetMessage())
			require.Equal(t, "codex", req.GetContext().GetAgentId())
			require.Equal(t, "model-override", req.GetContext().GetAgentModel())
			return want, nil
		},
	)

	got, err := application.SendApplicationMessage(ctx, request)
	require.NoError(t, err)
	require.Same(t, want, got)
	require.True(t, proto.Equal(&gatewayapi.MessageContext{AgentId: "codex", AgentModel: "model-override"}, request.GetContext()))
}

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

	client := NewMockGatewayServiceClient(gomock.NewController(t))
	application := NewApplication(client, ApplicationConfig{
		InstanceID:       "ticket-manager",
		WorkingDirectory: "/workspace",
	})
	require.NoError(t, application.RegisterApplicationTaskEventHandler("workspace.timeline", func(context.Context, *gatewayv1.TaskEvent) error {
		return nil
	}))
	sessionKey, err := EncodeApplicationSessionKey("workspace.timeline", "workspace/1")
	require.NoError(t, err)
	client.EXPECT().ProcessMessage(gomock.Any(), gomock.Any()).DoAndReturn(
		func(_ context.Context, req *gatewayapi.ProcessMessageRequest, _ ...grpc.CallOption) (*gatewayapi.ProcessMessageResponse, error) {
			require.Equal(t, "ticket-manager", req.GetContext().GetInstanceId())
			require.Equal(t, "/projects/workspace-1", req.GetContext().GetWorkingDirectory())
			return &gatewayapi.ProcessMessageResponse{}, nil
		},
	)

	_, err = application.SendApplicationMessage(context.Background(), &gatewayapi.ProcessMessageRequest{
		Context: &gatewayapi.MessageContext{InstanceId: "other", WorkingDirectory: "/projects/workspace-1"},
		Message: &gatewayv1.ApplicationMessage{SessionKey: sessionKey},
	})
	require.NoError(t, err)
}

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

	application := NewApplication(NewMockGatewayServiceClient(gomock.NewController(t)), ApplicationConfig{})

	_, err := application.SendApplicationMessage(context.Background(), &gatewayapi.ProcessMessageRequest{})
	require.ErrorIs(t, err, gatewayerrors.ErrInvalidApplicationMessage)

	_, err = application.SendApplicationMessage(context.Background(), &gatewayapi.ProcessMessageRequest{
		Message: &gatewayv1.ApplicationMessage{SessionKey: "invalid"},
	})
	require.ErrorIs(t, err, gatewayerrors.ErrInvalidApplicationMessage)

	sessionKey, encodeErr := EncodeApplicationSessionKey("missing", "session")
	require.NoError(t, encodeErr)
	_, err = application.SendApplicationMessage(context.Background(), &gatewayapi.ProcessMessageRequest{
		Message: &gatewayv1.ApplicationMessage{SessionKey: sessionKey},
	})
	require.ErrorIs(t, err, gatewayerrors.ErrTaskEventHandlerNotFound)
}

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

	client := NewMockGatewayServiceClient(gomock.NewController(t))
	application := NewApplication(client, ApplicationConfig{})
	require.NoError(t, application.RegisterApplicationTaskEventHandler("dispatch.triage", func(context.Context, *gatewayv1.TaskEvent) error {
		return nil
	}))
	sessionKey, err := EncodeApplicationSessionKey("dispatch.triage", "execution/42")
	require.NoError(t, err)
	want := errors.New("gateway unavailable")
	client.EXPECT().ProcessMessage(gomock.Any(), gomock.Any()).Return(nil, want)

	_, err = application.SendApplicationMessage(context.Background(), &gatewayapi.ProcessMessageRequest{
		Message: &gatewayv1.ApplicationMessage{SessionKey: sessionKey},
	})
	require.ErrorIs(t, err, want)
}

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

	application := NewApplication(NewMockGatewayServiceClient(gomock.NewController(t)), ApplicationConfig{InstanceID: "ticket-manager"})
	called := false
	require.NoError(t, application.RegisterApplicationTaskEventHandler("dispatch.triage", func(ctx context.Context, event *gatewayv1.TaskEvent) error {
		called = true
		require.Equal(t, "value", ctx.Value(testContextKey{}))
		require.Equal(t, "task-1", event.GetTaskId())
		return nil
	}))
	event := applicationTaskEvent("dispatch.triage:execution/42")

	err := application.HandleApplicationTaskEvent(context.WithValue(context.Background(), testContextKey{}, "value"), event)
	require.NoError(t, err)
	require.True(t, called)
}

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

	application := NewApplication(NewMockGatewayServiceClient(gomock.NewController(t)), ApplicationConfig{InstanceID: "ticket-manager"})
	want := errors.New("persist failed")
	require.NoError(t, application.RegisterApplicationTaskEventHandler("workspace.timeline", func(context.Context, *gatewayv1.TaskEvent) error {
		return want
	}))

	err := application.HandleApplicationTaskEvent(context.Background(), applicationTaskEvent("workspace.timeline:workspace/1"))
	require.ErrorIs(t, err, want)
}

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

	application := NewApplication(NewMockGatewayServiceClient(gomock.NewController(t)), ApplicationConfig{InstanceID: "ticket-manager"})
	tests := []*gatewayv1.TaskEvent{
		nil,
		applicationTaskEvent("invalid"),
		applicationTaskEvent("missing:session"),
		applicationTaskEvent("handler:session"),
	}
	tests[1].SessionSource.Source.InstanceId = "other"
	tests[3].SessionSource.Source.Type = sessionv1.MessageSourceType_MESSAGE_SOURCE_TYPE_LARK

	for _, event := range tests {
		err := application.HandleApplicationTaskEvent(context.Background(), event)
		require.Error(t, err)
	}
	for _, event := range []*gatewayv1.TaskEvent{tests[0], tests[1], tests[3]} {
		require.ErrorIs(t, application.HandleApplicationTaskEvent(context.Background(), event), gatewayerrors.ErrInvalidTaskEvent)
	}
	require.ErrorIs(t, application.HandleApplicationTaskEvent(context.Background(), tests[2]), gatewayerrors.ErrTaskEventHandlerNotFound)
}

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

	application := NewApplication(NewMockGatewayServiceClient(gomock.NewController(t)), ApplicationConfig{})
	handler := func(context.Context, *gatewayv1.TaskEvent) error { return nil }
	require.NoError(t, application.RegisterApplicationTaskEventHandler("workspace.timeline", handler))
	require.ErrorIs(t, application.RegisterApplicationTaskEventHandler("workspace.timeline", handler), gatewayerrors.ErrTaskEventHandlerAlreadyRegistered)
	require.ErrorIs(t, application.RegisterApplicationTaskEventHandler("invalid:name", handler), gatewayerrors.ErrInvalidTaskEvent)
	require.ErrorIs(t, application.RegisterApplicationTaskEventHandler("", handler), gatewayerrors.ErrInvalidTaskEvent)
	require.ErrorIs(t, application.RegisterApplicationTaskEventHandler("nil", nil), gatewayerrors.ErrInvalidTaskEvent)
}

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

	encoded, err := EncodeApplicationSessionKey("dispatch.triage", "execution:42")
	require.NoError(t, err)
	require.Equal(t, "dispatch.triage:execution:42", encoded)

	handlerName, sessionKey, err := DecodeApplicationSessionKey(encoded)
	require.NoError(t, err)
	require.Equal(t, "dispatch.triage", handlerName)
	require.Equal(t, "execution:42", sessionKey)

	_, err = EncodeApplicationSessionKey("bad:name", "session")
	require.ErrorIs(t, err, gatewayerrors.ErrInvalidApplicationMessage)
	_, _, err = DecodeApplicationSessionKey("invalid")
	require.ErrorIs(t, err, gatewayerrors.ErrInvalidApplicationMessage)
}

func applicationTaskEvent(sessionKey string) *gatewayv1.TaskEvent {
	return &gatewayv1.TaskEvent{
		Type:      gatewayv1.TaskEventType_TASK_EVENT_TYPE_FINISHED,
		TaskId:    "task-1",
		SessionId: "session-1",
		TaskSource: &sessionv1.TaskSource{
			SourceMessageId: "message-1",
			SenderEmail:     "user@example.com",
		},
		Status:  sessionv1.TaskStatus_TASK_STATUS_SUCCEEDED,
		Content: "done",
		SessionSource: &sessionv1.SessionSource{
			Source: &sessionv1.MessageSourceRef{
				Type:       sessionv1.MessageSourceType_MESSAGE_SOURCE_TYPE_APPLICATION,
				InstanceId: "ticket-manager",
			},
			SourceSessionKey: sessionKey,
			InitiatorEmail:   "user@example.com",
		},
	}
}

type testContextKey struct{}

type testContext struct {
	context.Context
}