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"
	grpcclient "github.com/rwadurian/mpc-system/services/server-party/adapters/output/grpc"
	"github.com/rwadurian/mpc-system/services/server-party/application/use_cases"
	"github.com/rwadurian/mpc-system/services/server-party/infrastructure/cache"
	"go.uber.org/zap"
)

// Global share cache for delegate parties
var globalShareCache *cache.ShareCache

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 (Delegate Mode)",
		zap.String("environment", cfg.Server.Environment),
		zap.Int("http_port", cfg.Server.HTTPPort))

	// Initialize share cache for delegate parties (15 minute TTL)
	globalShareCache = cache.NewShareCache(15 * time.Minute)
	logger.Info("Share cache initialized", zap.Duration("ttl", 15*time.Minute))

	// 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()

	// Get party ID from environment (or use default)
	partyID := os.Getenv("PARTY_ID")
	if partyID == "" {
		partyID = "server-party-api"
	}

	// Force PARTY_ROLE to delegate for this service
	os.Setenv("PARTY_ROLE", "delegate")

	// Register this party as a delegate party with Message Router
	logger.Info("Registering party with Message Router",
		zap.String("party_id", partyID),
		zap.String("role", "delegate"))

	if err := messageRouter.RegisterParty(ctx, partyID, "delegate", "1.0.0"); err != nil {
		logger.Fatal("Failed to register party", zap.Error(err))
	}

	logger.Info("Party registered successfully",
		zap.String("party_id", partyID),
		zap.String("role", "delegate"))

	// Initialize use cases with nil keyShareRepo (delegate doesn't use DB)
	// The use cases check PARTY_ROLE env var to determine behavior
	participateKeygenUC := use_cases.NewParticipateKeygenUseCase(
		nil, // No database storage for delegate
		sessionClient,
		messageRouter,
		cryptoService,
	)
	participateSigningUC := use_cases.NewParticipateSigningUseCase(
		nil, // No database storage for delegate
		sessionClient,
		messageRouter,
		cryptoService,
	)

	// Start HTTP server
	errChan := make(chan error, 1)
	go func() {
		if err := startHTTPServer(cfg, participateKeygenUC, participateSigningUC, 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,
	participateKeygenUC *use_cases.ParticipateKeygenUseCase,
	participateSigningUC *use_cases.ParticipateSigningUseCase,
	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",
			"role":    "delegate",
		})
	})

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

	{
		// Keygen participation - same as server-party but returns share
		api.POST("/keygen/participate", 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"`
			}

			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("Starting keygen participation (delegate)",
				zap.String("session_id", req.SessionID),
				zap.String("party_id", req.PartyID))

			// Execute keygen synchronously for delegate party
			ctx, cancel := context.WithTimeout(c.Request.Context(), 10*time.Minute)
			defer cancel()

			input := use_cases.ParticipateKeygenInput{
				SessionID: sessionID,
				PartyID:   req.PartyID,
				JoinToken: req.JoinToken,
			}

			output, err := participateKeygenUC.Execute(ctx, input)
			if err != nil {
				logger.Error("Keygen participation failed",
					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("Keygen participation completed (delegate)",
				zap.String("session_id", req.SessionID),
				zap.String("party_id", req.PartyID),
				zap.Bool("success", output.Success))

			// For delegate party, ShareForUser contains the encrypted share
			if len(output.ShareForUser) == 0 {
				c.JSON(http.StatusInternalServerError, gin.H{
					"error": "share not generated for delegate party",
				})
				return
			}

			// Store in cache for retrieval (optional, for async pattern)
			globalShareCache.Store(sessionID, req.PartyID, output.ShareForUser, output.PublicKey)

			// Return share directly
			c.JSON(http.StatusOK, gin.H{
				"success":     true,
				"session_id":  req.SessionID,
				"party_id":    req.PartyID,
				"party_index": output.KeyShare.PartyIndex,
				"share_data":  hex.EncodeToString(output.ShareForUser),
				"public_key":  hex.EncodeToString(output.PublicKey),
			})
		})

		// Signing with user-provided share
		api.POST("/sign/participate", 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"` // User's encrypted share
				MessageHash string `json:"message_hash"`
			}

			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
			}

			var messageHash []byte
			if req.MessageHash != "" {
				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("Starting signing participation (delegate)",
				zap.String("session_id", req.SessionID),
				zap.String("party_id", req.PartyID))

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

			input := use_cases.ParticipateSigningInput{
				SessionID:     sessionID,
				PartyID:       req.PartyID,
				JoinToken:     req.JoinToken,
				MessageHash:   messageHash,
				UserShareData: shareData, // Pass user's share
			}

			output, err := participateSigningUC.Execute(ctx, input)
			if err != nil {
				logger.Error("Signing participation failed",
					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 participation completed (delegate)",
				zap.String("session_id", req.SessionID),
				zap.String("party_id", req.PartyID),
				zap.Bool("success", output.Success))

			// Return signature
			var rHex, sHex string
			if output.R != nil {
				rHex = hex.EncodeToString(output.R.Bytes())
			}
			if output.S != nil {
				sHex = hex.EncodeToString(output.S.Bytes())
			}

			c.JSON(http.StatusOK, gin.H{
				"success":    true,
				"session_id": req.SessionID,
				"party_id":   req.PartyID,
				"signature":  hex.EncodeToString(output.Signature),
				"r":          rHex,
				"s":          sHex,
			})
		})

		// Get user share from cache (for async keygen pattern)
		api.GET("/sessions/:session_id/user-share", func(c *gin.Context) {
			sessionIDStr := c.Param("session_id")

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

			// Retrieve and delete share from cache (one-time retrieval)
			entry, exists := globalShareCache.GetAndDelete(sessionID)
			if !exists {
				c.JSON(http.StatusNotFound, gin.H{
					"error": "Share not found or already retrieved",
					"note":  "Shares can only be retrieved once and expire after 15 minutes",
				})
				return
			}

			logger.Info("User share retrieved successfully",
				zap.String("session_id", sessionIDStr),
				zap.String("party_id", entry.PartyID))

			c.JSON(http.StatusOK, gin.H{
				"session_id": sessionIDStr,
				"party_id":   entry.PartyID,
				"share":      hex.EncodeToString(entry.Share),
				"public_key": hex.EncodeToString(entry.PublicKey),
				"note":       "This share has been deleted from memory and cannot be retrieved again",
			})
		})
	}

	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()
	}
}