package clientv2

import (
	"context"
	"encoding/json"
	"net/http"
	"net/url"

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

type bgTaskClient interface {
	// GetBgTaskConfig retrieves the configuration of the background task.
	GetBgTaskConfig(ctx context.Context, workerType string) (workerCount int, autoScaleEnabled bool, err error)
	// RegisterBgTask notifies the manager that a background sub-task has been created.
	RegisterBgTask(ctx context.Context, taskType, taskKey string, gTaskID, subTaskID int64, execID string) error
	// RecycleBgTask notifies the manager that a background task has been completed.
	RecycleBgTask(ctx context.Context, gTaskID, subTaskID int64) error
	// UpdateBgTaskExecID updates the execId of the background task.
	UpdateBgTaskExecID(ctx context.Context, gTaskID int64, subTaskIDs []int64, execIDs []string) error
}

var _ bgTaskClient = &client{}

func (c *client) GetBgTaskConfig(ctx context.Context, workerType string) (int, bool, error) {
	// TODO: Use grpc after server refactor.
	workerCount, autoScaleEnabled, err := c.loadBgTaskConfigV2(ctx, workerType)
	if err != nil {
		// try v1 version for old scaler.
		autoScaleEnabled = false
		workerCount, err = c.loadBgTaskConfigV1(ctx, workerType)
		if err != nil {
			return 0, false, err
		}
	}
	return workerCount, autoScaleEnabled, nil
}

func (c *client) getWithContext(ctx context.Context, addr string) (*http.Response, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, addr, nil)
	if err != nil {
		return nil, err
	}
	return c.httpClient.Do(req)
}

func (c *client) loadBgTaskConfigV1(ctx context.Context, workerType string) (int, error) {
	type response struct {
		Paused      bool `json:"paused"`
		WorkerCount int  `json:"worker-count"`
	}
	addr, err := url.JoinPath(c.scalerAddr, "scaler/api/v1/bgtask", workerType+"-worker")
	if err != nil {
		return 0, err
	}
	res, err := c.getWithContext(ctx, addr)
	if err != nil {
		return 0, err
	}
	defer res.Body.Close()

	var cfg response
	err = json.NewDecoder(res.Body).Decode(&cfg)
	if err != nil {
		return 0, err
	}
	if cfg.Paused {
		return 0, ErrHandlerPaused
	}

	return cfg.WorkerCount, nil
}

func (c *client) loadBgTaskConfigV2(ctx context.Context, workerType string) (int, bool, error) {
	type response struct {
		Paused          bool `json:"paused"`
		WorkerCount     int  `json:"worker-count"`
		EnableAutoScale bool `json:"enable-auto-scale"`
	}
	addr, err := url.JoinPath(c.scalerAddr, "scaler/api/v1/bgtask/"+workerType+"-worker", c.keyspaceName)
	if err != nil {
		return 0, false, err
	}
	res, err := c.getWithContext(ctx, addr)
	if err != nil {
		return 0, false, err
	}
	defer res.Body.Close()

	var cfg response
	err = json.NewDecoder(res.Body).Decode(&cfg)
	if err != nil {
		return 0, false, err
	}
	if cfg.Paused {
		return 0, false, ErrHandlerPaused
	}
	if cfg.WorkerCount == 0 { // rewrite worker count if no worker is activated.
		cfg.WorkerCount = 1
	}
	return cfg.WorkerCount, cfg.EnableAutoScale, nil
}

func (c *client) RegisterBgTask(ctx context.Context, taskType, taskKey string, gTaskID, subTaskID int64, execID string) error {
	req := &scalerpb.RegisterBgTaskRequest{
		Header:       requestHeader(c),
		TaskType:     taskType,
		TaskKey:      taskKey,
		GlobalTaskId: gTaskID,
		SubTaskId:    subTaskID,
		ExecId:       execID,
	}
	resp, err := c.grpcClient.RegisterBgTask(ctx, req)
	return getError("RegisterBgTask", err, resp)
}

func (c *client) RecycleBgTask(ctx context.Context, gTaskID, subTaskID int64) error {
	req := &scalerpb.RecycleBgTaskRequest{
		Header:       requestHeader(c),
		GlobalTaskId: gTaskID,
		SubTaskId:    subTaskID,
	}
	resp, err := c.grpcClient.RecycleBgTask(ctx, req)
	return getError("RecycleBgTask", err, resp)
}

func (c *client) UpdateBgTaskExecID(ctx context.Context, gTaskID int64, subtaskIDs []int64, execIDs []string) error {
	req := &scalerpb.UpdateBgTaskExecIDRequest{
		Header:       requestHeader(c),
		GlobalTaskId: gTaskID,
		SubTaskIds:   subtaskIDs,
		ExecIds:      execIDs,
	}
	resp, err := c.grpcClient.UpdateBgTaskExecID(ctx, req)
	return getError("UpdateBgTaskExecID", err, resp)
}