package az

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

	opv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	corev1 "k8s.io/api/core/v1"
	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"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	caapiv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/cluster-agent/v1"
	apicommon "github.com/tidbcloud/infra-provider/api/tidbcloud/common"
	"github.com/tidbcloud/infra-provider/pkg/namer"
)

const (
	defaultGracefulTimeout = 10 * time.Minute

	lastSubStepKey = "last-sub-step"

	offlineCordonedNodesKey = "cordoned-nodes"
	offlineDeletedPodsKey   = "deleted-pods"
	offlinePausedASGsKey    = "paused-asgs"

	recoverUncordonedNodesKey = "uncordoned-nodes"
	recoverUnpausedASGsKey    = "unpaused-asgs"

	offlineSubStepPauseASGs   = "pause-asgs"
	recoverSubStepUnpauseASGs = "unpause-asgs"
)

var (
	sqlPlaneComponentTypes = []apicommon.ComponentType{
		apicommon.ComponentType_TIDB,
		apicommon.ComponentType_TIPROXY,
	}
	allTargetComponentTypes = []apicommon.ComponentType{
		apicommon.ComponentType_PD,
		apicommon.ComponentType_TIDB,
		apicommon.ComponentType_TIKV,
		apicommon.ComponentType_TIFLASH,
		apicommon.ComponentType_TIPROXY,
	}
)

type azResources struct {
	Zone      string // empty means all zones (no AZ filter)
	NodeNames []string
	Pods      []*caapiv1.Pod
	NodeZone  map[string]string // nodeName -> zone, always populated for pod zone lookup
}

type pdLeaderInfo struct {
	Name string `json:"name"`
}

// pdStoresInfo matches pd-ctl `store` / PD API /pd/api/v1/stores output.
type pdStoresInfo struct {
	Stores []pdStoreEntry `json:"stores"`
}

type pdStoreEntry struct {
	Store pdStoreMeta `json:"store"`
}

type pdStoreMeta struct {
	ID uint64 `json:"id"`
	//nolint:tagliatelle // pd-ctl JSON uses snake_case
	StateName string         `json:"state_name"`
	Labels    []pdStoreLabel `json:"labels"`
}

type pdStoreLabel struct {
	Key   string `json:"key"`
	Value string `json:"value"`
}

func (s pdStoreMeta) labelValue(key string) string {
	for _, l := range s.Labels {
		if l.Key == key {
			return l.Value
		}
	}
	return ""
}

func parsePDStores(stdout string) ([]pdStoreMeta, error) {
	var info pdStoresInfo
	if err := json.Unmarshal([]byte(stdout), &info); err != nil {
		return nil, fmt.Errorf("unmarshal pd store output: %w; stdout=%q", err, truncateForErr(stdout, 256))
	}
	stores := make([]pdStoreMeta, 0, len(info.Stores))
	for _, entry := range info.Stores {
		stores = append(stores, entry.Store)
	}
	return stores, nil
}

func parseGracefulTimeout(raw string) time.Duration {
	if raw == "" {
		return defaultGracefulTimeout
	}
	d, err := time.ParseDuration(raw)
	if err != nil || d <= 0 {
		return defaultGracefulTimeout
	}
	return d
}

func componentTypeName(ct apicommon.ComponentType) string {
	switch ct {
	case apicommon.ComponentType_PD:
		return "pd"
	case apicommon.ComponentType_TIKV:
		return "tikv"
	case apicommon.ComponentType_TIDB:
		return "tidb"
	case apicommon.ComponentType_TIFLASH:
		return "tiflash"
	case apicommon.ComponentType_TIPROXY:
		return "tiproxy"
	default:
		return strings.ToLower(ct.String())
	}
}

func isTargetComponent(ct apicommon.ComponentType, targets []apicommon.ComponentType) bool {
	for _, target := range targets {
		if ct == target {
			return true
		}
	}
	return false
}

func discoverAZResources(ctx context.Context, caClient clusteragent.Client, clusterID, zone string, group Group) (*azResources, error) {
	nodesResp, err := caClient.ListNodes(ctx, clusterID)
	if err != nil {
		return nil, fmt.Errorf("list nodes: %w", err)
	}
	podsResp, err := caClient.ListPods(ctx, clusterID)
	if err != nil {
		return nil, fmt.Errorf("list pods: %w", err)
	}

	targets := group.Components()
	nodeZone := make(map[string]string, len(nodesResp.Nodes))
	nodeClusterLabel := make(map[string]string, len(nodesResp.Nodes))
	nodeNames := make([]string, 0)
	for _, node := range nodesResp.Nodes {
		if node.GetK8SMetadata() == nil {
			continue
		}
		name := node.GetK8SMetadata().GetName()
		nodeZone[name] = node.GetZone()
		if labels := node.GetK8SMetadata().GetLabels(); labels != nil {
			nodeClusterLabel[name] = labels[namer.ClusterSelectorKey]
		}
		// Empty zone means return resources across all AZs.
		if zone != "" && node.GetZone() != zone {
			continue
		}
		if !isTargetComponent(node.GetComponentType(), targets) {
			continue
		}
		if !group.MatchesNode(node.GetComponentType(), nodeClusterLabel[name]) {
			continue
		}
		nodeNames = append(nodeNames, name)
	}

	pods := make([]*caapiv1.Pod, 0)
	for _, pod := range podsResp.Pods {
		if pod.GetNodeName() == "" {
			continue
		}
		podZone, ok := nodeZone[pod.GetNodeName()]
		if !ok {
			continue
		}
		if zone != "" && podZone != zone {
			continue
		}
		if !isTargetComponent(pod.GetComponentType(), targets) {
			continue
		}
		if !group.MatchesNode(pod.GetComponentType(), nodeClusterLabel[pod.GetNodeName()]) {
			continue
		}
		pods = append(pods, pod)
	}

	return &azResources{
		Zone:      zone,
		NodeNames: nodeNames,
		Pods:      pods,
		NodeZone:  nodeZone,
	}, nil
}

func podsByComponent(pods []*caapiv1.Pod, components []apicommon.ComponentType) map[apicommon.ComponentType][]string {
	out := make(map[apicommon.ComponentType][]string)
	for _, pod := range pods {
		if !isTargetComponent(pod.GetComponentType(), components) {
			continue
		}
		out[pod.GetComponentType()] = append(out[pod.GetComponentType()], pod.GetK8SMetadata().GetName())
	}
	return out
}

func getPDURL(ctx context.Context, caClient clusteragent.Client, clusterID string) (string, error) {
	topo, err := caClient.GetClusterTopology(ctx, clusterID)
	if err != nil {
		return "", fmt.Errorf("get cluster topology: %w", err)
	}
	if len(topo.PdNodes) == 0 {
		return "", fmt.Errorf("no PD nodes found for cluster %s", clusterID)
	}
	// Prefer the PD client Service over a specific pod peer so traffic stays on
	// Ready endpoints after an AZ offline has shut down some PD members.
	return pdClientServiceURL(topo.PdNodes[0]), nil
}

// pdClientServiceURL converts a pod peer host like
// "db-pd-0.db-pd-peer.tidb123.svc" into "https://db-pd.tidb123.svc:2379".
func pdClientServiceURL(pd *caapiv1.GetClusterTopologyResp_PDNodeInfo) string {
	parts := strings.Split(pd.GetHost(), ".")
	if len(parts) >= 3 && strings.HasSuffix(parts[1], "-peer") {
		svc := strings.TrimSuffix(parts[1], "-peer")
		ns := parts[2]
		return fmt.Sprintf("https://%s.%s.svc:%d", svc, ns, pd.GetPort())
	}
	return fmt.Sprintf("https://%s:%d", pd.GetHost(), pd.GetPort())
}

func execPDCTL(ctx context.Context, caClient clusteragent.Client, clusterID, pdURL, subCmd string) (string, error) {
	cmd := fmt.Sprintf("pd-ctl -u %s %s", pdURL, subCmd)
	resp, err := caClient.ExecCommand(ctx, &caapiv1.ExecuteCommandReq{
		ClusterId: clusterID,
		Command:   cmd,
	})
	if err != nil {
		return "", fmt.Errorf("execute %q: %w", cmd, err)
	}
	if resp.ExitCode != nil && *resp.ExitCode != 0 {
		return "", fmt.Errorf("execute %q failed: code=%d stderr=%s stdout=%s",
			cmd, *resp.ExitCode, strings.TrimSpace(resp.StderrMsg), truncateForErr(resp.StdoutMsg, 256))
	}
	return resp.StdoutMsg, nil
}

func truncateForErr(s string, max int) string {
	s = strings.TrimSpace(s)
	if len(s) <= max {
		return s
	}
	return s[:max] + "..."
}

func joinNames(names []string) string {
	if len(names) == 0 {
		return ""
	}
	sorted := slices.Clone(names)
	slices.Sort(sorted)
	return strings.Join(sorted, ",")
}

func parseNames(raw string) []string {
	if raw == "" {
		return nil
	}
	return strings.Split(raw, ",")
}

func mergeNames(existing string, added []string) []string {
	seen := make(map[string]struct{})
	var out []string
	for _, name := range append(parseNames(existing), added...) {
		if name == "" {
			continue
		}
		if _, ok := seen[name]; ok {
			continue
		}
		seen[name] = struct{}{}
		out = append(out, name)
	}
	return out
}

func transferPDLeaderIfNeeded(ctx context.Context, caClient clusteragent.Client, clusterID, pdURL, faultZone string, resources *azResources) error {
	logger := log.FromContext(ctx)
	stdout, err := execPDCTL(ctx, caClient, clusterID, pdURL, "member leader show")
	if err != nil {
		return err
	}
	var leader pdLeaderInfo
	if err := json.Unmarshal([]byte(stdout), &leader); err != nil {
		return fmt.Errorf("unmarshal pd leader show output: %w", err)
	}
	if leader.Name == "" {
		return fmt.Errorf("pd leader name is empty")
	}

	leaderInFaultAZ := false
	for _, pod := range resources.Pods {
		if pod.GetComponentType() != apicommon.ComponentType_PD {
			continue
		}
		if pod.GetK8SMetadata() == nil || pod.GetK8SMetadata().GetName() != leader.Name {
			continue
		}
		if resources.NodeZone[pod.GetNodeName()] == faultZone {
			leaderInFaultAZ = true
		}
		break
	}
	if !leaderInFaultAZ {
		logger.Info("PD leader is not in fault AZ, skip transfer", "leader", leader.Name, "zone", faultZone)
		return nil
	}

	var target string
	for _, pod := range resources.Pods {
		if pod.GetComponentType() != apicommon.ComponentType_PD {
			continue
		}
		if pod.GetK8SMetadata() == nil {
			continue
		}
		name := pod.GetK8SMetadata().GetName()
		if name == "" || name == leader.Name {
			continue
		}
		if resources.NodeZone[pod.GetNodeName()] == faultZone {
			continue
		}
		if phase := pod.GetPhase(); phase != "" && phase != string(corev1.PodRunning) {
			continue
		}
		target = name
		break
	}
	if target == "" {
		return fmt.Errorf("no healthy PD member found outside fault AZ %s (leader %s)", faultZone, leader.Name)
	}
	logger.Info("transferring PD leader away from fault AZ", "leader", leader.Name, "target", target, "zone", faultZone)
	_, err = execPDCTL(ctx, caClient, clusterID, pdURL, fmt.Sprintf("member leader transfer %s", target))
	return err
}

func cordonNodes(ctx context.Context, remoteCli client.Client, nodeNames []string) ([]string, error) {
	logger := log.FromContext(ctx)
	cordoned := make([]string, 0, len(nodeNames))
	for _, nodeName := range nodeNames {
		node := &corev1.Node{}
		if err := remoteCli.Get(ctx, client.ObjectKey{Name: nodeName}, node); err != nil {
			if apierrors.IsNotFound(err) {
				logger.Info("skip cordon node, node not found", "nodeName", nodeName)
				continue
			}
			return nil, fmt.Errorf("get node %s: %w", nodeName, err)
		}
		if node.Spec.Unschedulable {
			logger.Info("node already cordoned", "nodeName", nodeName)
			cordoned = append(cordoned, nodeName)
			continue
		}
		logger.Info("cordoning node", "nodeName", nodeName)
		node.Spec.Unschedulable = true
		if err := remoteCli.Update(ctx, node); err != nil {
			return nil, fmt.Errorf("cordon node %s: %w", nodeName, err)
		}
		cordoned = append(cordoned, nodeName)
	}
	return cordoned, nil
}

func uncordonNodes(ctx context.Context, remoteCli client.Client, nodeNames []string) ([]string, error) {
	logger := log.FromContext(ctx)
	uncordoned := make([]string, 0, len(nodeNames))
	for _, nodeName := range nodeNames {
		node := &corev1.Node{}
		if err := remoteCli.Get(ctx, client.ObjectKey{Name: nodeName}, node); err != nil {
			if apierrors.IsNotFound(err) {
				logger.Info("skip uncordon node, node not found", "nodeName", nodeName)
				continue
			}
			return nil, fmt.Errorf("get node %s: %w", nodeName, err)
		}
		if !node.Spec.Unschedulable {
			logger.Info("node already uncordoned", "nodeName", nodeName)
			uncordoned = append(uncordoned, nodeName)
			continue
		}
		logger.Info("uncordoning node", "nodeName", nodeName)
		node.Spec.Unschedulable = false
		if err := remoteCli.Update(ctx, node); err != nil {
			return nil, fmt.Errorf("uncordon node %s: %w", nodeName, err)
		}
		uncordoned = append(uncordoned, nodeName)
	}
	return uncordoned, nil
}

func deletePod(ctx context.Context, remoteCli client.Client, namespace, podName string, force bool) error {
	logger := log.FromContext(ctx).WithValues("namespace", namespace, "podName", podName, "force", force)
	pod := &corev1.Pod{}
	if err := remoteCli.Get(ctx, client.ObjectKey{Namespace: namespace, Name: podName}, pod); err != nil {
		if apierrors.IsNotFound(err) {
			logger.Info("skip delete pod, pod not found")
			return nil
		}
		return fmt.Errorf("get pod %s/%s: %w", namespace, podName, err)
	}
	if pod.DeletionTimestamp != nil && !force {
		logger.Info("skip delete pod, pod already terminating")
		return nil
	}
	var deleteOpts []client.DeleteOption
	if force {
		deleteOpts = append(deleteOpts, client.GracePeriodSeconds(0))
	}
	logger.Info("deleting pod")
	if err := remoteCli.Delete(ctx, pod, deleteOpts...); err != nil && !apierrors.IsNotFound(err) {
		if force {
			return fmt.Errorf("force delete pod %s/%s: %w", namespace, podName, err)
		}
		return fmt.Errorf("delete pod %s/%s: %w", namespace, podName, err)
	}
	logger.Info("deleted pod")
	return nil
}

func gracefulShutdownPods(ctx context.Context, remoteCli client.Client, namespace string, resources *azResources, targets []apicommon.ComponentType) ([]string, error) {
	logger := log.FromContext(ctx)
	deleted := make([]string, 0)
	grouped := podsByComponent(resources.Pods, targets)
	for _, comp := range targets {
		for _, podName := range grouped[comp] {
			logger.Info("graceful shutdown pod", "component", componentTypeName(comp), "podName", podName)
			if err := deletePod(ctx, remoteCli, namespace, podName, false); err != nil {
				return deleted, err
			}
			deleted = append(deleted, podName)
		}
	}
	return deleted, nil
}

func isPodShutdownComplete(pod *corev1.Pod) bool {
	if pod.DeletionTimestamp != nil {
		return false
	}
	// After graceful delete, the operator recreates the pod with the same name.
	// Cordoned nodes block scheduling, so the replacement stays Pending without
	// a node assignment and should be treated as shutdown complete.
	return pod.Status.Phase == corev1.PodPending && pod.Spec.NodeName == ""
}

func remainingShutdownPods(ctx context.Context, remoteCli client.Client, namespace string, resources *azResources, targets []apicommon.ComponentType) ([]string, error) {
	var remaining []string
	grouped := podsByComponent(resources.Pods, targets)
	for _, comp := range targets {
		for _, podName := range grouped[comp] {
			pod := &corev1.Pod{}
			if err := remoteCli.Get(ctx, client.ObjectKey{Namespace: namespace, Name: podName}, pod); err != nil {
				if apierrors.IsNotFound(err) {
					continue
				}
				return nil, fmt.Errorf("get pod %s/%s: %w", namespace, podName, err)
			}
			if isPodShutdownComplete(pod) {
				continue
			}
			remaining = append(remaining, podName)
		}
	}
	return remaining, nil
}

func forceShutdownPods(ctx context.Context, remoteCli client.Client, namespace string, podNames []string) ([]string, error) {
	logger := log.FromContext(ctx)
	deleted := make([]string, 0, len(podNames))
	for _, podName := range podNames {
		logger.Info("force shutdown pod", "podName", podName)
		if err := deletePod(ctx, remoteCli, namespace, podName, true); err != nil {
			return deleted, err
		}
		deleted = append(deleted, podName)
	}
	return deleted, nil
}

// waitPodsRunning lists all cluster pods via cluster-agent, filters to the
// group's components, and requires every matching pod to be Running. This
// deliberately ignores node/zone so unscheduled Pending pods after AZ offline
// are not skipped.
func waitPodsRunning(ctx context.Context, caClient clusteragent.Client, clusterID string, group Group) ([]string, error) {
	podsResp, err := caClient.ListPods(ctx, clusterID)
	if err != nil {
		return nil, fmt.Errorf("list pods: %w", err)
	}
	nodesResp, err := caClient.ListNodes(ctx, clusterID)
	if err != nil {
		return nil, fmt.Errorf("list nodes: %w", err)
	}
	nodeClusterLabel := make(map[string]string, len(nodesResp.Nodes))
	for _, node := range nodesResp.Nodes {
		if node.GetK8SMetadata() == nil {
			continue
		}
		if labels := node.GetK8SMetadata().GetLabels(); labels != nil {
			nodeClusterLabel[node.GetK8SMetadata().GetName()] = labels[namer.ClusterSelectorKey]
		}
	}
	targets := group.Components()
	runningPods := make([]string, 0)
	for _, pod := range podsResp.Pods {
		if !isTargetComponent(pod.GetComponentType(), targets) {
			continue
		}
		if !group.MatchesNode(pod.GetComponentType(), nodeClusterLabel[pod.GetNodeName()]) {
			continue
		}
		podName := ""
		if pod.GetK8SMetadata() != nil {
			podName = pod.GetK8SMetadata().GetName()
		}
		if !strings.EqualFold(pod.GetPhase(), string(corev1.PodRunning)) {
			return runningPods, fmt.Errorf("pod %s phase is %s, waiting for Running", podName, pod.GetPhase())
		}
		runningPods = append(runningPods, podName)
	}
	if len(runningPods) == 0 {
		return runningPods, fmt.Errorf("no target component pods found for group %s", group.ID())
	}
	return runningPods, nil
}

func isTCReady(tc *opv1alpha1.TidbCluster) bool {
	for _, cond := range tc.Status.Conditions {
		if cond.Type == opv1alpha1.TidbClusterReady && cond.Status == corev1.ConditionTrue {
			return true
		}
	}
	return false
}

func verifyGroupTCReady(ctx context.Context, remoteCli client.Client, namespace string, group Group) error {
	tcName := group.Namer().GetTiDBName()
	var tc opv1alpha1.TidbCluster
	if err := remoteCli.Get(ctx, client.ObjectKey{Namespace: namespace, Name: tcName}, &tc); err != nil {
		return fmt.Errorf("get tidbcluster %s for group %s: %w", tcName, group.ID(), err)
	}
	if !isTCReady(&tc) {
		return fmt.Errorf("tidbcluster %s for group %s is not ready, waiting", tcName, group.ID())
	}
	return nil
}

func ensureComponentPodsRunning(pods []*caapiv1.Pod) ([]string, error) {
	runningPods := make([]string, 0)
	for _, pod := range pods {
		if !isTargetComponent(pod.GetComponentType(), allTargetComponentTypes) {
			continue
		}
		podName := ""
		if pod.GetK8SMetadata() != nil {
			podName = pod.GetK8SMetadata().GetName()
		}
		if !strings.EqualFold(pod.GetPhase(), string(corev1.PodRunning)) {
			return runningPods, fmt.Errorf("pod %s phase is %s, waiting for Running", podName, pod.GetPhase())
		}
		runningPods = append(runningPods, podName)
	}
	if len(runningPods) == 0 {
		return runningPods, fmt.Errorf("no target component pods found")
	}
	return runningPods, nil
}

func verifyStoresOffline(ctx context.Context, caClient clusteragent.Client, clusterID, pdURL, zone string) error {
	stdout, err := execPDCTL(ctx, caClient, clusterID, pdURL, "store")
	if err != nil {
		return err
	}
	stores, err := parsePDStores(stdout)
	if err != nil {
		return err
	}
	for _, store := range stores {
		if store.labelValue("zone") != zone {
			continue
		}
		// Only Up means the store is still serving. After AZ shutdown it typically
		// becomes Disconnected/Down; Offline/Tombstone come from store delete.
		if strings.EqualFold(store.StateName, "up") {
			return fmt.Errorf("store %d in zone %s still in state %s", store.ID, zone, store.StateName)
		}
	}
	return nil
}

// subStepResult drives the sub-step loop in Handle:
//   - err != nil  → keep the workflow step Running and retry later
//   - done        → Handle returns nil (workflow step completes)
//   - next != ""  → checkpoint and continue with the next sub-step in-process
type subStepResult struct {
	next  string
	extra map[string]string
	done  bool
}

func advanceTo(next string, extra map[string]string) (subStepResult, error) {
	return subStepResult{next: next, extra: extra}, nil
}

func subStepDone() (subStepResult, error) {
	return subStepResult{done: true}, nil
}

func updateSubStep(ctx context.Context, wfCli client.Client, wf *cloudv1.Workflow, step cloudv1.WorkflowStep, subStep string, extra map[string]string) error {
	stepStatus, ok := wf.StepStatus(ctx, step.Name)
	if !ok {
		return fmt.Errorf("step status %s not found", step.Name)
	}
	outputs := map[string]string{lastSubStepKey: subStep}
	for k, v := range extra {
		outputs[k] = v
	}
	if err := stepStatus.UpdateOutput(ctx, outputs); err != nil {
		return err
	}
	return wfCli.Status().Update(ctx, wf)
}

// runSubSteps executes sub-steps starting from current until one returns an error
// (wait/retry), marks done, or fails. Returning nil from Handle only happens when
// a sub-step reports done — intermediate advances must not return nil directly.
func runSubSteps(
	ctx context.Context,
	wfCli client.Client,
	wf *cloudv1.Workflow,
	step cloudv1.WorkflowStep,
	current string,
	run func(ctx context.Context, subStep string) (subStepResult, error),
) error {
	logger := log.FromContext(ctx)
	for {
		logger.Info("handling AZ sub-step", "subStep", current)
		result, err := run(ctx, current)
		if err != nil {
			return err
		}
		if result.done {
			return nil
		}
		if result.next == "" {
			return cloudv1.NewWFNonRetriableError(fmt.Errorf("sub-step %s returned empty next", current))
		}
		if err := updateSubStep(ctx, wfCli, wf, step, result.next, result.extra); err != nil {
			return err
		}
		current = result.next
	}
}

func getCluster(ctx context.Context, cli client.Client, namespace, name string) (*cloudv1.Cluster, error) {
	cluster := &cloudv1.Cluster{}
	if err := cli.Get(ctx, client.ObjectKey{Namespace: namespace, Name: name}, cluster); err != nil {
		if apierrors.IsNotFound(err) {
			return nil, cloudv1.NewWFNonRetriableError(fmt.Errorf("cluster %s not found", name))
		}
		return nil, fmt.Errorf("get cluster: %w", err)
	}
	return cluster, nil
}