package interceptor

import (
	"context"

	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/grpc"

	globalserverlessv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/serverless/v1"
)

// ProjectIdInterceptor is designed to handle the conversion between ProjectId and OriginalProjectId
// across the serverless-global-server service.
//
// Reason for this implementation:
// When tidb-mgmt-service calls serverless-global-server, it passes both project_id and original_project_id.
// However, within serverless-global-server, all project_id values should be set to the value of original_project_id.
// This interceptor provides a unified way to handle this conversion automatically for all gRPC requests and responses.
//
// Key functionality:
// 1. On incoming requests: Set ProjectId to the value of OriginalProjectId
// 2. On outgoing responses: Set OriginalProjectId to the value of ProjectId and reset ProjectId to 0
//
// This ensures consistency across the service and avoids manual handling of these fields in every method.

func ProjectIdInterceptor() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
		// Handle request
		handleRequest(ctx, req)

		// Call handler
		resp, err := handler(ctx, req)
		if err != nil {
			return resp, err
		}

		// Handle response
		handleResponse(ctx, resp)

		return resp, nil
	}
}

// handleRequest handles ProjectId and OriginalProjectId in the request object
func handleRequest(ctx context.Context, req any) {
	if req == nil {
		return
	}
	// Check if it's ServerlessServiceCreateRequest
	if createReq, ok := req.(*globalserverlessv1.ServerlessServiceCreateRequest); ok {
		if createReq.OriginalProjectId != 0 {
			oldProjectId := createReq.ProjectId
			createReq.ProjectId = createReq.OriginalProjectId
			tilog.Info(ctx, "ProjectIdInterceptor: Updated ProjectId from OriginalProjectId",
				zap.Uint64("req.project_id", oldProjectId),
				zap.Uint64("req.original_project_id", createReq.OriginalProjectId),
				zap.Uint64("rewrite.project_id", createReq.ProjectId),
				zap.String("request_type", "ServerlessServiceCreateRequest"))
		}
		return
	}

	// Check if it's ServerlessServiceListRequest
	if listReq, ok := req.(*globalserverlessv1.ServerlessServiceListRequest); ok {
		if listReq.Filter != nil && len(listReq.Filter.OriginalProjectIds) > 0 {
			oldProjectIds := listReq.Filter.ProjectIds
			listReq.Filter.ProjectIds = listReq.Filter.OriginalProjectIds
			tilog.Info(ctx, "ProjectIdInterceptor: Updated ProjectIds from OriginalProjectIds",
				zap.Any("req.project_ids", oldProjectIds),
				zap.Any("req.original_project_ids", listReq.Filter.OriginalProjectIds),
				zap.Any("rewrite.project_ids", listReq.Filter.ProjectIds),
				zap.String("request_type", "ServerlessServiceListRequest"))
		}
		return
	}

	// Check if it's ServerlessServiceDeleteRequest
	if deleteReq, ok := req.(*globalserverlessv1.ServerlessServiceDeleteRequest); ok {
		if deleteReq.OriginalProjectId != 0 {
			oldProjectId := deleteReq.ProjectId
			deleteReq.ProjectId = deleteReq.OriginalProjectId
			tilog.Info(ctx, "ProjectIdInterceptor: Updated ProjectId from OriginalProjectId",
				zap.Uint64("req.project_id", oldProjectId),
				zap.Uint64("req.original_project_id", deleteReq.OriginalProjectId),
				zap.Uint64("rewrite.project_id", deleteReq.ProjectId),
				zap.String("request_type", "ServerlessServiceDeleteRequest"))
		}
		return
	}

	// Check if it's ServerlessServiceUpdateRequest
	if updateReq, ok := req.(*globalserverlessv1.ServerlessServiceUpdateRequest); ok {
		if updateReq.Cluster != nil && updateReq.Cluster.OriginalProjectId != 0 {
			oldProjectId := updateReq.Cluster.ProjectId
			updateReq.Cluster.ProjectId = updateReq.Cluster.OriginalProjectId
			tilog.Info(ctx, "ProjectIdInterceptor: Updated ProjectId from OriginalProjectId",
				zap.Uint64("req.project_id", oldProjectId),
				zap.Uint64("req.original_project_id", updateReq.Cluster.OriginalProjectId),
				zap.Uint64("rewrite.project_id", updateReq.Cluster.ProjectId),
				zap.String("request_type", "ServerlessServiceUpdateRequest"))
		}
		return
	}
}

// handleResponse handles ProjectId and OriginalProjectId in the response object
func handleResponse(ctx context.Context, resp any) {
	if resp == nil {
		return
	}

	// Check if it's ServerlessServiceCreateResponse
	if createResp, ok := resp.(*globalserverlessv1.ServerlessServiceCreateResponse); ok {
		if createResp.Cluster != nil && createResp.Cluster.ProjectId != 0 {
			oldProjectId := createResp.Cluster.ProjectId
			createResp.Cluster.OriginalProjectId = createResp.Cluster.ProjectId
			createResp.Cluster.ProjectId = 0
			tilog.Info(ctx, "ProjectIdInterceptor: Updated response ProjectId and OriginalProjectId",
				zap.Uint64("resp.project_id", oldProjectId),
				zap.Uint64("rewrite.original_project_id", createResp.Cluster.OriginalProjectId),
				zap.Uint64("rewrite.project_id", createResp.Cluster.ProjectId),
				zap.String("response_type", "ServerlessServiceCreateResponse"))
		}
		return
	}

	// Check if it's ServerlessServiceListResponse
	if listResp, ok := resp.(*globalserverlessv1.ServerlessServiceListResponse); ok {
		for i, cluster := range listResp.Items {
			if cluster.ProjectId != 0 {
				oldProjectId := cluster.ProjectId
				cluster.OriginalProjectId = cluster.ProjectId
				cluster.ProjectId = 0
				tilog.Info(ctx, "ProjectIdInterceptor: Updated response ProjectId and OriginalProjectId for cluster",
					zap.Int("cluster_index", i),
					zap.Uint64("resp.project_id", oldProjectId),
					zap.Uint64("rewrite.original_project_id", cluster.OriginalProjectId),
					zap.Uint64("rewrite.project_id", cluster.ProjectId),
					zap.String("response_type", "ServerlessServiceListResponse"))
			}
		}
		return
	}

	// Check if it's ServerlessServiceDeleteResponse
	if deleteResp, ok := resp.(*globalserverlessv1.ServerlessServiceDeleteResponse); ok {
		if deleteResp.Cluster != nil && deleteResp.Cluster.ProjectId != 0 {
			oldProjectId := deleteResp.Cluster.ProjectId
			deleteResp.Cluster.OriginalProjectId = deleteResp.Cluster.ProjectId
			deleteResp.Cluster.ProjectId = 0
			tilog.Info(ctx, "ProjectIdInterceptor: Updated response ProjectId and OriginalProjectId",
				zap.Uint64("resp.project_id", oldProjectId),
				zap.Uint64("rewrite.original_project_id", deleteResp.Cluster.OriginalProjectId),
				zap.Uint64("rewrite.project_id", deleteResp.Cluster.ProjectId),
				zap.String("response_type", "ServerlessServiceDeleteResponse"))
		}
		return
	}

	// Check if it's ServerlessServiceUpdateResponse
	if updateResp, ok := resp.(*globalserverlessv1.ServerlessServiceUpdateResponse); ok {
		if updateResp.Cluster != nil && updateResp.Cluster.ProjectId != 0 {
			oldProjectId := updateResp.Cluster.ProjectId
			updateResp.Cluster.OriginalProjectId = updateResp.Cluster.ProjectId
			updateResp.Cluster.ProjectId = 0
			tilog.Info(ctx, "ProjectIdInterceptor: Updated response ProjectId and OriginalProjectId",
				zap.Uint64("resp.project_id", oldProjectId),
				zap.Uint64("rewrite.original_project_id", updateResp.Cluster.OriginalProjectId),
				zap.Uint64("rewrite.project_id", updateResp.Cluster.ProjectId),
				zap.String("response_type", "ServerlessServiceUpdateResponse"))
		}
		return
	}
}