package service

import (
	"context"

	externalworkloadpb "github.com/pingcap/kvproto/pkg/externalworkloadpb"
	scalerpb "github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/proto/scaler"
)

type externalWorkloadService struct {
	externalworkloadpb.UnimplementedExternalWorkloadControllerServer
	scaler scalerpb.ScalerServer
}

// NewExternalWorkloadService adapts the scaler service to the upstream external workload API.
func NewExternalWorkloadService(scaler scalerpb.ScalerServer) externalworkloadpb.ExternalWorkloadControllerServer {
	return &externalWorkloadService{scaler: scaler}
}

func (s *externalWorkloadService) Ping(
	ctx context.Context,
	_ *externalworkloadpb.PingRequest,
) (*externalworkloadpb.Response, error) {
	response, err := s.scaler.Ping(ctx, &scalerpb.PingRequest{})
	return externalWorkloadResponse(response), err
}

func (s *externalWorkloadService) RegisterGC(
	ctx context.Context,
	request *externalworkloadpb.RegisterGCRequest,
) (*externalworkloadpb.Response, error) {
	response, err := s.scaler.RegisterGC(ctx, &scalerpb.RegisterGCRequest{
		Header:     scalerHeader(request.GetHeader()),
		DeletionTs: request.GetDeletionTs(),
	})
	return externalWorkloadResponse(response), err
}

func (s *externalWorkloadService) RecycleGC(
	ctx context.Context,
	request *externalworkloadpb.RecycleGCRequest,
) (*externalworkloadpb.Response, error) {
	response, err := s.scaler.RecycleGC(ctx, &scalerpb.RecycleGCRequest{
		Header:    scalerHeader(request.GetHeader()),
		SafePoint: request.GetSafePoint(),
	})
	return externalWorkloadResponse(response), err
}

func (s *externalWorkloadService) RegisterGCV2(
	ctx context.Context,
	request *externalworkloadpb.RegisterGCV2Request,
) (*externalworkloadpb.Response, error) {
	response, err := s.scaler.RegisterGCV2(ctx, &scalerpb.RegisterGCV2Request{
		Header:     scalerHeader(request.GetHeader()),
		SafePoint:  request.GetSafePoint(),
		GcLifeTime: request.GetGcLifeTime(),
	})
	return externalWorkloadResponse(response), err
}

func (s *externalWorkloadService) RecycleGCV2(
	ctx context.Context,
	request *externalworkloadpb.RecycleGCV2Request,
) (*externalworkloadpb.Response, error) {
	response, err := s.scaler.RecycleGCV2(ctx, &scalerpb.RecycleGCV2Request{
		Header:    scalerHeader(request.GetHeader()),
		SafePoint: request.GetSafePoint(),
	})
	return externalWorkloadResponse(response), err
}

func (s *externalWorkloadService) UpdateGCLifeTime(
	ctx context.Context,
	request *externalworkloadpb.UpdateGCLifeTimeRequest,
) (*externalworkloadpb.Response, error) {
	response, err := s.scaler.UpdateGCLifeTime(ctx, &scalerpb.UpdateGCLifeTimeRequest{
		Header:     scalerHeader(request.GetHeader()),
		GcLifeTime: request.GetGcLifeTime(),
	})
	return externalWorkloadResponse(response), err
}

func (s *externalWorkloadService) GetBgTaskConfig(
	ctx context.Context,
	request *externalworkloadpb.GetBgTaskConfigRequest,
) (*externalworkloadpb.GetBgTaskConfigResponse, error) {
	response, err := s.scaler.GetBgTaskConfig(ctx, &scalerpb.GetBgTaskConfigRequest{
		Header:     scalerHeader(request.GetHeader()),
		WorkerType: request.GetWorkerType(),
	})
	if response == nil {
		return nil, err
	}
	return &externalworkloadpb.GetBgTaskConfigResponse{
		Error:            externalWorkloadError(response.GetError()),
		WorkerCount:      response.GetWorkerCount(),
		AutoScaleEnabled: response.GetAutoScaleEnabled(),
	}, err
}

func (s *externalWorkloadService) RegisterBgTask(
	ctx context.Context,
	request *externalworkloadpb.RegisterBgTaskRequest,
) (*externalworkloadpb.Response, error) {
	response, err := s.scaler.RegisterBgTask(ctx, &scalerpb.RegisterBgTaskRequest{
		Header:       scalerHeader(request.GetHeader()),
		TaskType:     request.GetTaskType(),
		TaskKey:      request.GetTaskKey(),
		GlobalTaskId: request.GetGlobalTaskId(),
		SubTaskId:    request.GetSubTaskId(),
		ExecId:       request.GetExecId(),
	})
	return externalWorkloadResponse(response), err
}

func (s *externalWorkloadService) RecycleBgTask(
	ctx context.Context,
	request *externalworkloadpb.RecycleBgTaskRequest,
) (*externalworkloadpb.Response, error) {
	response, err := s.scaler.RecycleBgTask(ctx, &scalerpb.RecycleBgTaskRequest{
		Header:       scalerHeader(request.GetHeader()),
		GlobalTaskId: request.GetGlobalTaskId(),
		SubTaskId:    request.GetSubTaskId(),
	})
	return externalWorkloadResponse(response), err
}

func (s *externalWorkloadService) UpdateBgTaskExecID(
	ctx context.Context,
	request *externalworkloadpb.UpdateBgTaskExecIDRequest,
) (*externalworkloadpb.Response, error) {
	assignments := request.GetAssignments()
	subTaskIDs := make([]int64, 0, len(assignments))
	execIDs := make([]string, 0, len(assignments))
	for _, assignment := range assignments {
		subTaskIDs = append(subTaskIDs, assignment.GetSubTaskId())
		execIDs = append(execIDs, assignment.GetExecId())
	}
	response, err := s.scaler.UpdateBgTaskExecID(ctx, &scalerpb.UpdateBgTaskExecIDRequest{
		Header:       scalerHeader(request.GetHeader()),
		GlobalTaskId: request.GetGlobalTaskId(),
		SubTaskIds:   subTaskIDs,
		ExecIds:      execIDs,
	})
	return externalWorkloadResponse(response), err
}

func (s *externalWorkloadService) RegisterRemoteQuery(
	ctx context.Context,
	request *externalworkloadpb.RegisterRemoteQueryRequest,
) (*externalworkloadpb.Response, error) {
	response, err := s.scaler.RegisterRemoteQuery(ctx, &scalerpb.RegisterRemoteQueryRequest{
		Header:       scalerHeader(request.GetHeader()),
		QueryId:      request.GetQueryId(),
		QueryAddress: request.GetQueryAddress(),
	})
	return externalWorkloadResponse(response), err
}

func (s *externalWorkloadService) RegisterTTLTask(
	ctx context.Context,
	request *externalworkloadpb.RegisterTTLTaskRequest,
) (*externalworkloadpb.Response, error) {
	response, err := s.scaler.RegisterTTLTask(ctx, &scalerpb.RegisterTTLTaskRequest{
		Header:       scalerHeader(request.GetHeader()),
		TableId:      request.GetTableId(),
		TtlJobEnable: request.GetTtlJobEnable(),
	})
	return externalWorkloadResponse(response), err
}

func (s *externalWorkloadService) DeleteTTLTableInfo(
	ctx context.Context,
	request *externalworkloadpb.DeleteTTLTableInfoRequest,
) (*externalworkloadpb.Response, error) {
	response, err := s.scaler.DeleteTTLTableInfo(ctx, &scalerpb.DeleteTTLTableInfoRequest{
		Header:  scalerHeader(request.GetHeader()),
		TableId: request.GetTableId(),
	})
	return externalWorkloadResponse(response), err
}

func (s *externalWorkloadService) RecycleTTLTask(
	ctx context.Context,
	request *externalworkloadpb.RecycleTTLTaskRequest,
) (*externalworkloadpb.Response, error) {
	response, err := s.scaler.RecycleTTLTask(ctx, &scalerpb.RecycleTTLTaskRequest{
		Header:     scalerHeader(request.GetHeader()),
		FinishTime: request.GetCompletedJobCreateTime(),
	})
	return externalWorkloadResponse(response), err
}

func (s *externalWorkloadService) UpdateTTLJobEnable(
	ctx context.Context,
	request *externalworkloadpb.UpdateTTLJobEnableRequest,
) (*externalworkloadpb.Response, error) {
	response, err := s.scaler.UpdateTTLJobEnable(ctx, &scalerpb.UpdateTTLJobEnableRequest{
		Header:       scalerHeader(request.GetHeader()),
		TtlJobEnable: request.GetTtlJobEnable(),
	})
	return externalWorkloadResponse(response), err
}

func (s *externalWorkloadService) RegisterAutoAnalyze(
	ctx context.Context,
	request *externalworkloadpb.RegisterAutoAnalyzeRequest,
) (*externalworkloadpb.Response, error) {
	response, err := s.scaler.RegisterAutoAnalyze(ctx, &scalerpb.RegisterAutoAnalyzeRequest{
		Header: scalerHeader(request.GetHeader()),
		TaskId: request.GetTaskId(),
	})
	return externalWorkloadResponse(response), err
}

func (s *externalWorkloadService) RecycleAutoAnalyze(
	ctx context.Context,
	request *externalworkloadpb.RecycleAutoAnalyzeRequest,
) (*externalworkloadpb.Response, error) {
	response, err := s.scaler.RecycleAutoAnalyze(ctx, &scalerpb.RecycleAutoAnalyzeRequest{
		Header: scalerHeader(request.GetHeader()),
		TaskId: request.GetTaskId(),
	})
	return externalWorkloadResponse(response), err
}

func scalerHeader(header *externalworkloadpb.RequestHeader) *scalerpb.RequestHeader {
	if header == nil {
		return nil
	}
	return &scalerpb.RequestHeader{
		KeyspaceId:   header.GetKeyspaceId(),
		KeyspaceName: header.GetKeyspaceName(),
		TidbPool:     header.GetTidbPool(),
	}
}

func externalWorkloadResponse(response *scalerpb.Response) *externalworkloadpb.Response {
	if response == nil {
		return nil
	}
	return &externalworkloadpb.Response{Error: externalWorkloadError(response.GetError())}
}

func externalWorkloadError(err *scalerpb.Error) *externalworkloadpb.Error {
	if err == nil {
		return nil
	}
	return &externalworkloadpb.Error{
		Type:    externalworkloadpb.ErrorType(err.GetType()),
		Message: err.GetMessage(),
	}
}