package grpc

import (
	"context"
	"sync"
	"time"

	"github.com/google/uuid"
	pb "github.com/rwadurian/mpc-system/api/grpc/coordinator/v1"
	"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/use_cases"
	"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/value_objects"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

// DelegateShareEntry represents a delegate party's share for user retrieval
type DelegateShareEntry struct {
	PartyID        string
	EncryptedShare []byte
	PublicKey      []byte
	PartyIndex     int32
	CreatedAt      time.Time
}

// DelegateShareCache stores delegate shares temporarily for user retrieval
type DelegateShareCache struct {
	mu     sync.RWMutex
	shares map[string]*DelegateShareEntry // sessionID -> share
	ttl    time.Duration
}

// NewDelegateShareCache creates a new cache with TTL
func NewDelegateShareCache(ttl time.Duration) *DelegateShareCache {
	cache := &DelegateShareCache{
		shares: make(map[string]*DelegateShareEntry),
		ttl:    ttl,
	}
	// Start cleanup goroutine
	go cache.cleanupLoop()
	return cache
}

// Store stores a delegate share
func (c *DelegateShareCache) Store(sessionID string, entry *DelegateShareEntry) {
	c.mu.Lock()
	defer c.mu.Unlock()
	entry.CreatedAt = time.Now()
	c.shares[sessionID] = entry
	logger.Info("Delegate share stored",
		zap.String("session_id", sessionID),
		zap.String("party_id", entry.PartyID))
}

// Get retrieves and deletes a delegate share (one-time retrieval)
func (c *DelegateShareCache) Get(sessionID string) (*DelegateShareEntry, bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	entry, exists := c.shares[sessionID]
	if exists {
		delete(c.shares, sessionID)
		logger.Info("Delegate share retrieved and deleted",
			zap.String("session_id", sessionID),
			zap.String("party_id", entry.PartyID))
	}
	return entry, exists
}

// cleanupLoop removes expired entries
func (c *DelegateShareCache) cleanupLoop() {
	ticker := time.NewTicker(c.ttl / 2)
	defer ticker.Stop()
	for range ticker.C {
		c.mu.Lock()
		now := time.Now()
		for sessionID, entry := range c.shares {
			if now.Sub(entry.CreatedAt) > c.ttl {
				delete(c.shares, sessionID)
				logger.Debug("Expired delegate share removed",
					zap.String("session_id", sessionID))
			}
		}
		c.mu.Unlock()
	}
}

// Global delegate share cache (15 minute TTL)
var delegateShareCache = NewDelegateShareCache(15 * time.Minute)

// SessionCoordinatorServer implements the gRPC SessionCoordinator service
type SessionCoordinatorServer struct {
	pb.UnimplementedSessionCoordinatorServer
	createSessionUC    *use_cases.CreateSessionUseCase
	joinSessionUC      *use_cases.JoinSessionUseCase
	getSessionStatusUC *use_cases.GetSessionStatusUseCase
	reportCompletionUC *use_cases.ReportCompletionUseCase
	closeSessionUC     *use_cases.CloseSessionUseCase
	sessionRepo        repositories.SessionRepository
}

// NewSessionCoordinatorServer creates a new gRPC server
func NewSessionCoordinatorServer(
	createSessionUC *use_cases.CreateSessionUseCase,
	joinSessionUC *use_cases.JoinSessionUseCase,
	getSessionStatusUC *use_cases.GetSessionStatusUseCase,
	reportCompletionUC *use_cases.ReportCompletionUseCase,
	closeSessionUC *use_cases.CloseSessionUseCase,
	sessionRepo repositories.SessionRepository,
) *SessionCoordinatorServer {
	return &SessionCoordinatorServer{
		createSessionUC:    createSessionUC,
		joinSessionUC:      joinSessionUC,
		getSessionStatusUC: getSessionStatusUC,
		reportCompletionUC: reportCompletionUC,
		closeSessionUC:     closeSessionUC,
		sessionRepo:        sessionRepo,
	}
}

// CreateSession creates a new MPC session
func (s *SessionCoordinatorServer) CreateSession(
	ctx context.Context,
	req *pb.CreateSessionRequest,
) (*pb.CreateSessionResponse, error) {
	logger.Info("gRPC CreateSession received",
		zap.String("session_type", req.SessionType),
		zap.Int32("threshold_n", req.ThresholdN),
		zap.Int32("threshold_t", req.ThresholdT),
		zap.Int("num_participants", len(req.Participants)))

	// Convert request to input
	participants := make([]input.ParticipantInfo, len(req.Participants))
	for i, p := range req.Participants {
		var deviceInfo entities.DeviceInfo
		if p.DeviceInfo != nil {
			deviceInfo = entities.DeviceInfo{
				DeviceType: entities.DeviceType(p.DeviceInfo.DeviceType),
				DeviceID:   p.DeviceInfo.DeviceId,
				Platform:   p.DeviceInfo.Platform,
				AppVersion: p.DeviceInfo.AppVersion,
			}
		}
		participants[i] = input.ParticipantInfo{
			PartyID:    p.PartyId,
			DeviceInfo: deviceInfo,
			PartyIndex: int(p.PartyIndex),
		}
	}

	inputData := input.CreateSessionInput{
		InitiatorID:  "",
		SessionType:  req.SessionType,
		ThresholdN:   int(req.ThresholdN),
		ThresholdT:   int(req.ThresholdT),
		Participants: participants,
		MessageHash:  req.MessageHash,
		ExpiresIn:    time.Duration(req.ExpiresInSeconds) * time.Second,
		WalletName:   req.WalletName,
		InviteCode:   req.InviteCode,
	}

	// Parse and set keygen_session_id if provided (for signing sessions)
	if req.KeygenSessionId != "" {
		keygenSessionID, err := uuid.Parse(req.KeygenSessionId)
		if err == nil {
			inputData.KeygenSessionID = keygenSessionID
		} else {
			logger.Warn("Invalid keygen_session_id format",
				zap.String("keygen_session_id", req.KeygenSessionId),
				zap.Error(err))
		}
	}

	// Add party composition if provided (for automatic party selection)
	if req.PartyComposition != nil {
		inputData.PartyComposition = &input.PartyComposition{
			PersistentCount: int(req.PartyComposition.PersistentCount),
			DelegateCount:   int(req.PartyComposition.DelegateCount),
			TemporaryCount:  int(req.PartyComposition.TemporaryCount),
		}
		logger.Info("Party composition specified in gRPC request",
			zap.Int("persistent_count", inputData.PartyComposition.PersistentCount),
			zap.Int("delegate_count", inputData.PartyComposition.DelegateCount),
			zap.Int("temporary_count", inputData.PartyComposition.TemporaryCount))
	}

	// Add delegate user share if provided (for sign sessions with delegate party)
	if req.DelegateUserShare != nil {
		inputData.DelegateUserShare = &input.DelegateUserShare{
			DelegatePartyID: req.DelegateUserShare.DelegatePartyId,
			EncryptedShare:  req.DelegateUserShare.EncryptedShare,
			PartyIndex:      req.DelegateUserShare.PartyIndex,
		}
	}

	// Execute use case
	output, err := s.createSessionUC.Execute(ctx, inputData)
	if err != nil {
		logger.Error("gRPC CreateSession use case failed", zap.Error(err))
		return nil, toGRPCError(err)
	}

	logger.Info("gRPC CreateSession completed successfully",
		zap.String("session_id", output.SessionID.String()),
		zap.Int("num_join_tokens", len(output.JoinTokens)),
		zap.String("keygen_session_id", inputData.KeygenSessionID.String()),
		zap.String("session_type", req.SessionType))

	// Load session to get delegate party ID
	session, err := s.sessionRepo.FindByUUID(ctx, output.SessionID)
	if err != nil {
		logger.Error("Failed to load session after creation", zap.Error(err))
		return nil, toGRPCError(err)
	}

	// Extract selected party IDs from join tokens
	selectedParties := make([]string, 0, len(output.JoinTokens))
	for partyID := range output.JoinTokens {
		if partyID != "*" { // Exclude universal token
			selectedParties = append(selectedParties, partyID)
		}
	}

	// Convert output to response
	return &pb.CreateSessionResponse{
		SessionId:       output.SessionID.String(),
		JoinTokens:      output.JoinTokens,
		ExpiresAt:       output.ExpiresAt.UnixMilli(),
		SelectedParties: selectedParties,
		DelegatePartyId: session.DelegatePartyID,
	}, nil
}

// JoinSession allows a participant to join a session
func (s *SessionCoordinatorServer) JoinSession(
	ctx context.Context,
	req *pb.JoinSessionRequest,
) (*pb.JoinSessionResponse, error) {
	sessionID, err := uuid.Parse(req.SessionId)
	if err != nil {
		return nil, status.Error(codes.InvalidArgument, "invalid session ID")
	}

	var deviceInfo entities.DeviceInfo
	if req.DeviceInfo != nil {
		deviceInfo = entities.DeviceInfo{
			DeviceType: entities.DeviceType(req.DeviceInfo.DeviceType),
			DeviceID:   req.DeviceInfo.DeviceId,
			Platform:   req.DeviceInfo.Platform,
			AppVersion: req.DeviceInfo.AppVersion,
		}
	}

	inputData := input.JoinSessionInput{
		SessionID:  sessionID,
		PartyID:    req.PartyId,
		JoinToken:  req.JoinToken,
		DeviceInfo: deviceInfo,
	}

	output, err := s.joinSessionUC.Execute(ctx, inputData)
	if err != nil {
		return nil, toGRPCError(err)
	}

	// Convert other parties to response format
	otherParties := make([]*pb.PartyInfo, len(output.OtherParties))
	for i, p := range output.OtherParties {
		otherParties[i] = &pb.PartyInfo{
			PartyId:    p.PartyID,
			PartyIndex: int32(p.PartyIndex),
			DeviceInfo: &pb.DeviceInfo{
				DeviceType: string(p.DeviceInfo.DeviceType),
				DeviceId:   p.DeviceInfo.DeviceID,
				Platform:   p.DeviceInfo.Platform,
				AppVersion: p.DeviceInfo.AppVersion,
			},
		}
		// Debug: Log what we're about to send in gRPC response
		logger.Info("gRPC JoinSession - setting party_index in protobuf response",
			zap.String("party_id", p.PartyID),
			zap.Int("source_party_index", p.PartyIndex),
			zap.Int32("proto_party_index", otherParties[i].PartyIndex))
	}

	// Convert keygen_session_id to string (empty if nil)
	keygenSessionIDStr := ""
	if output.SessionInfo.KeygenSessionID != uuid.Nil {
		keygenSessionIDStr = output.SessionInfo.KeygenSessionID.String()
	}

	logger.Info("gRPC JoinSession completed",
		zap.String("session_id", output.SessionInfo.SessionID.String()),
		zap.String("party_id", req.PartyId),
		zap.Int("party_index", output.PartyIndex),
		zap.String("keygen_session_id", keygenSessionIDStr),
		zap.String("session_type", output.SessionInfo.SessionType),
		zap.Int("threshold_n", output.SessionInfo.ThresholdN),
		zap.Int("threshold_t", output.SessionInfo.ThresholdT))

	return &pb.JoinSessionResponse{
		Success: output.Success,
		SessionInfo: &pb.SessionInfo{
			SessionId:       output.SessionInfo.SessionID.String(),
			SessionType:     output.SessionInfo.SessionType,
			ThresholdN:      int32(output.SessionInfo.ThresholdN),
			ThresholdT:      int32(output.SessionInfo.ThresholdT),
			MessageHash:     output.SessionInfo.MessageHash,
			Status:          output.SessionInfo.Status,
			KeygenSessionId: keygenSessionIDStr,
		},
		OtherParties: otherParties,
		PartyIndex:   int32(output.PartyIndex),
	}, nil
}

// GetSessionStatus retrieves the status of a session
func (s *SessionCoordinatorServer) GetSessionStatus(
	ctx context.Context,
	req *pb.GetSessionStatusRequest,
) (*pb.GetSessionStatusResponse, error) {
	sessionID, err := uuid.Parse(req.SessionId)
	if err != nil {
		return nil, status.Error(codes.InvalidArgument, "invalid session ID")
	}

	output, err := s.getSessionStatusUC.Execute(ctx, sessionID)
	if err != nil {
		return nil, toGRPCError(err)
	}

	// Calculate completed parties from participants
	completedParties := 0
	for _, p := range output.Participants {
		if p.Status == "completed" {
			completedParties++
		}
	}

	// Convert participants to proto format
	protoParticipants := make([]*pb.ParticipantStatus, len(output.Participants))
	for i, p := range output.Participants {
		protoParticipants[i] = &pb.ParticipantStatus{
			PartyId:    p.PartyID,
			PartyIndex: int32(p.PartyIndex),
			Status:     p.Status,
		}
	}

	resp := &pb.GetSessionStatusResponse{
		Status:           output.Status,
		CompletedParties: int32(completedParties),
		TotalParties:     int32(len(output.Participants)),
		SessionType:      output.SessionType,
		PublicKey:        output.PublicKey,
		Signature:        output.Signature,
		HasDelegate:      output.HasDelegate, // Only meaningful for keygen sessions
		Participants:     protoParticipants,  // Include participant details with party_index
		ThresholdN:       int32(output.ThresholdN),
		ThresholdT:       int32(output.ThresholdT),
	}

	// Try to get delegate share from cache
	// Only include if session has delegate and is completed
	if output.HasDelegate && output.Status == "completed" {
		entry, found := delegateShareCache.Get(req.SessionId)
		if found {
			resp.DelegateShare = &pb.DelegateShareInfo{
				EncryptedShare: entry.EncryptedShare,
				PartyIndex:     entry.PartyIndex,
				PartyId:        entry.PartyID,
			}
			logger.Info("Delegate share included in GetSessionStatus response",
				zap.String("session_id", req.SessionId),
				zap.String("party_id", entry.PartyID))
		}
		// If has_delegate=true but delegate_share=nil, share was already retrieved
	}

	return resp, nil
}

// ReportCompletion reports that a participant has completed
func (s *SessionCoordinatorServer) ReportCompletion(
	ctx context.Context,
	req *pb.ReportCompletionRequest,
) (*pb.ReportCompletionResponse, error) {
	sessionID, err := uuid.Parse(req.SessionId)
	if err != nil {
		return nil, status.Error(codes.InvalidArgument, "invalid session ID")
	}

	inputData := input.ReportCompletionInput{
		SessionID: sessionID,
		PartyID:   req.PartyId,
		PublicKey: req.PublicKey,
		Signature: req.Signature,
	}

	output, err := s.reportCompletionUC.Execute(ctx, inputData)
	if err != nil {
		return nil, toGRPCError(err)
	}

	return &pb.ReportCompletionResponse{
		Success:      output.Success,
		AllCompleted: output.AllCompleted,
	}, nil
}

// CloseSession closes a session
func (s *SessionCoordinatorServer) CloseSession(
	ctx context.Context,
	req *pb.CloseSessionRequest,
) (*pb.CloseSessionResponse, error) {
	sessionID, err := uuid.Parse(req.SessionId)
	if err != nil {
		return nil, status.Error(codes.InvalidArgument, "invalid session ID")
	}

	err = s.closeSessionUC.Execute(ctx, sessionID)
	if err != nil {
		return nil, toGRPCError(err)
	}

	return &pb.CloseSessionResponse{
		Success: true,
	}, nil
}

// MarkPartyReady marks a party as ready
func (s *SessionCoordinatorServer) MarkPartyReady(
	ctx context.Context,
	req *pb.MarkPartyReadyRequest,
) (*pb.MarkPartyReadyResponse, error) {
	parsedID, err := uuid.Parse(req.SessionId)
	if err != nil {
		return nil, status.Error(codes.InvalidArgument, "invalid session ID")
	}
	sessionID := value_objects.SessionIDFromUUID(parsedID)

	session, err := s.sessionRepo.FindByID(ctx, sessionID)
	if err != nil {
		return nil, toGRPCError(err)
	}
	if session == nil {
		return nil, status.Error(codes.NotFound, "session not found")
	}

	// Mark party as ready
	if err := session.MarkPartyReady(req.PartyId); err != nil {
		return nil, toGRPCError(err)
	}

	// Save session
	if err := s.sessionRepo.Update(ctx, session); err != nil {
		return nil, toGRPCError(err)
	}

	// Check if all parties are ready
	allReady := session.AllPartiesReady()

	return &pb.MarkPartyReadyResponse{
		Success:  true,
		AllReady: allReady,
	}, nil
}

// StartSession starts a session
func (s *SessionCoordinatorServer) StartSession(
	ctx context.Context,
	req *pb.StartSessionRequest,
) (*pb.StartSessionResponse, error) {
	parsedID, err := uuid.Parse(req.SessionId)
	if err != nil {
		return nil, status.Error(codes.InvalidArgument, "invalid session ID")
	}
	sessionID := value_objects.SessionIDFromUUID(parsedID)

	session, err := s.sessionRepo.FindByID(ctx, sessionID)
	if err != nil {
		return nil, toGRPCError(err)
	}
	if session == nil {
		return nil, status.Error(codes.NotFound, "session not found")
	}

	// Start the session
	if err := session.Start(); err != nil {
		return nil, toGRPCError(err)
	}

	// Save session
	if err := s.sessionRepo.Update(ctx, session); err != nil {
		return nil, toGRPCError(err)
	}

	return &pb.StartSessionResponse{
		Success: true,
	}, nil
}

// SubmitDelegateShare stores a delegate party's share for user retrieval
func (s *SessionCoordinatorServer) SubmitDelegateShare(
	ctx context.Context,
	req *pb.SubmitDelegateShareRequest,
) (*pb.SubmitDelegateShareResponse, error) {
	if req.SessionId == "" {
		return nil, status.Error(codes.InvalidArgument, "session_id is required")
	}
	if req.PartyId == "" {
		return nil, status.Error(codes.InvalidArgument, "party_id is required")
	}
	if len(req.EncryptedShare) == 0 {
		return nil, status.Error(codes.InvalidArgument, "encrypted_share is required")
	}

	// Store in cache
	entry := &DelegateShareEntry{
		PartyID:        req.PartyId,
		EncryptedShare: req.EncryptedShare,
		PublicKey:      req.PublicKey,
		PartyIndex:     req.PartyIndex,
	}
	delegateShareCache.Store(req.SessionId, entry)

	logger.Info("Delegate share submitted",
		zap.String("session_id", req.SessionId),
		zap.String("party_id", req.PartyId),
		zap.Int("share_size", len(req.EncryptedShare)))

	return &pb.SubmitDelegateShareResponse{
		Success: true,
	}, nil
}

// toGRPCError converts domain errors to gRPC errors
func toGRPCError(err error) error {
	switch err {
	case entities.ErrSessionExpired:
		return status.Error(codes.DeadlineExceeded, err.Error())
	case entities.ErrSessionFull:
		return status.Error(codes.ResourceExhausted, err.Error())
	case entities.ErrParticipantNotFound:
		return status.Error(codes.NotFound, err.Error())
	case entities.ErrSessionNotInProgress:
		return status.Error(codes.FailedPrecondition, err.Error())
	case entities.ErrInvalidSessionType:
		return status.Error(codes.InvalidArgument, err.Error())
	default:
		return status.Error(codes.Internal, err.Error())
	}
}
