package illegalboundary

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/mitchellh/mapstructure"
	"golang.org/x/exp/slices"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"

	clusteragent "github.com/tidbcloud/eks-provider/controllers/utils/cluster_agent"
	"github.com/tidbcloud/eks-provider/controllers/workflow/action"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	apiv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/cluster-agent/v1"
)

const (
	peer_status_error   = "error"
	peer_status_invalid = "invalid"
	peer_status_fixed   = "fixed"
	peer_status_checked = "checked"
)

type PeerOutput struct {
	// Key         string `json:"key"`
	OriginPeer  uint64 `json:"-"`
	RegionID    uint64 `json:"-"`
	StoreID     uint64 `json:"-"`
	Status      string `json:"status"`
	Description string `json:"description"`
}

type Peer struct {
	// Role      uint64 `json:"role"`
	IsLearner bool   `json:"is_learner"`
	ID        uint64 `json:"id"`
	StoreID   uint64 `json:"store_id"`
	RoleName  string `json:"role_name"`
}

type RegionKeyOutput struct {
	ID       uint64 `json:"id"`
	StartKey string `json:"start_key"`
	EndKey   string `json:"end_key"`
	Peers    []Peer `json:"peers"`
}

// repairInvalidLearner
type repairInvalidLearner struct {
	cli client.Client
}

type repairInvalidLearnerInputs struct {
	CheckOnly   bool   `mapstructure:"checkOnly" json:"checkOnly"`
	Namespace   string `mapstructure:"namespace" json:"namespace" validate:"required"`
	ClusterName string `mapstructure:"clusterName" json:"clusterName" validate:"required"`
	InvalidKeys string `mapstructure:"invalidKeys" json:"invalidKeys" validate:"required"`
}

var (
	_ action.Action = &repairInvalidLearner{}
)

func NewRepairInvalidLearner(cli client.Client) action.Action {
	return &repairInvalidLearner{cli: cli}
}

func (p *repairInvalidLearner) Name() cloudv1.WorkflowActionName {
	return "RepairInvalidLearnerAction"
}

func (p *repairInvalidLearner) PreCheck(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep) error {
	inputs := &repairInvalidLearnerInputs{}
	if err := mapstructure.Decode(step.Inputs.Object, inputs); err != nil {
		return fmt.Errorf("decode inputs: %w", err)
	}
	if err := validate.Struct(inputs); err != nil {
		return fmt.Errorf("validate inputs: %w", err)
	}

	return nil
}

func (p *repairInvalidLearner) Handle(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep) error {

	stepStatus, _ := wf.StepStatus(ctx, step.Name)

	outputs := stepStatus.Outputs
	if outputs == nil {
		outputs = map[string]string{}
	}

	var fixedKeys []string
	if _, ok := outputs["fixed_keys"]; ok {
		fixedKeys = strings.Split(outputs["fixed_keys"], ",")
	}

	inputs := &mergeTFRegionKeysInputs{}
	_ = mapstructure.Decode(step.Inputs.Object, inputs)

	cluster := &cloudv1.Cluster{}
	if err := p.cli.Get(ctx, client.ObjectKey{Namespace: inputs.Namespace, Name: inputs.ClusterName}, cluster); err != nil {
		if apierrors.IsNotFound(err) {
			return cloudv1.NewWFNonRetriableError(fmt.Errorf("cluster %s not found", inputs.ClusterName))
		}
		return fmt.Errorf("get cluster: %w", err)
	}

	cli, err := clusteragent.NewClient(cluster)
	if err != nil {
		return fmt.Errorf("create cluster-agent client: %w", err)
	}

	// Decode all keys need to be fixed recorded by the first check-step.
	regionKeys, err := decodeRegionKeys(inputs.InvalidKeys)
	if err != nil {
		return err
	}

	// Execute `pd-ctl region key` for each invalid key, to check if any learner peers exist.
	gotError := false
	keyMapResult := map[string]*PeerOutput{}
	for _, k := range regionKeys {

		// Skip fixed-keys recorded by last Handle.
		if slices.Contains(fixedKeys, k) {
			continue
		}

		fixer := newInvalidLearnerFixer(ctx,
			k,
			inputs.ClusterName,
			inputs.DryRun,
			func(clusterName, command string) (*apiv1.ExecuteCommandResp, error) {
				return cli.ExecCommand(ctx, &apiv1.ExecuteCommandReq{ClusterId: clusterName, Command: command})
			},
		)

		keyMapResult[k] = fixer.Fix()
		if keyMapResult[k].Status == peer_status_fixed || keyMapResult[k].Status == peer_status_checked {
			fixedKeys = append(fixedKeys, k)

		} else {
			gotError = true
		}
	}

	outputs["fixed_keys"] = strings.Join(fixedKeys, ",")
	for k, v := range keyMapResult {
		strBytes, _ := json.Marshal(v)
		outputs[k] = string(strBytes)
	}

	_ = stepStatus.UpdateOutput(ctx, outputs)

	if gotError {
		return fmt.Errorf("repair invalid learner error encountered, wait for next try")
	}
	return nil
}

// Pick out TiFlash learners from all peers.
func locateLearnerPeer(peers []Peer) (learnerPeers []Peer) {
	learnerPeers = []Peer{}
	for _, peer := range peers {
		if strings.ToLower(peer.RoleName) != PEER_ROLE_NAME_LEARNER_LC {
			continue
		}
		learnerPeers = append(learnerPeers, peer)
	}
	return learnerPeers
}

func parseRegionKeysResp(resp *apiv1.ExecuteCommandResp) (r RegionKeyOutput, err error) {

	err = parsePDCTLRespError(resp)
	if err != nil {
		return
	}

	stdout := resp.StdoutMsg
	err = json.Unmarshal([]byte(stdout), &r)
	if err != nil {
		err = fmt.Errorf("unmarshal pd-ctl region key resp failed: %w", err)
		return
	}

	return
}

type invalidLearnerFixer struct {
	ctx         context.Context
	dryRun      bool
	key         string
	clusterName string
	regionId    uint64
	agentCmder  func(clusterName, command string) (*apiv1.ExecuteCommandResp, error)
}

func newInvalidLearnerFixer(
	ctx context.Context,
	key, clusterName string,
	dryRun bool,
	agentCmder func(clusterName, command string) (*apiv1.ExecuteCommandResp, error),
) *invalidLearnerFixer {
	return &invalidLearnerFixer{
		ctx:         ctx,
		dryRun:      dryRun,
		key:         key,
		clusterName: clusterName,
		agentCmder:  agentCmder,
	}
}

// Checking whether there are any learner peers among all peers (including both TiKV and TiFlash peers).
// If any exist, extract peers which are of the TiFlash-learner type.
func (i *invalidLearnerFixer) checkRegionPeers() ([]Peer, error) {
	resp, err := i.agentCmder(i.clusterName, strings.Join([]string{
		PD_PATROL_CTL_BIN,
		fmt.Sprintf("-u%s", PD_HTTPS_URL),
		"region", "key", i.key,
	}, " "))
	if err != nil {
		return nil, fmt.Errorf("failed execute PD-CTL region key: %w", err)
	}

	respOut, err := parseRegionKeysResp(resp)
	if err != nil {
		return nil, err
	}
	i.regionId = respOut.ID

	return locateLearnerPeer(respOut.Peers), nil
}

func (i *invalidLearnerFixer) removePeer(peerId, storeId uint64) error {
	resp, err := i.agentCmder(i.clusterName, strings.Join([]string{
		PD_PATROL_CTL_BIN, fmt.Sprintf("-u%s", PD_HTTPS_URL),
		"operator", "add", "remove-peer",
		fmt.Sprintf("%d", i.regionId),
		fmt.Sprintf("%d", storeId)}, " "))
	if err != nil {
		return fmt.Errorf("failed execute PD-CTL region key: %w", err)
	}

	return parsePDCTLRespError(resp)
}

func (i *invalidLearnerFixer) Fix() (p *PeerOutput) {
	p = &PeerOutput{}

	// Retrieve the original learners first.
	originLearners, err := i.checkRegionPeers()
	if err != nil {
		p.Status = peer_status_error
		p.Description = err.Error()
		return
	}

	if len(originLearners) <= 0 {
		p.Status = peer_status_checked
		return
	}

	// If there is any learner peer, mark this key as Invalid
	p.Status = peer_status_invalid

	// If dry-run set, return after check.
	if i.dryRun {
		p.Description = "Skipped fix due to dry-run set"
		return
	}

	// Used to verify if any remain unremoved or unupdated after the fix
	var originPeerIds []uint64
	for _, peer := range originLearners {

		peerId := peer.ID
		storeId := peer.StoreID

		originPeerIds = append(originPeerIds, peerId)

		// remove-peer for each origin-peer.
		err := i.removePeer(peerId, storeId)
		if err != nil {
			// Exit immediately when errors occur while removing all TiFlash peers for this key
			// And wait for the next reconcile to retry the fix.
			p.Status = peer_status_error
			p.Description = fmt.Sprintf("error encountered on one of peers: [key: %s], [err: %s]", i.key, err.Error())
			return
		}
	}

	// compare new learners with origin, if there are no learners or learners have changed, means fix succeed.
	checkStillExistFunc := func() (stillExistPeers []string) {
		newLearners, _ := i.checkRegionPeers()
		for _, peer := range newLearners {
			if slices.Contains(originPeerIds, peer.ID) {
				// Gather all still existed learners
				stillExistPeers = append(stillExistPeers, fmt.Sprintf("%d", peer.ID))
			}
		}
		return
	}

	// Since remove-peer is asynchronous, a short loop is used here to check whether the fix has taken effect.
	//
	// If some peers are still not updated after the loop, the fix is considered failed,
	// and wait for the next reconcile to re-fix.
	retry, maxRetry := 0, 3
	sleepInterval := time.Second * 5
	for {
		failedLearners := checkStillExistFunc()
		if len(failedLearners) <= 0 {
			p.Status = peer_status_fixed
			break
		}

		retry++
		if retry > maxRetry {
			p.Status = peer_status_invalid
			p.Description = fmt.Sprintf("peer [%s] still exist", strings.Join(failedLearners, ","))
			break
		}

		log.FromContext(i.ctx).Info("retrying to check learner is fixed ......", "times",
			fmt.Sprintf("%d/%d", retry, maxRetry))
		time.Sleep(sleepInterval)
	}

	return
}