package main

import (
	"context"
	"encoding/hex"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/google/uuid"

	"github.com/rwadurian/mpc-system/pkg/config"
	"github.com/rwadurian/mpc-system/pkg/crypto"
	"github.com/rwadurian/mpc-system/pkg/logger"
	"github.com/rwadurian/mpc-system/pkg/tss"
	grpcclient "github.com/rwadurian/mpc-system/services/server-party/adapters/output/grpc"
	"github.com/rwadurian/mpc-system/services/server-party/application/use_cases"
	"go.uber.org/zap"
)

func main() {
	// Parse flags
	configPath := flag.String("config", "", "Path to config file")
	flag.Parse()

	// Load configuration
	cfg, err := config.Load(*configPath)
	if err != nil {
		fmt.Printf("Failed to load config: %v\n", err)
		os.Exit(1)
	}

	// Initialize logger
	if err := logger.Init(&logger.Config{
		Level:    cfg.Logger.Level,
		Encoding: cfg.Logger.Encoding,
	}); err != nil {
		fmt.Printf("Failed to initialize logger: %v\n", err)
		os.Exit(1)
	}
	defer logger.Sync()

	logger.Info("Starting Server Party API Service",
		zap.String("environment", cfg.Server.Environment),
		zap.Int("http_port", cfg.Server.HTTPPort))

	// Initialize crypto service with master key from environment
	masterKeyHex := os.Getenv("MPC_CRYPTO_MASTER_KEY")
	if masterKeyHex == "" {
		masterKeyHex = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"
	}
	masterKey, err := hex.DecodeString(masterKeyHex)
	if err != nil {
		logger.Fatal("Invalid master key format", zap.Error(err))
	}
	cryptoService, err := crypto.NewCryptoService(masterKey)
	if err != nil {
		logger.Fatal("Failed to create crypto service", zap.Error(err))
	}

	// Get API key for authentication
	apiKey := os.Getenv("MPC_API_KEY")
	if apiKey == "" {
		logger.Warn("MPC_API_KEY not set, API will be unprotected")
	}

	// Get gRPC service addresses from environment
	coordinatorAddr := os.Getenv("SESSION_COORDINATOR_ADDR")
	if coordinatorAddr == "" {
		coordinatorAddr = "session-coordinator:50051"
	}
	routerAddr := os.Getenv("MESSAGE_ROUTER_ADDR")
	if routerAddr == "" {
		routerAddr = "message-router:50051"
	}

	// Initialize gRPC clients
	sessionClient, err := grpcclient.NewSessionCoordinatorClient(coordinatorAddr)
	if err != nil {
		logger.Fatal("Failed to connect to session coordinator", zap.Error(err))
	}
	defer sessionClient.Close()

	messageRouter, err := grpcclient.NewMessageRouterClient(routerAddr)
	if err != nil {
		logger.Fatal("Failed to connect to message router", zap.Error(err))
	}
	defer messageRouter.Close()

	// Create shutdown context
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Start HTTP server
	errChan := make(chan error, 1)
	go func() {
		if err := startHTTPServer(cfg, sessionClient, messageRouter, cryptoService, apiKey); err != nil {
			errChan <- fmt.Errorf("HTTP server error: %w", err)
		}
	}()

	// Wait for shutdown signal
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	select {
	case sig := <-sigChan:
		logger.Info("Received shutdown signal", zap.String("signal", sig.String()))
	case err := <-errChan:
		logger.Error("Server error", zap.Error(err))
	}

	// Graceful shutdown
	logger.Info("Shutting down...")
	cancel()

	time.Sleep(5 * time.Second)
	logger.Info("Shutdown complete")

	_ = ctx
}

func startHTTPServer(
	cfg *config.Config,
	sessionClient use_cases.SessionCoordinatorClient,
	messageRouter use_cases.MessageRouterClient,
	cryptoService *crypto.CryptoService,
	apiKey string,
) error {
	if cfg.Server.Environment == "production" {
		gin.SetMode(gin.ReleaseMode)
	}

	router := gin.New()
	router.Use(gin.Recovery())
	router.Use(gin.Logger())

	// Health check
	router.GET("/health", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{
			"status":  "healthy",
			"service": "server-party-api",
		})
	})

	// API routes with optional authentication
	api := router.Group("/api/v1")
	if apiKey != "" {
		api.Use(apiKeyAuth(apiKey))
	}

	{
		// Generate user share - synchronous endpoint that returns the share
		// This is the main endpoint for mpc-service to call
		api.POST("/keygen/generate-user-share", func(c *gin.Context) {
			var req struct {
				SessionID string `json:"session_id" binding:"required"`
				PartyID   string `json:"party_id" binding:"required"`
				JoinToken string `json:"join_token" binding:"required"`
				// Optional: encryption key for the share (provided by user)
				UserPublicKey string `json:"user_public_key"`
			}

			if err := c.ShouldBindJSON(&req); err != nil {
				c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
				return
			}

			sessionID, err := uuid.Parse(req.SessionID)
			if err != nil {
				c.JSON(http.StatusBadRequest, gin.H{"error": "invalid session_id format"})
				return
			}

			logger.Info("Generating user share",
				zap.String("session_id", req.SessionID),
				zap.String("party_id", req.PartyID))

			// Execute keygen synchronously and return the share
			ctx, cancel := context.WithTimeout(c.Request.Context(), 10*time.Minute)
			defer cancel()

			result, err := generateUserShare(
				ctx,
				sessionClient,
				messageRouter,
				cryptoService,
				sessionID,
				req.PartyID,
				req.JoinToken,
				req.UserPublicKey,
			)
			if err != nil {
				logger.Error("Failed to generate user share",
					zap.String("session_id", req.SessionID),
					zap.String("party_id", req.PartyID),
					zap.Error(err))
				c.JSON(http.StatusInternalServerError, gin.H{
					"error":      "keygen failed",
					"details":    err.Error(),
					"session_id": req.SessionID,
					"party_id":   req.PartyID,
				})
				return
			}

			logger.Info("User share generated successfully",
				zap.String("session_id", req.SessionID),
				zap.String("party_id", req.PartyID))

			c.JSON(http.StatusOK, gin.H{
				"success":     true,
				"session_id":  req.SessionID,
				"party_id":    req.PartyID,
				"party_index": result.PartyIndex,
				"share_data":  result.ShareData,
				"public_key":  result.PublicKey,
			})
		})

		// Sign with user share - synchronous endpoint
		api.POST("/sign/with-user-share", func(c *gin.Context) {
			var req struct {
				SessionID   string `json:"session_id" binding:"required"`
				PartyID     string `json:"party_id" binding:"required"`
				JoinToken   string `json:"join_token" binding:"required"`
				ShareData   string `json:"share_data" binding:"required"`
				MessageHash string `json:"message_hash" binding:"required"`
			}

			if err := c.ShouldBindJSON(&req); err != nil {
				c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
				return
			}

			sessionID, err := uuid.Parse(req.SessionID)
			if err != nil {
				c.JSON(http.StatusBadRequest, gin.H{"error": "invalid session_id format"})
				return
			}

			shareData, err := hex.DecodeString(req.ShareData)
			if err != nil {
				c.JSON(http.StatusBadRequest, gin.H{"error": "invalid share_data format (expected hex)"})
				return
			}

			messageHash, err := hex.DecodeString(req.MessageHash)
			if err != nil {
				c.JSON(http.StatusBadRequest, gin.H{"error": "invalid message_hash format (expected hex)"})
				return
			}

			logger.Info("Signing with user share",
				zap.String("session_id", req.SessionID),
				zap.String("party_id", req.PartyID))

			ctx, cancel := context.WithTimeout(c.Request.Context(), 5*time.Minute)
			defer cancel()

			result, err := signWithUserShare(
				ctx,
				sessionClient,
				messageRouter,
				cryptoService,
				sessionID,
				req.PartyID,
				req.JoinToken,
				shareData,
				messageHash,
			)
			if err != nil {
				logger.Error("Failed to sign with user share",
					zap.String("session_id", req.SessionID),
					zap.String("party_id", req.PartyID),
					zap.Error(err))
				c.JSON(http.StatusInternalServerError, gin.H{
					"error":      "signing failed",
					"details":    err.Error(),
					"session_id": req.SessionID,
					"party_id":   req.PartyID,
				})
				return
			}

			logger.Info("Signing completed successfully",
				zap.String("session_id", req.SessionID),
				zap.String("party_id", req.PartyID))

			c.JSON(http.StatusOK, gin.H{
				"success":    true,
				"session_id": req.SessionID,
				"party_id":   req.PartyID,
				"signature":  result.Signature,
				"r":          result.R,
				"s":          result.S,
				"v":          result.V,
			})
		})
	}

	logger.Info("Starting HTTP server", zap.Int("port", cfg.Server.HTTPPort))
	return router.Run(fmt.Sprintf(":%d", cfg.Server.HTTPPort))
}

func apiKeyAuth(expectedKey string) gin.HandlerFunc {
	return func(c *gin.Context) {
		apiKey := c.GetHeader("X-API-Key")
		if apiKey == "" {
			apiKey = c.Query("api_key")
		}
		if apiKey != expectedKey {
			c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid or missing API key"})
			c.Abort()
			return
		}
		c.Next()
	}
}

// UserShareResult contains the result of user share generation
type UserShareResult struct {
	PartyIndex int
	ShareData  string // hex encoded
	PublicKey  string // hex encoded
}

// generateUserShare generates a share for the user without storing it
func generateUserShare(
	ctx context.Context,
	sessionClient use_cases.SessionCoordinatorClient,
	messageRouter use_cases.MessageRouterClient,
	cryptoService *crypto.CryptoService,
	sessionID uuid.UUID,
	partyID string,
	joinToken string,
	userPublicKey string,
) (*UserShareResult, error) {
	// 1. Join session via coordinator
	sessionInfo, err := sessionClient.JoinSession(ctx, sessionID, partyID, joinToken)
	if err != nil {
		return nil, fmt.Errorf("failed to join session: %w", err)
	}

	if sessionInfo.SessionType != "keygen" {
		return nil, fmt.Errorf("invalid session type: expected keygen, got %s", sessionInfo.SessionType)
	}

	// 2. Find self in participants and build party index map
	var selfIndex int
	partyIndexMap := make(map[string]int)
	for _, p := range sessionInfo.Participants {
		partyIndexMap[p.PartyID] = p.PartyIndex
		if p.PartyID == partyID {
			selfIndex = p.PartyIndex
		}
	}

	// 3. Subscribe to messages
	msgChan, err := messageRouter.SubscribeMessages(ctx, sessionID, partyID)
	if err != nil {
		return nil, fmt.Errorf("failed to subscribe to messages: %w", err)
	}

	// 4. Run TSS Keygen protocol
	saveData, publicKey, err := runKeygenProtocol(
		ctx,
		sessionID,
		partyID,
		selfIndex,
		sessionInfo.Participants,
		sessionInfo.ThresholdN,
		sessionInfo.ThresholdT,
		msgChan,
		partyIndexMap,
		messageRouter,
	)
	if err != nil {
		return nil, fmt.Errorf("keygen protocol failed: %w", err)
	}

	// 5. Encrypt share (optionally with user's public key if provided)
	var encryptedShare []byte
	if userPublicKey != "" {
		// TODO: Encrypt with user's public key for end-to-end encryption
		encryptedShare, err = cryptoService.EncryptShare(saveData, partyID)
	} else {
		encryptedShare, err = cryptoService.EncryptShare(saveData, partyID)
	}
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt share: %w", err)
	}

	// 6. Report completion to coordinator
	if err := sessionClient.ReportCompletion(ctx, sessionID, partyID, publicKey); err != nil {
		logger.Error("failed to report completion", zap.Error(err))
		// Don't fail - share is generated
	}

	return &UserShareResult{
		PartyIndex: selfIndex,
		ShareData:  hex.EncodeToString(encryptedShare),
		PublicKey:  hex.EncodeToString(publicKey),
	}, nil
}

// SigningResult contains the result of signing
type SigningResult struct {
	Signature string
	R         string
	S         string
	V         int
}

// signWithUserShare signs using the user's share
func signWithUserShare(
	ctx context.Context,
	sessionClient use_cases.SessionCoordinatorClient,
	messageRouter use_cases.MessageRouterClient,
	cryptoService *crypto.CryptoService,
	sessionID uuid.UUID,
	partyID string,
	joinToken string,
	shareData []byte,
	messageHash []byte,
) (*SigningResult, error) {
	// 1. Join session via coordinator
	sessionInfo, err := sessionClient.JoinSession(ctx, sessionID, partyID, joinToken)
	if err != nil {
		return nil, fmt.Errorf("failed to join session: %w", err)
	}

	if sessionInfo.SessionType != "sign" {
		return nil, fmt.Errorf("invalid session type: expected sign, got %s", sessionInfo.SessionType)
	}

	// 2. Decrypt share
	decryptedShare, err := cryptoService.DecryptShare(shareData, partyID)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt share: %w", err)
	}

	// 3. Find self in participants
	var selfIndex int
	partyIndexMap := make(map[string]int)
	for _, p := range sessionInfo.Participants {
		partyIndexMap[p.PartyID] = p.PartyIndex
		if p.PartyID == partyID {
			selfIndex = p.PartyIndex
		}
	}

	// 4. Subscribe to messages
	msgChan, err := messageRouter.SubscribeMessages(ctx, sessionID, partyID)
	if err != nil {
		return nil, fmt.Errorf("failed to subscribe to messages: %w", err)
	}

	// 5. Run TSS Signing protocol
	signature, r, s, v, err := runSigningProtocol(
		ctx,
		sessionID,
		partyID,
		selfIndex,
		sessionInfo.Participants,
		sessionInfo.ThresholdN,
		sessionInfo.ThresholdT,
		msgChan,
		partyIndexMap,
		messageRouter,
		decryptedShare,
		messageHash,
	)
	if err != nil {
		return nil, fmt.Errorf("signing protocol failed: %w", err)
	}

	// 6. Report completion to coordinator
	if err := sessionClient.ReportCompletion(ctx, sessionID, partyID, signature); err != nil {
		logger.Error("failed to report completion", zap.Error(err))
	}

	return &SigningResult{
		Signature: hex.EncodeToString(signature),
		R:         hex.EncodeToString(r),
		S:         hex.EncodeToString(s),
		V:         v,
	}, nil
}

// runKeygenProtocol runs the TSS keygen protocol
func runKeygenProtocol(
	ctx context.Context,
	sessionID uuid.UUID,
	partyID string,
	selfIndex int,
	participants []use_cases.ParticipantInfo,
	n, t int,
	msgChan <-chan *use_cases.MPCMessage,
	partyIndexMap map[string]int,
	messageRouter use_cases.MessageRouterClient,
) ([]byte, []byte, error) {
	logger.Info("Running keygen protocol",
		zap.String("session_id", sessionID.String()),
		zap.String("party_id", partyID),
		zap.Int("self_index", selfIndex),
		zap.Int("n", n),
		zap.Int("t", t))

	// Create message handler adapter
	msgHandler := &messageHandler{
		sessionID:     sessionID,
		partyID:       partyID,
		messageRouter: messageRouter,
		msgChan:       make(chan *tss.ReceivedMessage, 100),
		partyIndexMap: partyIndexMap,
	}

	// Start message conversion goroutine
	go msgHandler.convertMessages(ctx, msgChan)

	// Create keygen config
	config := tss.KeygenConfig{
		Threshold:    t,
		TotalParties: n,
		Timeout:      10 * time.Minute,
	}

	// Create party list
	allParties := make([]tss.KeygenParty, len(participants))
	for i, p := range participants {
		allParties[i] = tss.KeygenParty{
			PartyID:    p.PartyID,
			PartyIndex: p.PartyIndex,
		}
	}

	selfParty := tss.KeygenParty{
		PartyID:    partyID,
		PartyIndex: selfIndex,
	}

	// Create keygen session
	session, err := tss.NewKeygenSession(config, selfParty, allParties, msgHandler)
	if err != nil {
		return nil, nil, err
	}

	// Run keygen
	result, err := session.Start(ctx)
	if err != nil {
		return nil, nil, err
	}

	logger.Info("Keygen completed successfully",
		zap.String("session_id", sessionID.String()),
		zap.String("party_id", partyID))

	return result.LocalPartySaveData, result.PublicKeyBytes, nil
}

// runSigningProtocol runs the TSS signing protocol
func runSigningProtocol(
	ctx context.Context,
	sessionID uuid.UUID,
	partyID string,
	selfIndex int,
	participants []use_cases.ParticipantInfo,
	n, t int,
	msgChan <-chan *use_cases.MPCMessage,
	partyIndexMap map[string]int,
	messageRouter use_cases.MessageRouterClient,
	shareData []byte,
	messageHash []byte,
) ([]byte, []byte, []byte, int, error) {
	logger.Info("Running signing protocol",
		zap.String("session_id", sessionID.String()),
		zap.String("party_id", partyID),
		zap.Int("self_index", selfIndex))

	// Create message handler adapter
	msgHandler := &messageHandler{
		sessionID:     sessionID,
		partyID:       partyID,
		messageRouter: messageRouter,
		msgChan:       make(chan *tss.ReceivedMessage, 100),
		partyIndexMap: partyIndexMap,
	}

	// Start message conversion goroutine
	go msgHandler.convertMessages(ctx, msgChan)

	// Create signing config
	config := tss.SigningConfig{
		Threshold:    t,
		TotalSigners: n,
		Timeout:      5 * time.Minute,
	}

	// Create party list
	allParties := make([]tss.SigningParty, len(participants))
	for i, p := range participants {
		allParties[i] = tss.SigningParty{
			PartyID:    p.PartyID,
			PartyIndex: p.PartyIndex,
		}
	}

	selfParty := tss.SigningParty{
		PartyID:    partyID,
		PartyIndex: selfIndex,
	}

	// Create signing session
	session, err := tss.NewSigningSession(config, selfParty, allParties, shareData, messageHash, msgHandler)
	if err != nil {
		return nil, nil, nil, 0, err
	}

	// Run signing
	result, err := session.Start(ctx)
	if err != nil {
		return nil, nil, nil, 0, err
	}

	logger.Info("Signing completed successfully",
		zap.String("session_id", sessionID.String()),
		zap.String("party_id", partyID))

	// Convert big.Int to []byte
	var rBytes, sBytes []byte
	if result.R != nil {
		rBytes = result.R.Bytes()
	}
	if result.S != nil {
		sBytes = result.S.Bytes()
	}

	return result.Signature, rBytes, sBytes, result.RecoveryID, nil
}

// messageHandler adapts MPCMessage channel to tss.MessageHandler
type messageHandler struct {
	sessionID     uuid.UUID
	partyID       string
	messageRouter use_cases.MessageRouterClient
	msgChan       chan *tss.ReceivedMessage
	partyIndexMap map[string]int
}

func (h *messageHandler) SendMessage(ctx context.Context, isBroadcast bool, toParties []string, msgBytes []byte) error {
	return h.messageRouter.RouteMessage(ctx, h.sessionID, h.partyID, toParties, 0, msgBytes)
}

func (h *messageHandler) ReceiveMessages() <-chan *tss.ReceivedMessage {
	return h.msgChan
}

func (h *messageHandler) convertMessages(ctx context.Context, inChan <-chan *use_cases.MPCMessage) {
	for {
		select {
		case <-ctx.Done():
			close(h.msgChan)
			return
		case msg, ok := <-inChan:
			if !ok {
				close(h.msgChan)
				return
			}

			fromIndex, exists := h.partyIndexMap[msg.FromParty]
			if !exists {
				continue
			}

			tssMsg := &tss.ReceivedMessage{
				FromPartyIndex: fromIndex,
				IsBroadcast:    msg.IsBroadcast,
				MsgBytes:       msg.Payload,
			}

			select {
			case h.msgChan <- tssMsg:
			case <-ctx.Done():
				return
			}
		}
	}
}