package gkeheterogeneousstatus

import (
	"sort"
	"strings"
	"time"
)

const (
	ConfigMapNamePrefix = "tidb-cluster-gke-heterogeneous-status-"
	ConfigMapDataKey    = "current-gke-heterogeneous-status"

	ProbeStatusSuccess       = "success"
	ProbeStatusFailed        = "failed"
	ProbeStatusDisabled      = "disabled"
	ProbeStatusDisabling     = "disabling"
	ProbeStatusInvalidIntent = "invalid_intent"

	PlacementStatusHeterogeneous = "heterogeneous"
	PlacementStatusBase          = "base"
	PlacementStatusMismatch      = "mismatch"
	PlacementStatusPending       = "pending"
	PlacementStatusUnknown       = "unknown"

	ScaleWorkflowOperationScaleOut = "scale_out"
	ScaleWorkflowOperationDisable  = "disable"
)

type Snapshot struct {
	ClusterID         string `json:"cluster_id" yaml:"clusterId"`
	Namespace         string `json:"namespace" yaml:"namespace"`
	TenantID          string `json:"tenant_id" yaml:"tenantId"`
	ProjectID         string `json:"project_id" yaml:"projectId"`
	ClusterGeneration int64  `json:"cluster_generation" yaml:"clusterGeneration"`
	// ClusterFullSynced records whether the Cluster was FullSynced when this
	// snapshot was last updated.
	ClusterFullSynced bool `json:"cluster_full_synced" yaml:"clusterFullSynced"`
	// LastPlacementProbeTime records the latest full placement probe attempt.
	// Successful and failed placement probes update it; throttled metadata
	// updates, disabled snapshots, and invalid intent snapshots do not.
	LastPlacementProbeTime *time.Time `json:"last_placement_probe_time,omitempty" yaml:"lastPlacementProbeTime,omitempty"`
	// LastSuccessfulPlacementProbeTime records the latest successful placement
	// probe. Failed placement probes keep the previous successful value.
	LastSuccessfulPlacementProbeTime *time.Time           `json:"last_successful_placement_probe_time,omitempty" yaml:"lastSuccessfulPlacementProbeTime,omitempty"`
	ProbeStatus                      string               `json:"probe_status" yaml:"probeStatus"`
	ProbeError                       string               `json:"probe_error" yaml:"probeError"`
	Enabled                          bool                 `json:"enabled" yaml:"enabled"`
	ValidIntent                      bool                 `json:"valid_intent" yaml:"validIntent"`
	ScaleWorkflow                    *ScaleWorkflowStatus `json:"scale_workflow,omitempty" yaml:"scaleWorkflow,omitempty"`
	Components                       []ComponentStatus    `json:"components,omitempty" yaml:"components,omitempty"`
	ClusterNodeUsages                []NodeUsage          `json:"cluster_node_usages,omitempty" yaml:"clusterNodeUsages,omitempty"`
	ClusterNodeSummary               *NodeSummary         `json:"cluster_node_summary,omitempty" yaml:"clusterNodeSummary,omitempty"`
}

type ScaleWorkflowStatus struct {
	WorkflowNamespace string     `json:"workflow_namespace,omitempty" yaml:"workflowNamespace,omitempty"`
	WorkflowName      string     `json:"workflow_name,omitempty" yaml:"workflowName,omitempty"`
	WorkflowUID       string     `json:"workflow_uid,omitempty" yaml:"workflowUID,omitempty"`
	WorkflowPhase     string     `json:"workflow_phase,omitempty" yaml:"workflowPhase,omitempty"`
	Operation         string     `json:"operation,omitempty" yaml:"operation,omitempty"`
	ComponentType     string     `json:"component_type,omitempty" yaml:"componentType,omitempty"`
	TiDBGroupID       string     `json:"tidb_group_id,omitempty" yaml:"tidbGroupID,omitempty"`
	BeginAt           *time.Time `json:"begin_at,omitempty" yaml:"beginAt,omitempty"`
	EndAt             *time.Time `json:"end_at,omitempty" yaml:"endAt,omitempty"`
	DurationSeconds   *float64   `json:"duration_seconds,omitempty" yaml:"durationSeconds,omitempty"`
	InProgress        bool       `json:"in_progress,omitempty" yaml:"inProgress,omitempty"`
	Stale             bool       `json:"stale,omitempty" yaml:"stale,omitempty"`
}

type ComponentStatus struct {
	ComponentType     string     `json:"component_type,omitempty" yaml:"componentType,omitempty"`
	TiDBGroupID       string     `json:"tidb_group_id,omitempty" yaml:"tidbGroupId,omitempty"`
	ComputeClassName  string     `json:"compute_class_name,omitempty" yaml:"computeClassName,omitempty"`
	BaseInstanceType  string     `json:"base_instance_type,omitempty" yaml:"baseInstanceType,omitempty"`
	DesiredReplicas   int32      `json:"desired_replicas,omitempty" yaml:"desiredReplicas,omitempty"`
	TotalPods         int32      `json:"total_pods,omitempty" yaml:"totalPods,omitempty"`
	ScheduledPods     int32      `json:"scheduled_pods,omitempty" yaml:"scheduledPods,omitempty"`
	HeterogeneousPods int32      `json:"heterogeneous_pods,omitempty" yaml:"heterogeneousPods,omitempty"`
	BasePods          int32      `json:"base_pods,omitempty" yaml:"basePods,omitempty"`
	MismatchPods      int32      `json:"mismatch_pods,omitempty" yaml:"mismatchPods,omitempty"`
	PendingPods       int32      `json:"pending_pods,omitempty" yaml:"pendingPods,omitempty"`
	UnknownPods       int32      `json:"unknown_pods,omitempty" yaml:"unknownPods,omitempty"`
	PlacementRatio    *float64   `json:"placement_ratio,omitempty" yaml:"placementRatio,omitempty"`
	EffectiveRatio    *float64   `json:"effective_ratio,omitempty" yaml:"effectiveRatio,omitempty"`
	PodGroups         []PodGroup `json:"pod_groups,omitempty" yaml:"podGroups,omitempty"`
	// Pods stores full placement details only before CompactSnapshot.
	Pods        []PodPlacement `json:"-" yaml:"-"`
	NodeUsages  []NodeUsage    `json:"node_usages,omitempty" yaml:"nodeUsages,omitempty"`
	NodeSummary *NodeSummary   `json:"node_summary,omitempty" yaml:"nodeSummary,omitempty"`
}

type PodGroup struct {
	PlacementStatus      string `json:"placement_status,omitempty" yaml:"placementStatus,omitempty"`
	ExpectedRole         string `json:"expected_role,omitempty" yaml:"expectedRole,omitempty"`
	ExpectedInstanceType string `json:"expected_instance_type,omitempty" yaml:"expectedInstanceType,omitempty"`
	ExpectedComputeClass string `json:"expected_compute_class,omitempty" yaml:"expectedComputeClass,omitempty"`
	ActualNodePool       string `json:"actual_nodepool,omitempty" yaml:"actualNodePool,omitempty"`
	ActualInstanceType   string `json:"actual_instance_type,omitempty" yaml:"actualInstanceType,omitempty"`
	ActualComputeClass   string `json:"actual_compute_class,omitempty" yaml:"actualComputeClass,omitempty"`
	PodCount             int32  `json:"pod_count,omitempty" yaml:"podCount,omitempty"`
	// Pods stores only pod names for debugging. PodCount is the exact total.
	Pods []string `json:"pods,omitempty" yaml:"pods,omitempty"`
}

type PodPlacement struct {
	Pod                  string `json:"pod,omitempty" yaml:"pod,omitempty"`
	Node                 string `json:"node,omitempty" yaml:"node,omitempty"`
	NodePool             string `json:"nodepool,omitempty" yaml:"nodepool,omitempty"`
	Zone                 string `json:"zone,omitempty" yaml:"zone,omitempty"`
	ActualInstanceType   string `json:"actual_instance_type,omitempty" yaml:"actualInstanceType,omitempty"`
	ActualComputeClass   string `json:"actual_compute_class,omitempty" yaml:"actualComputeClass,omitempty"`
	ExpectedRole         string `json:"expected_role,omitempty" yaml:"expectedRole,omitempty"`
	ExpectedInstanceType string `json:"expected_instance_type,omitempty" yaml:"expectedInstanceType,omitempty"`
	ExpectedComputeClass string `json:"expected_compute_class,omitempty" yaml:"expectedComputeClass,omitempty"`
	IsHeterogeneous      bool   `json:"is_heterogeneous,omitempty" yaml:"isHeterogeneous,omitempty"`
	PlacementStatus      string `json:"placement_status,omitempty" yaml:"placementStatus,omitempty"`
}

type NodeUsage struct {
	NodePool        string `json:"nodepool,omitempty" yaml:"nodepool,omitempty"`
	MachineType     string `json:"machine_type,omitempty" yaml:"machineType,omitempty"`
	MachineFamily   string `json:"machine_family,omitempty" yaml:"machineFamily,omitempty"`
	ExpectedRole    string `json:"expected_role,omitempty" yaml:"expectedRole,omitempty"`
	PlacementStatus string `json:"placement_status,omitempty" yaml:"placementStatus,omitempty"`
	UsedNodeCount   int32  `json:"used_node_count,omitempty" yaml:"usedNodeCount,omitempty"`
	PodCount        int32  `json:"pod_count,omitempty" yaml:"podCount,omitempty"`
}

type NodeSummary struct {
	TotalUsedNodeCount         int32    `json:"total_used_node_count,omitempty" yaml:"totalUsedNodeCount,omitempty"`
	HeterogeneousUsedNodeCount int32    `json:"heterogeneous_used_node_count,omitempty" yaml:"heterogeneousUsedNodeCount,omitempty"`
	BaseUsedNodeCount          int32    `json:"base_used_node_count,omitempty" yaml:"baseUsedNodeCount,omitempty"`
	MismatchUsedNodeCount      int32    `json:"mismatch_used_node_count,omitempty" yaml:"mismatchUsedNodeCount,omitempty"`
	UnknownUsedNodeCount       int32    `json:"unknown_used_node_count,omitempty" yaml:"unknownUsedNodeCount,omitempty"`
	HeterogeneousUsedNodeRatio *float64 `json:"heterogeneous_used_node_ratio,omitempty" yaml:"heterogeneousUsedNodeRatio,omitempty"`
}

type nodeUsageAgg struct {
	usage NodeUsage
	nodes map[string]struct{}
}

const unknownNodeUsageLabel = "unknown"

func MachineFamilyFromType(machineType string) string {
	machineType = strings.TrimSpace(machineType)
	if machineType == "" {
		return unknownNodeUsageLabel
	}
	if i := strings.Index(machineType, "-"); i > 0 {
		return machineType[:i]
	}
	return machineType
}

// PopulateNodeUsages derives node usage from full pod placements before the
// snapshot is compacted. It intentionally does not inventory nodepools because
// nodepools can be shared across clusters and components.
func PopulateNodeUsages(components []ComponentStatus) ([]ComponentStatus, []NodeUsage) {
	out := make([]ComponentStatus, len(components))
	copy(out, components)
	for i := range out {
		out[i].NodeUsages = componentNodeUsages(out[i])
		out[i].NodeSummary = NodeSummaryFromComponents([]ComponentStatus{out[i]})
	}
	return out, clusterNodeUsages(out)
}

func NodeSummaryFromComponents(components []ComponentStatus) *NodeSummary {
	nodes := map[string]string{}
	for _, component := range components {
		for _, pod := range component.Pods {
			node := strings.TrimSpace(pod.Node)
			if node == "" {
				continue
			}
			status := nodeUsageLabel(pod.PlacementStatus)
			current, ok := nodes[node]
			if !ok || nodeSummaryStatusRank(status) < nodeSummaryStatusRank(current) {
				nodes[node] = status
			}
		}
	}
	if len(nodes) == 0 {
		return nil
	}

	summary := &NodeSummary{TotalUsedNodeCount: int32(len(nodes))}
	for _, placementStatus := range nodes {
		switch placementStatus {
		case PlacementStatusHeterogeneous:
			summary.HeterogeneousUsedNodeCount++
		case PlacementStatusBase:
			summary.BaseUsedNodeCount++
		case PlacementStatusMismatch:
			summary.MismatchUsedNodeCount++
		default:
			summary.UnknownUsedNodeCount++
		}
	}
	ratio := float64(summary.HeterogeneousUsedNodeCount) / float64(summary.TotalUsedNodeCount)
	summary.HeterogeneousUsedNodeRatio = &ratio
	return summary
}

func componentNodeUsages(component ComponentStatus) []NodeUsage {
	aggs := map[string]*nodeUsageAgg{}
	for _, pod := range component.Pods {
		node := strings.TrimSpace(pod.Node)
		if node == "" {
			continue
		}
		nodePool := nodeUsageLabel(pod.NodePool)
		machineType := nodeUsageLabel(pod.ActualInstanceType)
		machineFamily := MachineFamilyFromType(machineType)
		expectedRole := nodeUsageLabel(pod.ExpectedRole)
		placementStatus := nodeUsageLabel(pod.PlacementStatus)
		key := strings.Join([]string{
			nodePool,
			machineType,
			machineFamily,
			expectedRole,
			placementStatus,
		}, "\x00")
		agg := aggs[key]
		if agg == nil {
			agg = &nodeUsageAgg{
				usage: NodeUsage{
					NodePool:        nodePool,
					MachineType:     machineType,
					MachineFamily:   machineFamily,
					ExpectedRole:    expectedRole,
					PlacementStatus: placementStatus,
				},
				nodes: map[string]struct{}{},
			}
			aggs[key] = agg
		}
		agg.nodes[node] = struct{}{}
		agg.usage.PodCount++
	}
	return sortedNodeUsages(aggs)
}

func CompactSnapshot(snapshot Snapshot) Snapshot {
	for i := range snapshot.Components {
		component := &snapshot.Components[i]
		if len(component.Pods) > 0 {
			component.PodGroups = PodGroupsFromPlacements(component.Pods)
		}
		component.Pods = nil
	}
	return snapshot
}

func PodGroupsFromPlacements(placements []PodPlacement) []PodGroup {
	type key struct {
		placementStatus      string
		expectedRole         string
		expectedInstanceType string
		expectedComputeClass string
		actualNodePool       string
		actualInstanceType   string
		actualComputeClass   string
	}

	groups := map[key]*PodGroup{}
	for _, placement := range placements {
		k := key{
			placementStatus:      placement.PlacementStatus,
			expectedRole:         placement.ExpectedRole,
			expectedInstanceType: placement.ExpectedInstanceType,
			expectedComputeClass: placement.ExpectedComputeClass,
			actualNodePool:       placement.NodePool,
			actualInstanceType:   placement.ActualInstanceType,
			actualComputeClass:   placement.ActualComputeClass,
		}
		group := groups[k]
		if group == nil {
			group = &PodGroup{
				PlacementStatus:      placement.PlacementStatus,
				ExpectedRole:         placement.ExpectedRole,
				ExpectedInstanceType: placement.ExpectedInstanceType,
				ExpectedComputeClass: placement.ExpectedComputeClass,
				ActualNodePool:       placement.NodePool,
				ActualInstanceType:   placement.ActualInstanceType,
				ActualComputeClass:   placement.ActualComputeClass,
			}
			groups[k] = group
		}
		group.PodCount++
		group.Pods = append(group.Pods, placement.Pod)
	}

	out := make([]PodGroup, 0, len(groups))
	for _, group := range groups {
		sort.Strings(group.Pods)
		out = append(out, *group)
	}
	sort.Slice(out, func(i, j int) bool {
		return comparePodGroups(out[i], out[j]) < 0
	})
	return out
}

func comparePodGroups(a, b PodGroup) int {
	if rankA, rankB := placementStatusRank(a.PlacementStatus), placementStatusRank(b.PlacementStatus); rankA != rankB {
		return rankA - rankB
	}
	for _, pair := range [][2]string{
		{a.PlacementStatus, b.PlacementStatus},
		{a.ExpectedRole, b.ExpectedRole},
		{a.ExpectedInstanceType, b.ExpectedInstanceType},
		{a.ExpectedComputeClass, b.ExpectedComputeClass},
		{a.ActualNodePool, b.ActualNodePool},
		{a.ActualInstanceType, b.ActualInstanceType},
		{a.ActualComputeClass, b.ActualComputeClass},
	} {
		if pair[0] < pair[1] {
			return -1
		}
		if pair[0] > pair[1] {
			return 1
		}
	}
	return 0
}

func placementStatusRank(status string) int {
	switch status {
	case PlacementStatusPending:
		return 0
	case PlacementStatusMismatch:
		return 1
	case PlacementStatusUnknown:
		return 2
	case PlacementStatusHeterogeneous:
		return 3
	case PlacementStatusBase:
		return 4
	default:
		return 5
	}
}

func nodeSummaryStatusRank(status string) int {
	switch status {
	case PlacementStatusHeterogeneous:
		return 0
	case PlacementStatusMismatch:
		return 1
	case PlacementStatusUnknown:
		return 2
	case PlacementStatusBase:
		return 3
	default:
		return 4
	}
}

func clusterNodeUsages(components []ComponentStatus) []NodeUsage {
	aggs := map[string]*nodeUsageAgg{}
	for _, component := range components {
		for _, pod := range component.Pods {
			node := strings.TrimSpace(pod.Node)
			if node == "" {
				continue
			}
			nodePool := nodeUsageLabel(pod.NodePool)
			machineType := nodeUsageLabel(pod.ActualInstanceType)
			machineFamily := MachineFamilyFromType(machineType)
			key := strings.Join([]string{
				nodePool,
				machineType,
				machineFamily,
			}, "\x00")
			agg := aggs[key]
			if agg == nil {
				agg = &nodeUsageAgg{
					usage: NodeUsage{
						NodePool:      nodePool,
						MachineType:   machineType,
						MachineFamily: machineFamily,
					},
					nodes: map[string]struct{}{},
				}
				aggs[key] = agg
			}
			agg.nodes[node] = struct{}{}
			agg.usage.PodCount++
		}
	}
	return sortedNodeUsages(aggs)
}

func nodeUsageLabel(value string) string {
	value = strings.TrimSpace(value)
	if value == "" {
		return unknownNodeUsageLabel
	}
	return value
}

func sortedNodeUsages(aggs map[string]*nodeUsageAgg) []NodeUsage {
	out := make([]NodeUsage, 0, len(aggs))
	for _, agg := range aggs {
		usage := agg.usage
		usage.UsedNodeCount = int32(len(agg.nodes))
		out = append(out, usage)
	}
	sort.Slice(out, func(i, j int) bool {
		if out[i].NodePool != out[j].NodePool {
			return out[i].NodePool < out[j].NodePool
		}
		if out[i].MachineType != out[j].MachineType {
			return out[i].MachineType < out[j].MachineType
		}
		if out[i].ExpectedRole != out[j].ExpectedRole {
			return out[i].ExpectedRole < out[j].ExpectedRole
		}
		return out[i].PlacementStatus < out[j].PlacementStatus
	})
	return out
}