package workernodeusage

import (
	"context"
	"slices"
	"strconv"
	"strings"

	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	corev1 "k8s.io/api/core/v1"
)

const (
	TenantLabelKey       = "tags.tidbcloud.com/tenant"
	SharedPoolLabelKey   = "tags.tidbcloud.com/shared-pool"
	ResourcePoolLabelKey = "tags.tidbcloud.com/resource-pool"
	NodeSpecLabelKey     = "node-spec"
	DedicatedTaintKey    = "dedicated"
	WorkerTaintValue     = "worker"

	DropReasonConflictingOwnership = "conflicting_ownership"
	DropReasonMissingIdentity      = "missing_identity"
	DropReasonInvalidNodeSpec      = "invalid_node_spec"
)

type NodeUsage struct {
	NodeName        string
	TenantID        string
	SharedPoolID    string
	VCPU            uint64
	DropReason      string
	RelevantPoolIDs []string
	HasCPUMismatch  bool
}

type Snapshot struct {
	Nodes []NodeUsage
}

type Summary struct {
	VCPUsUsed    uint64
	CountedNodes uint64
	DroppedNodes uint64
	DropReasons  map[string]uint64
}

func BuildSnapshot(ctx context.Context, nodes map[string]*corev1.Node, pods []*corev1.Pod) Snapshot {
	podsByNode := make(map[string][]*corev1.Pod, len(pods))
	for _, pod := range pods {
		if pod == nil || pod.Spec.NodeName == "" || !IsActiveNonDaemonPod(pod) {
			continue
		}
		podsByNode[pod.Spec.NodeName] = append(podsByNode[pod.Spec.NodeName], pod)
	}

	nodeNames := make([]string, 0, len(nodes))
	for nodeName := range nodes {
		nodeNames = append(nodeNames, nodeName)
	}
	slices.Sort(nodeNames)

	result := Snapshot{Nodes: make([]NodeUsage, 0, len(nodeNames))}
	for _, nodeName := range nodeNames {
		result.Nodes = append(result.Nodes, buildNodeUsage(ctx, nodes[nodeName], podsByNode[nodeName]))
	}

	return result
}

func (s Snapshot) AggregateUsageBySharedPool() map[string]map[string]uint64 {
	aggregated := make(map[string]map[string]uint64)
	for _, node := range s.Nodes {
		if node.SharedPoolID == "" || node.TenantID == "" || node.VCPU == 0 {
			continue
		}
		key := node.TenantID + ":" + node.SharedPoolID
		if aggregated[key] == nil {
			aggregated[key] = map[string]uint64{"vcpu": 0}
		}
		aggregated[key]["vcpu"] += node.VCPU
	}
	return aggregated
}

func (s Snapshot) SummaryForSharedPool(sharedPoolID string) Summary {
	summary := Summary{
		DropReasons: map[string]uint64{
			DropReasonConflictingOwnership: 0,
			DropReasonMissingIdentity:      0,
			DropReasonInvalidNodeSpec:      0,
		},
	}

	for _, node := range s.Nodes {
		if node.SharedPoolID == sharedPoolID && node.VCPU > 0 {
			summary.VCPUsUsed += node.VCPU
			summary.CountedNodes++
			continue
		}

		if node.DropReason == "" || !slices.Contains(node.RelevantPoolIDs, sharedPoolID) {
			continue
		}
		summary.DroppedNodes++
		summary.DropReasons[node.DropReason]++
	}

	return summary
}

func IsWorkerNode(node *corev1.Node) bool {
	for _, taint := range node.Spec.Taints {
		if taint.Key == DedicatedTaintKey && taint.Value == WorkerTaintValue {
			return true
		}
	}
	return false
}

func PoolLabelValue(pod *corev1.Pod) string {
	if poolID := pod.Labels[ResourcePoolLabelKey]; poolID != "" {
		return poolID
	}
	return pod.Labels[SharedPoolLabelKey]
}

func IsBillablePod(pod *corev1.Pod) bool {
	tenantLabel := pod.Labels[TenantLabelKey]
	poolLabel := PoolLabelValue(pod)
	return tenantLabel != "" && poolLabel != "" && IsActiveNonDaemonPod(pod)
}

func IsActiveNonDaemonPod(pod *corev1.Pod) bool {
	return IsNonDaemonSetPod(pod) && !IsInactivePod(pod)
}

func IsNonDaemonSetPod(pod *corev1.Pod) bool {
	for _, owner := range pod.OwnerReferences {
		if owner.Kind == "DaemonSet" {
			return false
		}
	}
	return true
}

func IsInactivePod(pod *corev1.Pod) bool {
	return pod.Status.Phase == corev1.PodSucceeded || pod.Status.Phase == corev1.PodFailed
}

func buildNodeUsage(ctx context.Context, node *corev1.Node, pods []*corev1.Pod) NodeUsage {
	usage := NodeUsage{
		NodeName:        node.Name,
		RelevantPoolIDs: relevantPoolIDs(pods),
	}

	billablePods := make([]*corev1.Pod, 0, len(pods))
	for _, pod := range pods {
		if IsBillablePod(pod) {
			billablePods = append(billablePods, pod)
		}
	}

	ownership := determineNodeOwnership(ctx, node.Name, billablePods)
	switch {
	case ownership.conflicts:
		usage.DropReason = DropReasonConflictingOwnership
		return usage
	case ownership.tenantID == "" || ownership.sharedPoolID == "":
		usage.DropReason = DropReasonMissingIdentity
		return usage
	}

	vcpu, mismatch := ParseNodeSpecCPU(ctx, node)
	usage.HasCPUMismatch = mismatch
	if vcpu == 0 {
		usage.DropReason = DropReasonInvalidNodeSpec
		return usage
	}

	usage.TenantID = ownership.tenantID
	usage.SharedPoolID = ownership.sharedPoolID
	usage.VCPU = uint64(vcpu)
	if len(usage.RelevantPoolIDs) == 0 {
		usage.RelevantPoolIDs = []string{ownership.sharedPoolID}
	}

	return usage
}

type nodeOwnership struct {
	tenantID     string
	sharedPoolID string
	conflicts    bool
}

type podIdentity struct {
	tenantID string
	poolID   string
}

func determineNodeOwnership(ctx context.Context, nodeName string, pods []*corev1.Pod) nodeOwnership {
	var ownership nodeOwnership

	for _, pod := range pods {
		identity, valid := podIdentityFromPod(pod)
		if !valid {
			continue
		}

		if ownership.tenantID == "" {
			ownership.tenantID = identity.tenantID
			ownership.sharedPoolID = identity.poolID
			continue
		}
		if ownership.tenantID == identity.tenantID && ownership.sharedPoolID == identity.poolID {
			continue
		}

		tilog.Warn(ctx, "conflicting resource pool ownership detected on node",
			zap.String("node", nodeName),
			zap.String("existing_tenant", ownership.tenantID),
			zap.String("existing_pool", ownership.sharedPoolID),
			zap.String("conflicting_tenant", identity.tenantID),
			zap.String("conflicting_pool", identity.poolID),
			zap.String("conflicting_pod", pod.Name))
		ownership.conflicts = true
		return ownership
	}

	return ownership
}

func podIdentityFromPod(pod *corev1.Pod) (podIdentity, bool) {
	tenantID := pod.Labels[TenantLabelKey]
	poolID := PoolLabelValue(pod)
	if tenantID == "" || poolID == "" {
		return podIdentity{}, false
	}
	return podIdentity{tenantID: tenantID, poolID: poolID}, true
}

func ParseNodeSpecCPU(ctx context.Context, node *corev1.Node) (int, bool) {
	nodeSpec := node.Labels[NodeSpecLabelKey]
	if nodeSpec == "" {
		tilog.Warn(ctx, "node missing node-spec label",
			zap.String("node", node.Name))
		return 0, false
	}

	idx := strings.Index(strings.ToUpper(nodeSpec), "C")
	if idx <= 0 {
		tilog.Warn(ctx, "invalid node-spec format",
			zap.String("node", node.Name),
			zap.String("node-spec", nodeSpec))
		return 0, false
	}

	vcpuStr := nodeSpec[:idx]
	vcpu, err := strconv.Atoi(vcpuStr)
	if err != nil {
		tilog.Warn(ctx, "failed to parse vCPU from node-spec",
			zap.String("node", node.Name),
			zap.String("node-spec", nodeSpec),
			zap.Error(err))
		return 0, false
	}

	mismatch := false
	capacityCPU := node.Status.Capacity.Cpu()
	if capacityCPU != nil {
		capacityCPUInt := int(capacityCPU.Value())
		if capacityCPUInt != vcpu {
			mismatch = true
			tilog.Warn(ctx, "node-spec vCPU does not match status.capacity.cpu",
				zap.String("node", node.Name),
				zap.String("node-spec", nodeSpec),
				zap.Int("node-spec-vcpu", vcpu),
				zap.Int("capacity-cpu", capacityCPUInt))
		}
	}

	return vcpu, mismatch
}

func relevantPoolIDs(pods []*corev1.Pod) []string {
	seen := make(map[string]struct{}, len(pods))
	values := make([]string, 0, len(pods))
	for _, pod := range pods {
		value := strings.TrimSpace(PoolLabelValue(pod))
		if value == "" {
			continue
		}
		if _, ok := seen[value]; ok {
			continue
		}
		seen[value] = struct{}{}
		values = append(values, value)
	}
	slices.Sort(values)
	return values
}