package grpc

import (
	"context"

	"github.com/google/uuid"
	router "github.com/rwadurian/mpc-system/api/grpc/router/v1"
	"github.com/rwadurian/mpc-system/pkg/grpcutil"
	"github.com/rwadurian/mpc-system/pkg/logger"
	"github.com/rwadurian/mpc-system/pkg/retry"
	"github.com/rwadurian/mpc-system/services/session-coordinator/application/use_cases"
	"go.uber.org/zap"
)

// MessageRouterClient wraps gRPC client for Message Router service
type MessageRouterClient struct {
	resilientConn *grpcutil.ResilientConn
	retryCfg      retry.Config
}

// NewMessageRouterClient creates a new Message Router gRPC client with auto-reconnection
func NewMessageRouterClient(address string) (*MessageRouterClient, error) {
	config := grpcutil.DefaultClientConfig(address)
	resilientConn, err := grpcutil.NewResilientConn(config)
	if err != nil {
		return nil, err
	}

	logger.Info("Connected to Message Router with keepalive",
		zap.String("address", address),
		zap.Duration("keepalive_time", config.KeepaliveTime),
		zap.Bool("auto_reconnect", config.EnableReconnect))

	return &MessageRouterClient{
		resilientConn: resilientConn,
		retryCfg:      retry.DefaultConfig(),
	}, nil
}

// getClient returns a new MessageRouterClient using the current connection
func (c *MessageRouterClient) getClient() router.MessageRouterClient {
	return router.NewMessageRouterClient(c.resilientConn.GetConn())
}

// Close closes the gRPC connection
func (c *MessageRouterClient) Close() error {
	if c.resilientConn != nil {
		return c.resilientConn.Close()
	}
	return nil
}

// IsConnected returns true if the connection is ready
func (c *MessageRouterClient) IsConnected() bool {
	return c.resilientConn != nil && c.resilientConn.IsConnected()
}

// PublishSessionEvent publishes a session event to all subscribed parties
// Includes automatic retry with exponential backoff for transient failures
func (c *MessageRouterClient) PublishSessionEvent(
	ctx context.Context,
	event *router.SessionEvent,
) error {
	logger.Info("Publishing session event to Message Router",
		zap.String("event_type", event.EventType),
		zap.String("session_id", event.SessionId),
		zap.Int("selected_parties_count", len(event.SelectedParties)))

	req := &router.PublishSessionEventRequest{
		Event: event,
	}

	return retry.DoVoid(ctx, c.retryCfg, "PublishSessionEvent", func() error {
		resp, err := c.getClient().PublishSessionEvent(ctx, req)
		if err != nil {
			return err
		}

		logger.Info("Session event published successfully",
			zap.String("event_type", event.EventType),
			zap.String("session_id", event.SessionId),
			zap.Int32("subscriber_count", resp.SubscriberCount))

		return nil
	})
}

// GetRegisteredParties retrieves registered parties from Message Router
// Includes automatic retry with exponential backoff for transient failures
func (c *MessageRouterClient) GetRegisteredParties(
	ctx context.Context,
	roleFilter string,
) ([]*router.RegisteredParty, error) {
	req := &router.GetRegisteredPartiesRequest{
		RoleFilter: roleFilter,
		OnlyOnline: true,
	}

	return retry.Do(ctx, c.retryCfg, "GetRegisteredParties", func() ([]*router.RegisteredParty, error) {
		resp, err := c.getClient().GetRegisteredParties(ctx, req)
		if err != nil {
			return nil, err
		}

		logger.Debug("Retrieved registered parties from Message Router",
			zap.String("role_filter", roleFilter),
			zap.Int32("count", resp.TotalCount))

		return resp.Parties, nil
	})
}

// PublishSessionCreated publishes a session_created event
func (c *MessageRouterClient) PublishSessionCreated(
	ctx context.Context,
	sessionID uuid.UUID,
	thresholdN int32,
	thresholdT int32,
	selectedParties []string,
	joinTokens map[string]string,
	messageHash []byte,
	createdAt int64,
	expiresAt int64,
	delegateUserShare *use_cases.DelegateUserShareInfo,
) error {
	event := &router.SessionEvent{
		EventId:         uuid.New().String(),
		EventType:       "session_created",
		SessionId:       sessionID.String(),
		ThresholdN:      thresholdN,
		ThresholdT:      thresholdT,
		SelectedParties: selectedParties,
		JoinTokens:      joinTokens,
		MessageHash:     messageHash,
		CreatedAt:       createdAt,
		ExpiresAt:       expiresAt,
	}

	// Add delegate user share if provided (for sign sessions with delegate)
	if delegateUserShare != nil {
		event.DelegateUserShare = &router.DelegateUserShare{
			DelegatePartyId: delegateUserShare.DelegatePartyID,
			EncryptedShare:  delegateUserShare.EncryptedShare,
			PartyIndex:      delegateUserShare.PartyIndex,
		}
	}

	return c.PublishSessionEvent(ctx, event)
}

// PublishSessionStarted publishes a session_started event when all parties have joined
// CRITICAL: participants contains the complete list of all parties with their indices
// Receivers should use this list for TSS protocol instead of JoinSession response
func (c *MessageRouterClient) PublishSessionStarted(
	ctx context.Context,
	sessionID string,
	thresholdN int32,
	thresholdT int32,
	selectedParties []string,
	joinTokens map[string]string,
	startedAt int64,
	participants []use_cases.SessionParticipantInfo,
) error {
	// Convert participants to proto format
	protoParticipants := make([]*router.PartyInfo, len(participants))
	for i, p := range participants {
		protoParticipants[i] = &router.PartyInfo{
			PartyId:    p.PartyID,
			PartyIndex: p.PartyIndex,
		}
	}

	event := &router.SessionEvent{
		EventId:         uuid.New().String(),
		EventType:       "session_started",
		SessionId:       sessionID,
		ThresholdN:      thresholdN,
		ThresholdT:      thresholdT,
		SelectedParties: selectedParties,
		JoinTokens:      joinTokens,
		CreatedAt:       startedAt,
		Participants:    protoParticipants,
	}

	logger.Info("Publishing session_started event with participants",
		zap.String("session_id", sessionID),
		zap.Int("participant_count", len(participants)))

	return c.PublishSessionEvent(ctx, event)
}

// PublishParticipantJoined publishes a participant_joined event to all parties in the session
// This notifies the initiator and other participants that a new party has joined
func (c *MessageRouterClient) PublishParticipantJoined(
	ctx context.Context,
	sessionID string,
	partyID string,
	selectedParties []string,
	joinedAt int64,
) error {
	logger.Info("Publishing participant_joined event to Message Router",
		zap.String("session_id", sessionID),
		zap.String("joined_party_id", partyID),
		zap.Strings("notify_parties", selectedParties),
		zap.Int64("joined_at", joinedAt))

	event := &router.SessionEvent{
		EventId:         uuid.New().String(),
		EventType:       "participant_joined",
		SessionId:       sessionID,
		SelectedParties: selectedParties,
		CreatedAt:       joinedAt,
		// Note: We could add a custom field for the joined party ID, but for now
		// the event itself indicates someone joined. The initiator can refresh
		// their participant list via API if needed.
	}

	return c.PublishSessionEvent(ctx, event)
}
