package grpc

import (
	"context"
	"fmt"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"

	coordinatorpb "github.com/rwadurian/mpc-system/api/grpc/coordinator/v1"
	"github.com/rwadurian/mpc-system/pkg/logger"
	"go.uber.org/zap"
)

// SessionCoordinatorClient wraps the gRPC client for session coordinator
type SessionCoordinatorClient struct {
	client coordinatorpb.SessionCoordinatorClient
	conn   *grpc.ClientConn
}

// NewSessionCoordinatorClient creates a new session coordinator gRPC client
func NewSessionCoordinatorClient(address string) (*SessionCoordinatorClient, error) {
	var conn *grpc.ClientConn
	var err error

	maxRetries := 5
	for i := 0; i < maxRetries; i++ {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		conn, err = grpc.DialContext(
			ctx,
			address,
			grpc.WithTransportCredentials(insecure.NewCredentials()),
			grpc.WithBlock(),
		)
		cancel()

		if err == nil {
			break
		}

		if i < maxRetries-1 {
			logger.Warn("Failed to connect to session coordinator, retrying...",
				zap.Int("attempt", i+1),
				zap.Int("max_retries", maxRetries),
				zap.Error(err))
			time.Sleep(time.Duration(i+1) * 2 * time.Second)
		}
	}

	if err != nil {
		return nil, fmt.Errorf("failed to connect to session coordinator after %d retries: %w", maxRetries, err)
	}

	logger.Info("Connected to session coordinator", zap.String("address", address))

	client := coordinatorpb.NewSessionCoordinatorClient(conn)

	return &SessionCoordinatorClient{
		client: client,
		conn:   conn,
	}, nil
}

// CreateKeygenSessionAuto creates a new keygen session with automatic party selection
// Coordinator will select parties from registered pool based on composition requirements
func (c *SessionCoordinatorClient) CreateKeygenSessionAuto(
	ctx context.Context,
	thresholdN int32,
	thresholdT int32,
	persistentCount int32,
	delegateCount int32,
	expiresInSeconds int64,
) (*CreateSessionAutoResponse, error) {
	req := &coordinatorpb.CreateSessionRequest{
		SessionType:      "keygen",
		ThresholdN:       thresholdN,
		ThresholdT:       thresholdT,
		Participants:     nil, // No participants - coordinator selects automatically
		ExpiresInSeconds: expiresInSeconds,
		PartyComposition: &coordinatorpb.PartyComposition{
			PersistentCount: persistentCount,
			DelegateCount:   delegateCount,
		},
	}

	logger.Info("Sending CreateSession gRPC request (auto party selection)",
		zap.String("session_type", "keygen"),
		zap.Int32("threshold_n", thresholdN),
		zap.Int32("threshold_t", thresholdT),
		zap.Int32("persistent_count", persistentCount),
		zap.Int32("delegate_count", delegateCount))

	resp, err := c.client.CreateSession(ctx, req)
	if err != nil {
		logger.Error("CreateSession gRPC call failed", zap.Error(err))
		return nil, fmt.Errorf("failed to create keygen session: %w", err)
	}

	// Extract selected parties and delegate party from response
	var selectedParties []string
	var delegateParty string
	for partyID := range resp.JoinTokens {
		selectedParties = append(selectedParties, partyID)
	}
	if resp.DelegatePartyId != "" {
		delegateParty = resp.DelegatePartyId
	}

	logger.Info("CreateSession gRPC call succeeded",
		zap.String("session_id", resp.SessionId),
		zap.Int("num_parties", len(selectedParties)),
		zap.String("delegate_party", delegateParty))

	return &CreateSessionAutoResponse{
		SessionID:       resp.SessionId,
		SelectedParties: selectedParties,
		DelegateParty:   delegateParty,
		JoinTokens:      resp.JoinTokens,
		ExpiresAt:       resp.ExpiresAt,
	}, nil
}

// DelegateUserShareInput contains user's share for delegate party in signing
type DelegateUserShareInput struct {
	DelegatePartyID string
	EncryptedShare  []byte
	PartyIndex      int32
}

// SigningPartyInfo contains party ID and its original index from keygen
type SigningPartyInfo struct {
	PartyID    string
	PartyIndex int32
}

// CreateSigningSessionAuto creates a new signing session with automatic party selection
// Coordinator will select parties from the provided party info (from account shares)
// delegateUserShare is required if any of the parties is a delegate party
// keygenThresholdN is the original threshold_n from the keygen session (required for TSS math)
//
// BREAKING CHANGE WARNING (for co-sign feature, commit 042212ea):
// Original code: ThresholdN = int32(len(parties)) - used participant count as N
// New code: ThresholdN = keygenThresholdN - uses original N from keygen session
// This change affects PERSISTENT SIGN flow. The original approach made threshold_n
// equal to participant count (T+1), which worked with the old N-based validation.
// If issues arise with persistent sign, REVERT to: ThresholdN: int32(len(parties))
// Related files: session_coordinator.go, mpc_session.go, account_handler.go
func (c *SessionCoordinatorClient) CreateSigningSessionAuto(
	ctx context.Context,
	thresholdT int32,
	keygenThresholdN int32,
	parties []SigningPartyInfo,
	messageHash []byte,
	expiresInSeconds int64,
	delegateUserShare *DelegateUserShareInput,
	keygenSessionID string,
) (*CreateSessionAutoResponse, error) {
	// Convert party info to participant info with original party index from keygen
	pbParticipants := make([]*coordinatorpb.ParticipantInfo, len(parties))
	for i, party := range parties {
		pbParticipants[i] = &coordinatorpb.ParticipantInfo{
			PartyId:    party.PartyID,
			PartyIndex: party.PartyIndex,
		}
	}

	// CRITICAL: Use keygenThresholdN (original n from keygen), NOT len(parties)
	// TSS signing requires the same n value used during keygen for correct mathematical operations
	req := &coordinatorpb.CreateSessionRequest{
		SessionType:      "sign",
		ThresholdN:       keygenThresholdN,
		ThresholdT:       thresholdT,
		Participants:     pbParticipants,
		MessageHash:      messageHash,
		ExpiresInSeconds: expiresInSeconds,
		KeygenSessionId:  keygenSessionID,
	}

	// Add delegate user share if provided
	if delegateUserShare != nil {
		req.DelegateUserShare = &coordinatorpb.DelegateUserShare{
			DelegatePartyId: delegateUserShare.DelegatePartyID,
			EncryptedShare:  delegateUserShare.EncryptedShare,
			PartyIndex:      delegateUserShare.PartyIndex,
		}
		logger.Info("Sending CreateSigningSession gRPC request with delegate user share",
			zap.Int32("threshold_t", thresholdT),
			zap.Int32("keygen_threshold_n", keygenThresholdN),
			zap.Int("num_signing_parties", len(parties)),
			zap.String("delegate_party_id", delegateUserShare.DelegatePartyID))
	} else {
		logger.Info("Sending CreateSigningSession gRPC request",
			zap.Int32("threshold_t", thresholdT),
			zap.Int32("keygen_threshold_n", keygenThresholdN),
			zap.Int("num_signing_parties", len(parties)))
	}

	resp, err := c.client.CreateSession(ctx, req)
	if err != nil {
		logger.Error("CreateSigningSession gRPC call failed", zap.Error(err))
		return nil, fmt.Errorf("failed to create signing session: %w", err)
	}

	// Extract selected parties from response
	var selectedParties []string
	for partyID := range resp.JoinTokens {
		selectedParties = append(selectedParties, partyID)
	}

	logger.Info("CreateSigningSession gRPC call succeeded",
		zap.String("session_id", resp.SessionId),
		zap.Int("num_parties", len(selectedParties)))

	return &CreateSessionAutoResponse{
		SessionID:       resp.SessionId,
		SelectedParties: selectedParties,
		DelegateParty:   resp.DelegatePartyId,
		JoinTokens:      resp.JoinTokens,
		ExpiresAt:       resp.ExpiresAt,
	}, nil
}

// GetSessionStatus gets the status of a session
func (c *SessionCoordinatorClient) GetSessionStatus(
	ctx context.Context,
	sessionID string,
) (*SessionStatusResponse, error) {
	req := &coordinatorpb.GetSessionStatusRequest{
		SessionId: sessionID,
	}

	resp, err := c.client.GetSessionStatus(ctx, req)
	if err != nil {
		return nil, fmt.Errorf("failed to get session status: %w", err)
	}

	result := &SessionStatusResponse{
		Status:           resp.Status,
		CompletedParties: resp.CompletedParties,
		TotalParties:     resp.TotalParties,
		ThresholdT:       resp.ThresholdT,
		ThresholdN:       resp.ThresholdN,
		SessionType:      resp.SessionType,
		PublicKey:        resp.PublicKey,
		Signature:        resp.Signature,
		HasDelegate:      resp.HasDelegate,
	}

	// Include delegate share if present
	if resp.DelegateShare != nil {
		result.DelegateShare = &DelegateShareInfo{
			EncryptedShare: resp.DelegateShare.EncryptedShare,
			PartyIndex:     resp.DelegateShare.PartyIndex,
			PartyID:        resp.DelegateShare.PartyId,
		}
	}

	// Include participants if present (for co_managed_keygen sessions)
	if len(resp.Participants) > 0 {
		result.Participants = make([]ParticipantStatusInfo, len(resp.Participants))
		for i, p := range resp.Participants {
			result.Participants[i] = ParticipantStatusInfo{
				PartyID:    p.PartyId,
				PartyIndex: p.PartyIndex,
				Status:     p.Status,
			}
		}
	}

	return result, nil
}

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

// ParticipantInfo contains participant information
type ParticipantInfo struct {
	PartyID    string
	DeviceType string
	DeviceID   string
	Platform   string
	AppVersion string
}

// CreateSessionResponse contains the created session information
type CreateSessionResponse struct {
	SessionID  string
	JoinTokens map[string]string
	ExpiresAt  int64
}

// CreateSessionAutoResponse contains the created session info with auto-selected parties
type CreateSessionAutoResponse struct {
	SessionID       string
	SelectedParties []string
	DelegateParty   string
	JoinTokens      map[string]string
	ExpiresAt       int64
}

// SessionStatusResponse contains session status information
type SessionStatusResponse struct {
	Status           string
	CompletedParties int32
	TotalParties     int32
	ThresholdT       int32  // Minimum parties needed to sign (e.g., 2 in 2-of-3)
	ThresholdN       int32  // Total number of parties required (e.g., 3 in 2-of-3)
	SessionType      string // "keygen" or "sign"
	PublicKey        []byte
	Signature        []byte
	// HasDelegate indicates whether this keygen session has a delegate party
	// Only meaningful for keygen sessions. Always false for sign sessions.
	HasDelegate bool
	// DelegateShare is non-nil when session_type="keygen" AND has_delegate=true AND session is completed
	// nil with has_delegate=true means share was already retrieved (one-time retrieval)
	DelegateShare *DelegateShareInfo
	// Participants contains detailed participant information including party_index
	// Used by service-party-app for co_managed_keygen sessions
	Participants []ParticipantStatusInfo
}

// ParticipantStatusInfo contains participant status information
type ParticipantStatusInfo struct {
	PartyID    string
	PartyIndex int32
	Status     string
}

// DelegateShareInfo contains the delegate party's share for user
type DelegateShareInfo struct {
	EncryptedShare []byte
	PartyIndex     int32
	PartyID        string
}

// CreateCoManagedSessionResponse contains the created co-managed session info
type CreateCoManagedSessionResponse struct {
	SessionID            string
	InviteCode           string
	WalletName           string
	SelectedServerParties []string          // Auto-selected server parties
	JoinTokens           map[string]string  // Includes wildcard token for external parties
	ExpiresAt            int64
	ThresholdN           int32
	ThresholdT           int32
}

// CreateCoManagedKeygenSession creates a new co-managed keygen session
// This session waits for external participants to join via invite code
func (c *SessionCoordinatorClient) CreateCoManagedKeygenSession(
	ctx context.Context,
	walletName string,
	inviteCode string,
	thresholdN int32,
	thresholdT int32,
	persistentCount int32,   // Number of server parties to auto-select
	externalCount int32,     // Number of external parties (Service Party Apps)
	expiresInSeconds int64,
) (*CreateCoManagedSessionResponse, error) {
	req := &coordinatorpb.CreateSessionRequest{
		SessionType:      "co_managed_keygen",
		ThresholdN:       thresholdN,
		ThresholdT:       thresholdT,
		Participants:     nil, // External participants will join later
		ExpiresInSeconds: expiresInSeconds,
		PartyComposition: &coordinatorpb.PartyComposition{
			PersistentCount: persistentCount,
			DelegateCount:   0,
			TemporaryCount:  externalCount, // External parties treated as temporary
		},
		WalletName: walletName,
		InviteCode: inviteCode,
	}

	logger.Info("Sending CreateCoManagedKeygenSession gRPC request",
		zap.String("session_type", "co_managed_keygen"),
		zap.String("wallet_name", walletName),
		zap.Int32("threshold_n", thresholdN),
		zap.Int32("threshold_t", thresholdT),
		zap.Int32("persistent_count", persistentCount),
		zap.Int32("external_count", externalCount))

	resp, err := c.client.CreateSession(ctx, req)
	if err != nil {
		logger.Error("CreateCoManagedKeygenSession gRPC call failed", zap.Error(err))
		return nil, fmt.Errorf("failed to create co-managed keygen session: %w", err)
	}

	// Extract selected server parties
	var serverParties []string
	for partyID := range resp.JoinTokens {
		if partyID != "*" { // Exclude wildcard token
			serverParties = append(serverParties, partyID)
		}
	}

	logger.Info("CreateCoManagedKeygenSession gRPC call succeeded",
		zap.String("session_id", resp.SessionId),
		zap.Int("num_server_parties", len(serverParties)))

	return &CreateCoManagedSessionResponse{
		SessionID:            resp.SessionId,
		InviteCode:           inviteCode,
		WalletName:           walletName,
		SelectedServerParties: serverParties,
		JoinTokens:           resp.JoinTokens,
		ExpiresAt:            resp.ExpiresAt,
		ThresholdN:           thresholdN,
		ThresholdT:           thresholdT,
	}, nil
}

// JoinSessionResponse contains join session result
type JoinSessionResponse struct {
	Success      bool
	PartyIndex   int32
	SessionInfo  *SessionInfoResponse
	OtherParties []PartyInfoResponse
}

// SessionInfoResponse contains session info from join response
type SessionInfoResponse struct {
	SessionID      string
	SessionType    string
	ThresholdN     int32
	ThresholdT     int32
	Status         string
	WalletName     string
	InviteCode     string
	KeygenSessionID string
}

// PartyInfoResponse contains party info
type PartyInfoResponse struct {
	PartyID    string
	PartyIndex int32
}

// JoinSession joins an existing session
func (c *SessionCoordinatorClient) JoinSession(
	ctx context.Context,
	sessionID string,
	partyID string,
	joinToken string,
	deviceType string,
	deviceID string,
) (*JoinSessionResponse, error) {
	req := &coordinatorpb.JoinSessionRequest{
		SessionId: sessionID,
		PartyId:   partyID,
		JoinToken: joinToken,
		DeviceInfo: &coordinatorpb.DeviceInfo{
			DeviceType: deviceType,
			DeviceId:   deviceID,
		},
	}

	logger.Info("Sending JoinSession gRPC request",
		zap.String("session_id", sessionID),
		zap.String("party_id", partyID))

	resp, err := c.client.JoinSession(ctx, req)
	if err != nil {
		logger.Error("JoinSession gRPC call failed", zap.Error(err))
		return nil, fmt.Errorf("failed to join session: %w", err)
	}

	result := &JoinSessionResponse{
		Success:    resp.Success,
		PartyIndex: resp.PartyIndex,
	}

	if resp.SessionInfo != nil {
		result.SessionInfo = &SessionInfoResponse{
			SessionID:      resp.SessionInfo.SessionId,
			SessionType:    resp.SessionInfo.SessionType,
			ThresholdN:     resp.SessionInfo.ThresholdN,
			ThresholdT:     resp.SessionInfo.ThresholdT,
			Status:         resp.SessionInfo.Status,
			WalletName:     resp.SessionInfo.WalletName,
			InviteCode:     resp.SessionInfo.InviteCode,
			KeygenSessionID: resp.SessionInfo.KeygenSessionId,
		}
	}

	for _, p := range resp.OtherParties {
		result.OtherParties = append(result.OtherParties, PartyInfoResponse{
			PartyID:    p.PartyId,
			PartyIndex: p.PartyIndex,
		})
	}

	logger.Info("JoinSession gRPC call succeeded",
		zap.Bool("success", resp.Success),
		zap.Int32("party_index", resp.PartyIndex))

	return result, nil
}
