//go:generate ${GOBIN}/mockgen -source application.go -destination application_mock.go -package=application
//go:generate ${GOBIN}/mockgen -destination gateway_client_mock.go -package=application github.com/tidbcloud/aisupport/packages/sdk-go/contracts/api/agent-gateway/gateway/v1 GatewayServiceClient
package application

import (
	"context"
	"fmt"
	"strings"
	"sync"

	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"

	"google.golang.org/protobuf/proto"
)

const applicationSessionKeySeparator = ":"

type ApplicationConfig struct {
	InstanceID       string
	WorkingDirectory string
}

type ApplicationTaskEventHandler func(context.Context, *gatewayv1.TaskEvent) error

type Application interface {
	SendApplicationMessage(context.Context, *gatewayapi.ProcessMessageRequest) (*gatewayapi.ProcessMessageResponse, error)
	RegisterApplicationTaskEventHandler(string, ApplicationTaskEventHandler) error
	HandleApplicationTaskEvent(context.Context, *gatewayv1.TaskEvent) error
}

type application struct {
	config   ApplicationConfig
	client   gatewayapi.GatewayServiceClient
	mu       sync.RWMutex
	handlers map[string]ApplicationTaskEventHandler
}

func NewApplication(client gatewayapi.GatewayServiceClient, config ApplicationConfig) Application {
	return &application{
		config:   config,
		client:   client,
		handlers: make(map[string]ApplicationTaskEventHandler),
	}
}

func (a *application) SendApplicationMessage(ctx context.Context, req *gatewayapi.ProcessMessageRequest) (*gatewayapi.ProcessMessageResponse, error) {
	if req.GetMessage() == nil {
		return nil, fmt.Errorf("%w: message is required", gatewayerrors.ErrInvalidApplicationMessage)
	}
	handlerName, _, err := DecodeApplicationSessionKey(req.GetMessage().GetSessionKey())
	if err != nil {
		return nil, fmt.Errorf("decode application session key: %w", err)
	}
	if !a.hasHandler(handlerName) {
		return nil, fmt.Errorf("%w: %s", gatewayerrors.ErrTaskEventHandlerNotFound, handlerName)
	}

	request := proto.Clone(req).(*gatewayapi.ProcessMessageRequest)
	if request.Context == nil {
		request.Context = &gatewayapi.MessageContext{}
	}
	request.Context.InstanceId = a.config.InstanceID
	if request.Context.WorkingDirectory == "" {
		request.Context.WorkingDirectory = a.config.WorkingDirectory
	}

	resp, err := a.client.ProcessMessage(ctx, request)
	if err != nil {
		return nil, fmt.Errorf("process application message with agent gateway: %w", err)
	}
	return resp, nil
}

func (a *application) RegisterApplicationTaskEventHandler(name string, handler ApplicationTaskEventHandler) error {
	if name == "" || strings.Contains(name, applicationSessionKeySeparator) || handler == nil {
		return fmt.Errorf("%w: invalid handler registration", gatewayerrors.ErrInvalidTaskEvent)
	}

	a.mu.Lock()
	defer a.mu.Unlock()
	if _, ok := a.handlers[name]; ok {
		return fmt.Errorf("%w: %s", gatewayerrors.ErrTaskEventHandlerAlreadyRegistered, name)
	}
	a.handlers[name] = handler
	return nil
}

func (a *application) HandleApplicationTaskEvent(ctx context.Context, event *gatewayv1.TaskEvent) error {
	if event.GetTaskId() == "" || event.GetSessionId() == "" || event.GetTaskSource().GetSourceMessageId() == "" ||
		event.GetSessionSource().GetSource().GetType() != sessionv1.MessageSourceType_MESSAGE_SOURCE_TYPE_APPLICATION ||
		event.GetSessionSource().GetSource().GetInstanceId() != a.config.InstanceID ||
		event.GetType() == gatewayv1.TaskEventType_TASK_EVENT_TYPE_UNSPECIFIED {
		return fmt.Errorf("%w: missing or mismatched source fields", gatewayerrors.ErrInvalidTaskEvent)
	}

	handlerName, _, err := DecodeApplicationSessionKey(event.GetSessionSource().GetSourceSessionKey())
	if err != nil {
		return fmt.Errorf("%w: %v", gatewayerrors.ErrInvalidTaskEvent, err)
	}
	a.mu.RLock()
	handler := a.handlers[handlerName]
	a.mu.RUnlock()
	if handler == nil {
		return fmt.Errorf("%w: %s", gatewayerrors.ErrTaskEventHandlerNotFound, handlerName)
	}
	return handler(ctx, event)
}

func (a *application) hasHandler(name string) bool {
	a.mu.RLock()
	defer a.mu.RUnlock()
	return a.handlers[name] != nil
}

func EncodeApplicationSessionKey(handlerName, sessionKey string) (string, error) {
	if handlerName == "" || strings.Contains(handlerName, applicationSessionKeySeparator) || sessionKey == "" {
		return "", fmt.Errorf("%w: handler name and session key are required, and handler name cannot contain ':'", gatewayerrors.ErrInvalidApplicationMessage)
	}
	return handlerName + applicationSessionKeySeparator + sessionKey, nil
}

func DecodeApplicationSessionKey(value string) (string, string, error) {
	handlerName, sessionKey, ok := strings.Cut(value, applicationSessionKeySeparator)
	if !ok || handlerName == "" || sessionKey == "" {
		return "", "", fmt.Errorf("%w: application session key must use <handler>:<session-key>", gatewayerrors.ErrInvalidApplicationMessage)
	}
	return handlerName, sessionKey, nil
}

var _ Application = (*application)(nil)