package grpc

import (
	"context"

	"github.com/google/uuid"
	coordinator "github.com/rwadurian/mpc-system/api/grpc/coordinator/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"
)

// SessionCoordinatorClient implements use_cases.SessionCoordinatorClient
type SessionCoordinatorClient struct {
	resilientConn *grpcutil.ResilientConn
	address       string
	retryCfg      retry.Config
}

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

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

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

// getConn returns the current gRPC connection
func (c *SessionCoordinatorClient) getConn() *grpc.ClientConn {
	return c.resilientConn.GetConn()
}

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

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

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

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

	return retry.Do(ctx, c.retryCfg, "JoinSession", func() (*use_cases.SessionInfo, error) {
		// Make the gRPC call using the raw connection
		resp := &coordinator.JoinSessionResponse{}
		err := c.getConn().Invoke(ctx, "/mpc.coordinator.v1.SessionCoordinator/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
		// Note: OtherParties should include ALL participants (including self) from coordinator
		participants := make([]use_cases.ParticipantInfo, len(resp.OtherParties))
		for i, p := range resp.OtherParties {
			// Debug: Log what we received from gRPC
			logger.Debug("gRPC client - received party_index from protobuf response",
				zap.String("party_id", p.PartyId),
				zap.Int32("proto_party_index", p.PartyIndex),
				zap.Int("converted_party_index", int(p.PartyIndex)))

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

		sessionInfo := &use_cases.SessionInfo{
			SessionID:    sessionID,
			SessionType:  resp.SessionInfo.SessionType,
			ThresholdN:   int(resp.SessionInfo.ThresholdN),
			ThresholdT:   int(resp.SessionInfo.ThresholdT),
			MessageHash:  resp.SessionInfo.MessageHash,
			Participants: participants,
		}

		logger.Info("Joined session successfully",
			zap.String("session_id", sessionID.String()),
			zap.String("party_id", partyID),
			zap.String("session_type", sessionInfo.SessionType))

		return sessionInfo, nil
	})
}

// ReportCompletion reports that a party has completed the MPC protocol
// Includes automatic retry with exponential backoff for transient failures
func (c *SessionCoordinatorClient) ReportCompletion(
	ctx context.Context,
	sessionID uuid.UUID,
	partyID string,
	resultData []byte,
) error {
	req := &coordinator.ReportCompletionRequest{
		SessionId: sessionID.String(),
		PartyId:   partyID,
		PublicKey: resultData, // For keygen, this is public key; for signing, this is signature
	}

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

		logger.Info("Reported completion",
			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
// Includes automatic retry with exponential backoff for transient failures
func (c *SessionCoordinatorClient) MarkPartyReady(
	ctx context.Context,
	sessionID uuid.UUID,
	partyID string,
) (bool, error) {
	req := &coordinator.MarkPartyReadyRequest{
		SessionId: sessionID.String(),
		PartyId:   partyID,
	}

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

		logger.Info("Marked party ready",
			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
// Includes automatic retry with exponential backoff for transient failures
func (c *SessionCoordinatorClient) GetSessionStatus(
	ctx context.Context,
	sessionID uuid.UUID,
) (string, error) {
	req := &coordinator.GetSessionStatusRequest{
		SessionId: sessionID.String(),
	}

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

		return resp.Status, nil
	})
}

// findPartyIndex finds the party index from the list of parties
func findPartyIndex(parties []*coordinator.PartyInfo, partyID string) int {
	for _, p := range parties {
		if p.PartyId == partyID {
			return int(p.PartyIndex)
		}
	}
	return 0
}
