package use_cases

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/google/uuid"
	"github.com/rwadurian/mpc-system/pkg/jwt"
	"github.com/rwadurian/mpc-system/pkg/logger"
	"github.com/rwadurian/mpc-system/services/session-coordinator/application/ports/input"
	"github.com/rwadurian/mpc-system/services/session-coordinator/application/ports/output"
	"github.com/rwadurian/mpc-system/services/session-coordinator/domain/entities"
	"github.com/rwadurian/mpc-system/services/session-coordinator/domain/repositories"
	"github.com/rwadurian/mpc-system/services/session-coordinator/domain/services"
	"github.com/rwadurian/mpc-system/services/session-coordinator/domain/value_objects"
	"go.uber.org/zap"
)

// Maximum retries for optimistic lock conflicts during join session
const joinSessionMaxRetries = 3

// SessionParticipantInfo contains party ID and index for session_started event
type SessionParticipantInfo struct {
	PartyID    string
	PartyIndex int32
}

// JoinSessionMessageRouterClient defines the interface for publishing session events via gRPC
type JoinSessionMessageRouterClient interface {
	// PublishSessionStarted publishes session_started event with complete participants list
	// CRITICAL: participants contains all parties with their indices for TSS protocol
	PublishSessionStarted(
		ctx context.Context,
		sessionID string,
		thresholdN int32,
		thresholdT int32,
		selectedParties []string,
		joinTokens map[string]string,
		startedAt int64,
		participants []SessionParticipantInfo,
	) error

	// PublishParticipantJoined broadcasts a participant_joined event to all parties in the session
	// This allows the initiator's waiting screen to update in real-time when participants join
	PublishParticipantJoined(
		ctx context.Context,
		sessionID string,
		partyID string,
		selectedParties []string,
		joinedAt int64,
	) error
}

// JoinSessionUseCase implements the join session use case
type JoinSessionUseCase struct {
	sessionRepo         repositories.SessionRepository
	tokenValidator      jwt.TokenValidator
	eventPublisher      output.MessageBrokerPort
	messageRouterClient JoinSessionMessageRouterClient
	coordinatorSvc      *services.SessionCoordinatorService
}

// NewJoinSessionUseCase creates a new join session use case
func NewJoinSessionUseCase(
	sessionRepo repositories.SessionRepository,
	tokenValidator jwt.TokenValidator,
	eventPublisher output.MessageBrokerPort,
	messageRouterClient JoinSessionMessageRouterClient,
) *JoinSessionUseCase {
	return &JoinSessionUseCase{
		sessionRepo:         sessionRepo,
		tokenValidator:      tokenValidator,
		eventPublisher:      eventPublisher,
		messageRouterClient: messageRouterClient,
		coordinatorSvc:      services.NewSessionCoordinatorService(),
	}
}

// Execute executes the join session use case with retry logic for optimistic lock conflicts
func (uc *JoinSessionUseCase) Execute(
	ctx context.Context,
	inputData input.JoinSessionInput,
) (*input.JoinSessionOutput, error) {
	return uc.executeWithRetry(ctx, inputData, 0)
}

// executeWithRetry executes the join session with retry logic for optimistic lock conflicts
func (uc *JoinSessionUseCase) executeWithRetry(
	ctx context.Context,
	inputData input.JoinSessionInput,
	retry int,
) (*input.JoinSessionOutput, error) {
	if retry >= joinSessionMaxRetries {
		logger.Error("max retries exceeded for optimistic lock in join session",
			zap.String("session_id", inputData.SessionID.String()),
			zap.String("party_id", inputData.PartyID),
			zap.Int("retry_count", retry))
		return nil, fmt.Errorf("max retries exceeded: %w", entities.ErrOptimisticLockConflict)
	}

	if retry > 0 {
		logger.Info("retrying join session due to optimistic lock conflict",
			zap.String("session_id", inputData.SessionID.String()),
			zap.String("party_id", inputData.PartyID),
			zap.Int("retry_attempt", retry))
	}

	// Debug: log token info
	tokenLen := len(inputData.JoinToken)
	tokenPreview := ""
	if tokenLen > 20 {
		tokenPreview = inputData.JoinToken[:20] + "..."
	} else if tokenLen > 0 {
		tokenPreview = inputData.JoinToken
	}
	logger.Debug("JoinSession: parsing token",
		zap.String("session_id", inputData.SessionID.String()),
		zap.String("party_id", inputData.PartyID),
		zap.Int("token_length", tokenLen),
		zap.String("token_preview", tokenPreview))

	// 1. Parse join token to extract session ID (in case not provided)
	claims, err := uc.tokenValidator.ParseJoinTokenClaims(inputData.JoinToken)
	if err != nil {
		logger.Error("JoinSession: failed to parse token",
			zap.String("session_id", inputData.SessionID.String()),
			zap.String("party_id", inputData.PartyID),
			zap.Error(err))
		return nil, err
	}

	// Extract session ID from token if not provided in input
	sessionID := inputData.SessionID
	if sessionID == uuid.Nil {
		sessionID, err = uuid.Parse(claims.SessionID)
		if err != nil {
			return nil, err
		}
	}

	// 2. Validate join token with session ID and party ID
	_, err = uc.tokenValidator.ValidateJoinToken(
		inputData.JoinToken,
		sessionID,
		inputData.PartyID,
	)
	if err != nil {
		return nil, err
	}

	// 3. Load session (fresh read for each retry attempt)
	session, err := uc.sessionRepo.FindByUUID(ctx, sessionID)
	if err != nil {
		return nil, err
	}

	// 4. Create party ID value object
	partyID, err := value_objects.NewPartyID(inputData.PartyID)
	if err != nil {
		return nil, err
	}

	// 5. Get participant (must already exist from CreateSession)
	// For co_managed_keygen with wildcard token, allow dynamic participant addition
	participant, err := session.GetParticipant(partyID)
	if err != nil {
		// Check if this is a co_managed_keygen session with wildcard token
		// Wildcard tokens allow any party to join dynamically
		if session.SessionType == entities.SessionTypeCoManagedKeygen && claims.PartyID == "*" {
			// Dynamic participant addition for co_managed_keygen
			if len(session.Participants) >= session.Threshold.N() {
				logger.Warn("session is full, cannot add more participants",
					zap.String("session_id", session.ID.String()),
					zap.String("party_id", inputData.PartyID),
					zap.Int("current_participants", len(session.Participants)),
					zap.Int("threshold_n", session.Threshold.N()))
				return nil, entities.ErrSessionFull
			}

			// Create new participant with the next available index
			newIndex := len(session.Participants)
			participant, err = entities.NewParticipant(partyID, newIndex, inputData.DeviceInfo)
			if err != nil {
				logger.Error("failed to create new participant",
					zap.String("session_id", session.ID.String()),
					zap.String("party_id", inputData.PartyID),
					zap.Error(err))
				return nil, err
			}

			// Add participant to session
			if err := session.AddParticipant(participant); err != nil {
				logger.Error("failed to add participant to session",
					zap.String("session_id", session.ID.String()),
					zap.String("party_id", inputData.PartyID),
					zap.Error(err))
				return nil, err
			}

			logger.Info("dynamically added participant to co_managed_keygen session",
				zap.String("session_id", session.ID.String()),
				zap.String("party_id", inputData.PartyID),
				zap.Int("party_index", newIndex),
				zap.Int("total_participants", len(session.Participants)))
		} else {
			// Participant doesn't exist - this party was not invited to this session
			logger.Warn("party not found in session participants",
				zap.String("session_id", session.ID.String()),
				zap.String("party_id", inputData.PartyID),
				zap.Int("existing_participant_count", len(session.Participants)))
			return nil, entities.ErrPartyNotInvited
		}
	}

	logger.Debug("participant found in session",
		zap.String("party_id", participant.PartyID.String()),
		zap.Int("party_index", participant.PartyIndex),
		zap.String("current_status", participant.Status.String()))

	// 6. Update participant status to joined
	if err := session.UpdateParticipantStatus(partyID, value_objects.ParticipantStatusJoined); err != nil {
		return nil, err
	}

	// 6.1 Update party activity timestamp for timeout tracking
	if err := session.UpdatePartyActivity(partyID); err != nil {
		logger.Warn("failed to update party activity",
			zap.String("session_id", session.ID.String()),
			zap.String("party_id", inputData.PartyID),
			zap.Error(err))
	}

	// 7. Check if session should start (all participants joined)
	if uc.coordinatorSvc.ShouldStartSession(session) {
		if err := session.Start(); err != nil {
			return nil, err
		}

		startedAt := time.Now().UnixMilli()

		// Publish session started event to internal message broker (for logging/monitoring)
		startedEvent := output.SessionStartedEvent{
			SessionID: session.ID.String(),
			StartedAt: startedAt,
		}
		if err := uc.eventPublisher.PublishEvent(ctx, output.TopicSessionStarted, startedEvent); err != nil {
			logger.Error("failed to publish session started event",
				zap.String("session_id", session.ID.String()),
				zap.Error(err))
		}

		// Publish session started event via gRPC to message-router (for party notification)
		if uc.messageRouterClient != nil {
			selectedParties := session.GetPartyIDs()
			// Build join tokens map (empty for session_started, parties already have tokens)
			joinTokens := make(map[string]string)

			// CRITICAL: Build complete participants list with party indices
			// This ensures all parties have the same participant list for TSS protocol
			participants := make([]SessionParticipantInfo, len(session.Participants))
			for i, p := range session.Participants {
				participants[i] = SessionParticipantInfo{
					PartyID:    p.PartyID.String(),
					PartyIndex: int32(p.PartyIndex),
				}
			}

			if err := uc.messageRouterClient.PublishSessionStarted(
				ctx,
				session.ID.String(),
				int32(session.Threshold.N()),
				int32(session.Threshold.T()),
				selectedParties,
				joinTokens,
				startedAt,
				participants,
			); err != nil {
				logger.Error("failed to publish session started event to message router",
					zap.String("session_id", session.ID.String()),
					zap.Error(err))
			} else {
				logger.Info("published session started event to message router",
					zap.String("session_id", session.ID.String()),
					zap.Int("party_count", len(selectedParties)),
					zap.Int("participant_count", len(participants)))
			}
		}
	}

	// 8. Save updated session (with optimistic lock retry)
	if err := uc.sessionRepo.Update(ctx, session); err != nil {
		// Check if this is an optimistic lock conflict - if so, retry
		if errors.Is(err, entities.ErrOptimisticLockConflict) {
			logger.Warn("optimistic lock conflict detected in join session, retrying",
				zap.String("session_id", session.ID.String()),
				zap.String("party_id", inputData.PartyID),
				zap.Int("retry_attempt", retry+1))
			return uc.executeWithRetry(ctx, inputData, retry+1)
		}
		return nil, err
	}

	// 9. Publish participant joined event to internal message broker
	joinedAt := time.Now().UnixMilli()
	event := output.ParticipantJoinedEvent{
		SessionID: session.ID.String(),
		PartyID:   inputData.PartyID,
		JoinedAt:  joinedAt,
	}
	if err := uc.eventPublisher.PublishEvent(ctx, output.TopicParticipantJoined, event); err != nil {
		logger.Error("failed to publish participant joined event to internal broker",
			zap.String("session_id", session.ID.String()),
			zap.String("party_id", inputData.PartyID),
			zap.Error(err))
	}

	// 9.1 Publish participant joined event via gRPC to message-router (for real-time UI updates)
	// This notifies the initiator and other participants that a new party has joined
	if uc.messageRouterClient != nil {
		// Get all party IDs in the session to notify them
		allPartyIDs := session.GetPartyIDs()
		logger.Info("Broadcasting participant_joined event via gRPC",
			zap.String("session_id", session.ID.String()),
			zap.String("joined_party_id", inputData.PartyID),
			zap.Strings("notify_parties", allPartyIDs),
			zap.Int("total_participants", len(session.Participants)))

		if err := uc.messageRouterClient.PublishParticipantJoined(
			ctx,
			session.ID.String(),
			inputData.PartyID,
			allPartyIDs,
			joinedAt,
		); err != nil {
			logger.Error("failed to publish participant joined event to message router",
				zap.String("session_id", session.ID.String()),
				zap.String("party_id", inputData.PartyID),
				zap.Error(err))
		} else {
			logger.Info("Successfully published participant_joined event to message router",
				zap.String("session_id", session.ID.String()),
				zap.String("joined_party_id", inputData.PartyID),
				zap.Int("notify_count", len(allPartyIDs)))
		}
	} else {
		logger.Warn("messageRouterClient is nil, cannot broadcast participant_joined event",
			zap.String("session_id", session.ID.String()),
			zap.String("party_id", inputData.PartyID))
	}

	// 10. Build response with other parties info
	otherParties := session.GetOtherParties(partyID)
	partyInfos := make([]input.PartyInfo, len(otherParties))
	for i, p := range otherParties {
		partyInfos[i] = input.PartyInfo{
			PartyID:    p.PartyID.String(),
			PartyIndex: p.PartyIndex,
			DeviceInfo: p.DeviceInfo,
		}
	}

	// Debug logging
	logger.Info("join session - returning participant info",
		zap.String("party_id", inputData.PartyID),
		zap.Int("party_index", participant.PartyIndex),
		zap.Int("total_participants", len(session.Participants)))

	return &input.JoinSessionOutput{
		Success:    true,
		PartyIndex: participant.PartyIndex,
		SessionInfo: input.SessionInfo{
			SessionID:       session.ID.UUID(),
			SessionType:     string(session.SessionType),
			ThresholdN:      session.Threshold.N(),
			ThresholdT:      session.Threshold.T(),
			MessageHash:     session.MessageHash,
			Status:          session.Status.String(),
			KeygenSessionID: session.KeygenSessionID,
		},
		OtherParties: partyInfos,
	}, nil
}
