package grpc

import (
	"context"
	"io"
	"sync"
	"time"

	"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/server-party/application/use_cases"
	"go.uber.org/zap"
	"google.golang.org/grpc"
)

// MessageRouterClient implements use_cases.MessageRouterClient
type MessageRouterClient struct {
	resilientConn *grpcutil.ResilientConn
	address       string
	mu            sync.Mutex
	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,
		address:       address,
		retryCfg:      retry.DefaultConfig(),
	}, nil
}

// getConn returns the current gRPC connection
func (c *MessageRouterClient) getConn() *grpc.ClientConn {
	return 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()
}

// WaitForReady waits for the connection to be ready
func (c *MessageRouterClient) WaitForReady(ctx context.Context) bool {
	if c.resilientConn == nil {
		return false
	}
	return c.resilientConn.WaitForReady(ctx)
}

// RouteMessage sends an MPC protocol message to other parties
// Includes automatic retry with exponential backoff for transient failures
func (c *MessageRouterClient) RouteMessage(
	ctx context.Context,
	sessionID uuid.UUID,
	fromParty string,
	toParties []string,
	roundNumber int,
	payload []byte,
) error {
	req := &router.RouteMessageRequest{
		SessionId:   sessionID.String(),
		FromParty:   fromParty,
		ToParties:   toParties,
		RoundNumber: int32(roundNumber),
		MessageType: "tss",
		Payload:     payload,
	}

	return retry.DoVoid(ctx, c.retryCfg, "RouteMessage", func() error {
		resp := &router.RouteMessageResponse{}
		err := c.getConn().Invoke(ctx, "/mpc.router.v1.MessageRouter/RouteMessage", req, resp)
		if err != nil {
			return err
		}

		if !resp.Success {
			logger.Error("Message routing failed",
				zap.String("session_id", sessionID.String()))
			return use_cases.ErrKeygenFailed
		}

		logger.Debug("Message routed successfully",
			zap.String("session_id", sessionID.String()),
			zap.String("from", fromParty),
			zap.Int("to_count", len(toParties)),
			zap.Int("round", roundNumber))

		return nil
	})
}

// SubscribeMessages subscribes to MPC messages for a party
func (c *MessageRouterClient) SubscribeMessages(
	ctx context.Context,
	sessionID uuid.UUID,
	partyID string,
) (<-chan *use_cases.MPCMessage, error) {
	req := &router.SubscribeMessagesRequest{
		SessionId: sessionID.String(),
		PartyId:   partyID,
	}

	// Create a streaming connection
	stream, err := c.createSubscribeStream(ctx, req)
	if err != nil {
		logger.Error("Failed to subscribe to messages",
			zap.Error(err),
			zap.String("session_id", sessionID.String()),
			zap.String("party_id", partyID))
		return nil, err
	}

	// Create output channel
	msgChan := make(chan *use_cases.MPCMessage, 100)

	// Start goroutine to receive messages
	go func() {
		defer close(msgChan)

		for {
			select {
			case <-ctx.Done():
				logger.Debug("Message subscription context cancelled",
					zap.String("session_id", sessionID.String()),
					zap.String("party_id", partyID))
				return
			default:
				msg := &router.MPCMessage{}
				err := stream.RecvMsg(msg)
				if err == io.EOF {
					logger.Debug("Message stream ended",
						zap.String("session_id", sessionID.String()))
					return
				}
				if err != nil {
					logger.Error("Error receiving message",
						zap.Error(err),
						zap.String("session_id", sessionID.String()))
					return
				}

				// Convert to use_cases.MPCMessage
				mpcMsg := &use_cases.MPCMessage{
					MessageID:   msg.MessageId,
					SessionID:   msg.SessionId,
					FromParty:   msg.FromParty,
					IsBroadcast: msg.IsBroadcast,
					RoundNumber: int(msg.RoundNumber),
					Payload:     msg.Payload,
				}

				select {
				case msgChan <- mpcMsg:
					logger.Debug("Received MPC message",
						zap.String("message_id", msg.MessageId),
						zap.String("from", msg.FromParty),
						zap.Int("round", int(msg.RoundNumber)))

					// Send acknowledgment for the received message
					go func(messageID, sessionIDStr, pID string) {
						// Use parent context to allow proper cancellation
						ackCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
						defer cancel()
						sid, _ := uuid.Parse(sessionIDStr)
						if err := c.AcknowledgeMessage(ackCtx, messageID, sid, pID, true, ""); err != nil {
							logger.Warn("Failed to acknowledge message",
								zap.String("message_id", messageID),
								zap.Error(err))
						} else {
							logger.Debug("Message acknowledged",
								zap.String("message_id", messageID))
						}
					}(msg.MessageId, msg.SessionId, partyID)
				case <-ctx.Done():
					return
				}
			}
		}
	}()

	logger.Info("Subscribed to messages",
		zap.String("session_id", sessionID.String()),
		zap.String("party_id", partyID))

	return msgChan, nil
}

// createSubscribeStream creates a streaming connection for message subscription
func (c *MessageRouterClient) createSubscribeStream(
	ctx context.Context,
	req *router.SubscribeMessagesRequest,
) (grpc.ClientStream, error) {
	streamDesc := &grpc.StreamDesc{
		StreamName:    "SubscribeMessages",
		ServerStreams: true,
	}

	stream, err := c.getConn().NewStream(ctx, streamDesc, "/mpc.router.v1.MessageRouter/SubscribeMessages")
	if err != nil {
		return nil, err
	}

	if err := stream.SendMsg(req); err != nil {
		return nil, err
	}

	if err := stream.CloseSend(); err != nil {
		return nil, err
	}

	return stream, nil
}

// GetPendingMessages gets pending messages (polling alternative)
// Includes automatic retry with exponential backoff for transient failures
func (c *MessageRouterClient) GetPendingMessages(
	ctx context.Context,
	sessionID uuid.UUID,
	partyID string,
	afterTimestamp int64,
) ([]*use_cases.MPCMessage, error) {
	req := &router.GetPendingMessagesRequest{
		SessionId:      sessionID.String(),
		PartyId:        partyID,
		AfterTimestamp: afterTimestamp,
	}

	return retry.Do(ctx, c.retryCfg, "GetPendingMessages", func() ([]*use_cases.MPCMessage, error) {
		resp := &router.GetPendingMessagesResponse{}
		err := c.getConn().Invoke(ctx, "/mpc.router.v1.MessageRouter/GetPendingMessages", req, resp)
		if err != nil {
			return nil, err
		}

		messages := make([]*use_cases.MPCMessage, len(resp.Messages))
		for i, msg := range resp.Messages {
			messages[i] = &use_cases.MPCMessage{
				MessageID:   msg.MessageId,
				SessionID:   msg.SessionId,
				FromParty:   msg.FromParty,
				IsBroadcast: msg.IsBroadcast,
				RoundNumber: int(msg.RoundNumber),
				Payload:     msg.Payload,
			}

			// Send acknowledgment for each received message
			go func(messageID, sessionIDStr, pID string) {
				// Use parent context to allow proper cancellation
				ackCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
				defer cancel()
				sid, _ := uuid.Parse(sessionIDStr)
				if err := c.AcknowledgeMessage(ackCtx, messageID, sid, pID, true, ""); err != nil {
					logger.Warn("Failed to acknowledge pending message",
						zap.String("message_id", messageID),
						zap.Error(err))
				}
			}(msg.MessageId, msg.SessionId, partyID)
		}

		return messages, nil
	})
}

// NotificationConfig represents notification channel configuration for a party
// If any notification channel is set, party operates in OFFLINE mode (24h async)
// If no notification channels are set, party operates in REAL-TIME mode (Message Router push)
type NotificationConfig struct {
	Email     string // Optional: email address for session invitations
	Phone     string // Optional: phone number for SMS notifications
	PushToken string // Optional: push notification token (FCM/APNs)
}

// HasAnyChannel returns true if any notification channel is configured
func (nc *NotificationConfig) HasAnyChannel() bool {
	return nc != nil && (nc.Email != "" || nc.Phone != "" || nc.PushToken != "")
}

// RegisterParty registers this party with the message router
// This should be called on startup and whenever party configuration changes
func (c *MessageRouterClient) RegisterParty(
	ctx context.Context,
	partyID string,
	partyRole string,
	version string,
) error {
	return c.RegisterPartyWithNotification(ctx, partyID, partyRole, version, nil)
}

// RegisterPartyWithNotification registers party with optional notification channels
// If notification channels are provided, party operates in offline mode (24h async)
// If no notification channels, party operates in real-time mode (Message Router push)
// Includes automatic retry with exponential backoff for transient failures
//
// This method should be called:
// 1. On startup to register the party
// 2. When notification channels are updated (e.g., user binds email/phone)
// 3. When party configuration changes
func (c *MessageRouterClient) RegisterPartyWithNotification(
	ctx context.Context,
	partyID string,
	partyRole string,
	version string,
	notification *NotificationConfig,
) error {
	req := &router.RegisterPartyRequest{
		PartyId:   partyID,
		PartyRole: partyRole,
		Version:   version,
	}

	// Set notification channel if configured (enables offline mode)
	if notification != nil && notification.HasAnyChannel() {
		req.Notification = &router.NotificationChannel{
			Email:     notification.Email,
			Phone:     notification.Phone,
			PushToken: notification.PushToken,
		}
	}

	// Log registration mode
	mode := "real-time"
	if notification != nil && notification.HasAnyChannel() {
		mode = "offline"
	}

	return retry.DoVoid(ctx, c.retryCfg, "RegisterParty", func() error {
		resp := &router.RegisterPartyResponse{}
		err := c.getConn().Invoke(ctx, "/mpc.router.v1.MessageRouter/RegisterParty", req, resp)
		if err != nil {
			return err
		}

		logger.Info("Party registered successfully",
			zap.String("party_id", partyID),
			zap.String("role", partyRole),
			zap.String("mode", mode),
			zap.Bool("has_email", notification != nil && notification.Email != ""),
			zap.Bool("has_phone", notification != nil && notification.Phone != ""),
			zap.Bool("has_push", notification != nil && notification.PushToken != ""),
			zap.Bool("success", resp.Success))

		return nil
	})
}

// UpdateNotificationChannels re-registers party with new notification channels
// Call this when party's notification configuration changes
// TODO: Implement when notification management is needed
func (c *MessageRouterClient) UpdateNotificationChannels(
	ctx context.Context,
	partyID string,
	partyRole string,
	version string,
	notification *NotificationConfig,
) error {
	logger.Info("Updating party notification channels",
		zap.String("party_id", partyID),
		zap.Bool("has_email", notification != nil && notification.Email != ""),
		zap.Bool("has_phone", notification != nil && notification.Phone != ""),
		zap.Bool("has_push", notification != nil && notification.PushToken != ""))

	return c.RegisterPartyWithNotification(ctx, partyID, partyRole, version, notification)
}

// SubscribeSessionEvents subscribes to session lifecycle events with auto-reconnect
func (c *MessageRouterClient) SubscribeSessionEvents(
	ctx context.Context,
	partyID string,
	eventHandler func(*router.SessionEvent),
) error {
	req := &router.SubscribeSessionEventsRequest{
		PartyId:    partyID,
		EventTypes: []string{}, // Subscribe to all event types
	}

	// Create initial streaming connection
	stream, err := c.createSessionEventStream(ctx, req)
	if err != nil {
		logger.Error("Failed to subscribe to session events",
			zap.Error(err),
			zap.String("party_id", partyID))
		return err
	}

	logger.Info("Subscribed to session events",
		zap.String("party_id", partyID))

	// Start goroutine to receive events with auto-reconnect
	go func() {
		currentStream := stream
		reconnectBackoff := time.Second // Start with 1 second backoff
		maxBackoff := 30 * time.Second

		for {
			select {
			case <-ctx.Done():
				logger.Info("Session event subscription cancelled",
					zap.String("party_id", partyID))
				return
			default:
				event := &router.SessionEvent{}
				err := currentStream.RecvMsg(event)
				if err == io.EOF {
					logger.Warn("Session event stream ended, reconnecting...",
						zap.String("party_id", partyID))
				} else if err != nil {
					logger.Warn("Error receiving session event, reconnecting...",
						zap.Error(err),
						zap.String("party_id", partyID))
				} else {
					// Successfully received event, reset backoff
					reconnectBackoff = time.Second

					logger.Info("Received session event",
						zap.String("event_type", event.EventType),
						zap.String("session_id", event.SessionId),
						zap.String("party_id", partyID))

					// Call event handler
					if eventHandler != nil {
						eventHandler(event)
					}
					continue
				}

				// Reconnect with exponential backoff
				for {
					select {
					case <-ctx.Done():
						return
					case <-time.After(reconnectBackoff):
						logger.Info("Attempting to reconnect session event stream",
							zap.String("party_id", partyID),
							zap.Duration("backoff", reconnectBackoff))

						newStream, err := c.createSessionEventStream(ctx, req)
						if err != nil {
							logger.Error("Failed to reconnect session event stream",
								zap.Error(err),
								zap.String("party_id", partyID))
							// Increase backoff for next attempt
							reconnectBackoff = reconnectBackoff * 2
							if reconnectBackoff > maxBackoff {
								reconnectBackoff = maxBackoff
							}
							continue
						}

						logger.Info("Successfully reconnected to session events",
							zap.String("party_id", partyID))
						currentStream = newStream
						reconnectBackoff = time.Second // Reset backoff on success
						break
					}
					break
				}
			}
		}
	}()

	return nil
}

// createSessionEventStream creates a streaming connection for session events
func (c *MessageRouterClient) createSessionEventStream(
	ctx context.Context,
	req *router.SubscribeSessionEventsRequest,
) (grpc.ClientStream, error) {
	streamDesc := &grpc.StreamDesc{
		StreamName:    "SubscribeSessionEvents",
		ServerStreams: true,
	}

	stream, err := c.getConn().NewStream(ctx, streamDesc, "/mpc.router.v1.MessageRouter/SubscribeSessionEvents")
	if err != nil {
		return nil, err
	}

	if err := stream.SendMsg(req); err != nil {
		return nil, err
	}

	if err := stream.CloseSend(); err != nil {
		return nil, err
	}

	return stream, nil
}

// AcknowledgeMessage acknowledges receipt and processing of a message
// This should be called after successfully processing an MPC message
// Includes automatic retry with exponential backoff
func (c *MessageRouterClient) AcknowledgeMessage(
	ctx context.Context,
	messageID string,
	sessionID uuid.UUID,
	partyID string,
	success bool,
	errorMessage string,
) error {
	req := &router.AcknowledgeMessageRequest{
		MessageId:    messageID,
		PartyId:      partyID,
		SessionId:    sessionID.String(),
		Success:      success,
		ErrorMessage: errorMessage,
	}

	return retry.DoVoid(ctx, c.retryCfg, "AcknowledgeMessage", func() error {
		resp := &router.AcknowledgeMessageResponse{}
		err := c.getConn().Invoke(ctx, "/mpc.router.v1.MessageRouter/AcknowledgeMessage", req, resp)
		if err != nil {
			return err
		}

		if !resp.Success {
			logger.Warn("Message acknowledgment failed",
				zap.String("message_id", messageID),
				zap.String("response", resp.Message))
		}

		return nil
	})
}

// GetMessageStatus gets the delivery status of a message
// Includes automatic retry with exponential backoff
func (c *MessageRouterClient) GetMessageStatus(
	ctx context.Context,
	messageID string,
	sessionID uuid.UUID,
) (*router.GetMessageStatusResponse, error) {
	req := &router.GetMessageStatusRequest{
		MessageId: messageID,
		SessionId: sessionID.String(),
	}

	return retry.Do(ctx, c.retryCfg, "GetMessageStatus", func() (*router.GetMessageStatusResponse, error) {
		resp := &router.GetMessageStatusResponse{}
		err := c.getConn().Invoke(ctx, "/mpc.router.v1.MessageRouter/GetMessageStatus", req, resp)
		if err != nil {
			return nil, err
		}
		return resp, nil
	})
}

// Heartbeat sends a heartbeat to keep the party connection alive
// Returns the number of pending messages for this party
// Includes automatic retry with exponential backoff
func (c *MessageRouterClient) Heartbeat(
	ctx context.Context,
	partyID string,
) (int32, error) {
	req := &router.HeartbeatRequest{
		PartyId:   partyID,
		Timestamp: time.Now().UnixMilli(),
	}

	return retry.Do(ctx, c.retryCfg, "Heartbeat", func() (int32, error) {
		resp := &router.HeartbeatResponse{}
		err := c.getConn().Invoke(ctx, "/mpc.router.v1.MessageRouter/Heartbeat", req, resp)
		if err != nil {
			return 0, err
		}

		if resp.PendingMessages > 0 {
			logger.Debug("Heartbeat response",
				zap.String("party_id", partyID),
				zap.Int32("pending_messages", resp.PendingMessages))
		}

		return resp.PendingMessages, nil
	})
}

// StartHeartbeat starts a background goroutine that sends heartbeats periodically
// Returns a cancel function to stop the heartbeat
func (c *MessageRouterClient) StartHeartbeat(
	ctx context.Context,
	partyID string,
	interval time.Duration,
	onPendingMessages func(count int32),
) context.CancelFunc {
	heartbeatCtx, cancel := context.WithCancel(ctx)

	go func() {
		ticker := time.NewTicker(interval)
		defer ticker.Stop()

		for {
			select {
			case <-heartbeatCtx.Done():
				logger.Info("Heartbeat stopped", zap.String("party_id", partyID))
				return
			case <-ticker.C:
				pendingCount, err := c.Heartbeat(heartbeatCtx, partyID)
				if err != nil {
					logger.Warn("Heartbeat failed",
						zap.String("party_id", partyID),
						zap.Error(err))
					continue
				}

				if onPendingMessages != nil && pendingCount > 0 {
					onPendingMessages(pendingCount)
				}
			}
		}
	}()

	logger.Info("Heartbeat started",
		zap.String("party_id", partyID),
		zap.Duration("interval", interval))

	return cancel
}

// ============================================
// Session Operations (Proxied through Message Router)
// These methods allow server-parties to only connect to Message Router
// ============================================

// JoinSession joins an MPC session via Message Router
// Includes automatic retry with exponential backoff for transient failures
func (c *MessageRouterClient) JoinSession(
	ctx context.Context,
	sessionID uuid.UUID,
	partyID, joinToken string,
) (*use_cases.SessionInfo, error) {
	req := &router.JoinSessionRequest{
		SessionId: sessionID.String(),
		PartyId:   partyID,
		JoinToken: joinToken,
		DeviceInfo: &router.DeviceInfo{
			DeviceType: "server",
			DeviceId:   partyID,
			Platform:   "linux",
			AppVersion: "1.0.0",
		},
	}

	return retry.Do(ctx, c.retryCfg, "JoinSession", func() (*use_cases.SessionInfo, error) {
		resp := &router.JoinSessionResponse{}
		err := c.getConn().Invoke(ctx, "/mpc.router.v1.MessageRouter/JoinSession", req, resp)
		if err != nil {
			return nil, err
		}

		if !resp.Success {
			logger.Error("Join session failed", zap.String("session_id", sessionID.String()))
			return nil, use_cases.ErrInvalidSession
		}

		// Convert response to SessionInfo
		// Include both other parties AND self in the participants list
		participants := make([]use_cases.ParticipantInfo, 0, len(resp.OtherParties)+1)

		// Add self first
		participants = append(participants, use_cases.ParticipantInfo{
			PartyID:    partyID,
			PartyIndex: int(resp.PartyIndex),
		})
		logger.Debug("Added self to participants",
			zap.String("party_id", partyID),
			zap.Int32("party_index", resp.PartyIndex))

		// Add other parties
		for _, p := range resp.OtherParties {
			logger.Debug("Received party_index from Message Router",
				zap.String("party_id", p.PartyId),
				zap.Int32("party_index", p.PartyIndex))

			participants = append(participants, use_cases.ParticipantInfo{
				PartyID:    p.PartyId,
				PartyIndex: int(p.PartyIndex),
			})
		}

		sessionInfo := &use_cases.SessionInfo{
			SessionID:    sessionID,
			Participants: participants,
		}

		if resp.SessionInfo != nil {
			sessionInfo.SessionType = resp.SessionInfo.SessionType
			sessionInfo.ThresholdN = int(resp.SessionInfo.ThresholdN)
			sessionInfo.ThresholdT = int(resp.SessionInfo.ThresholdT)
			sessionInfo.MessageHash = resp.SessionInfo.MessageHash

			// Parse keygen_session_id if provided (for signing sessions)
			if resp.SessionInfo.KeygenSessionId != "" {
				if keygenID, err := uuid.Parse(resp.SessionInfo.KeygenSessionId); err == nil {
					sessionInfo.KeygenSessionID = keygenID
				}
			}
		}

		logger.Info("Joined session via Message Router",
			zap.String("session_id", sessionID.String()),
			zap.String("party_id", partyID),
			zap.String("session_type", sessionInfo.SessionType),
			zap.String("keygen_session_id", sessionInfo.KeygenSessionID.String()),
			zap.Int("threshold_n", sessionInfo.ThresholdN),
			zap.Int("threshold_t", sessionInfo.ThresholdT))

		return sessionInfo, nil
	})
}

// ReportCompletion reports protocol completion via Message Router
// Includes automatic retry with exponential backoff for transient failures
func (c *MessageRouterClient) ReportCompletion(
	ctx context.Context,
	sessionID uuid.UUID,
	partyID string,
	resultData []byte,
) error {
	req := &router.ReportCompletionRequest{
		SessionId: sessionID.String(),
		PartyId:   partyID,
		PublicKey: resultData, // For keygen: public key; for signing: signature
	}

	return retry.DoVoid(ctx, c.retryCfg, "ReportCompletion", func() error {
		resp := &router.ReportCompletionResponse{}
		err := c.getConn().Invoke(ctx, "/mpc.router.v1.MessageRouter/ReportCompletion", req, resp)
		if err != nil {
			return err
		}

		logger.Info("Reported completion via Message Router",
			zap.String("session_id", sessionID.String()),
			zap.String("party_id", partyID),
			zap.Bool("all_completed", resp.AllCompleted))

		return nil
	})
}

// MarkPartyReady marks the party as ready to start the protocol via Message Router
// Includes automatic retry with exponential backoff for transient failures
func (c *MessageRouterClient) MarkPartyReady(
	ctx context.Context,
	sessionID uuid.UUID,
	partyID string,
) (bool, error) {
	req := &router.MarkPartyReadyRequest{
		SessionId: sessionID.String(),
		PartyId:   partyID,
	}

	return retry.Do(ctx, c.retryCfg, "MarkPartyReady", func() (bool, error) {
		resp := &router.MarkPartyReadyResponse{}
		err := c.getConn().Invoke(ctx, "/mpc.router.v1.MessageRouter/MarkPartyReady", req, resp)
		if err != nil {
			return false, err
		}

		logger.Info("Marked party ready via Message Router",
			zap.String("session_id", sessionID.String()),
			zap.String("party_id", partyID),
			zap.Bool("all_ready", resp.AllReady))

		return resp.AllReady, nil
	})
}

// GetSessionStatus gets the current session status via Message Router
// Includes automatic retry with exponential backoff for transient failures
func (c *MessageRouterClient) GetSessionStatus(
	ctx context.Context,
	sessionID uuid.UUID,
) (string, error) {
	req := &router.GetSessionStatusRequest{
		SessionId: sessionID.String(),
	}

	return retry.Do(ctx, c.retryCfg, "GetSessionStatus", func() (string, error) {
		resp := &router.GetSessionStatusResponse{}
		err := c.getConn().Invoke(ctx, "/mpc.router.v1.MessageRouter/GetSessionStatus", req, resp)
		if err != nil {
			return "", err
		}

		return resp.Status, nil
	})
}

// SubmitDelegateShare submits user's share from delegate party to Session Coordinator
// This is called after keygen by delegate parties to return the share for user retrieval
// Includes automatic retry with exponential backoff for transient failures
func (c *MessageRouterClient) SubmitDelegateShare(
	ctx context.Context,
	sessionID string,
	partyID string,
	encryptedShare []byte,
	publicKey []byte,
	partyIndex int32,
) error {
	req := &router.SubmitDelegateShareRequest{
		SessionId:      sessionID,
		PartyId:        partyID,
		EncryptedShare: encryptedShare,
		PublicKey:      publicKey,
		PartyIndex:     partyIndex,
	}

	return retry.DoVoid(ctx, c.retryCfg, "SubmitDelegateShare", func() error {
		resp := &router.SubmitDelegateShareResponse{}
		err := c.getConn().Invoke(ctx, "/mpc.router.v1.MessageRouter/SubmitDelegateShare", req, resp)
		if err != nil {
			return err
		}

		if !resp.Success {
			logger.Error("Submit delegate share failed",
				zap.String("session_id", sessionID),
				zap.String("party_id", partyID))
			return use_cases.ErrKeygenFailed
		}

		logger.Info("Delegate share submitted successfully via Message Router",
			zap.String("session_id", sessionID),
			zap.String("party_id", partyID),
			zap.Int("share_size", len(encryptedShare)))

		return nil
	})
}

// GetSessionStatusFull gets the full session status including participants via Message Router
// This is used for co_managed_keygen sessions to wait for all parties to join
// Includes automatic retry with exponential backoff for transient failures
func (c *MessageRouterClient) GetSessionStatusFull(
	ctx context.Context,
	sessionID uuid.UUID,
) (*use_cases.SessionStatusInfo, error) {
	req := &router.GetSessionStatusRequest{
		SessionId: sessionID.String(),
	}

	return retry.Do(ctx, c.retryCfg, "GetSessionStatusFull", func() (*use_cases.SessionStatusInfo, error) {
		resp := &router.GetSessionStatusResponse{}
		err := c.getConn().Invoke(ctx, "/mpc.router.v1.MessageRouter/GetSessionStatus", req, resp)
		if err != nil {
			return nil, err
		}

		// Convert participants from response
		participants := make([]use_cases.ParticipantInfo, len(resp.Participants))
		for i, p := range resp.Participants {
			participants[i] = use_cases.ParticipantInfo{
				PartyID:    p.PartyId,
				PartyIndex: int(p.PartyIndex),
			}
		}

		logger.Debug("GetSessionStatusFull response",
			zap.String("session_id", sessionID.String()),
			zap.String("status", resp.Status),
			zap.Int32("threshold_n", resp.ThresholdN),
			zap.Int("participants_count", len(participants)))

		return &use_cases.SessionStatusInfo{
			Status:       resp.Status,
			ThresholdN:   int(resp.ThresholdN),
			ThresholdT:   int(resp.ThresholdT),
			Participants: participants,
		}, nil
	})
}
