package main

import (
	"context"
	"fmt"
	"sort"
	"time"

	"github.com/dustin/go-humanize"
	corev1 "k8s.io/api/core/v1"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	metricsv "k8s.io/metrics/pkg/client/clientset/versioned"
)

type PodInfo struct {
	Name            string
	Namespace       string
	Status          string
	ReadyContainers int
	ContainerCount  int
	RestartCount    int
	CPUUsage        string
	MemoryUsage     string
	Age             time.Duration
	Node            string
	IsStarting      bool // Added to track if pod is starting or restarting
}

func refreshData(clientset *kubernetes.Clientset, metricsClient *metricsv.Clientset, namespaces []string) ([]PodInfo, error) {
	pods, err := clientset.CoreV1().Pods("").List(context.Background(), metav1.ListOptions{})
	if err != nil {
		return nil, fmt.Errorf("error listing pods: %v", err)
	}

	// Create a map for faster namespace lookups
	namespaceMap := make(map[string]bool)
	for _, ns := range namespaces {
		namespaceMap[ns] = true
	}

	podMetricsMap := make(map[string]map[string]struct {
		CPU    int64
		Memory int64
	})
	if metricsClient != nil {
		podMetrics, err := metricsClient.MetricsV1beta1().PodMetricses("").List(context.Background(), metav1.ListOptions{})
		if err == nil {
			for _, metric := range podMetrics.Items {
				nsKey := metric.Namespace
				// Only process metrics for specified namespaces
				if !namespaceMap[nsKey] {
					continue
				}

				if _, ok := podMetricsMap[nsKey]; !ok {
					podMetricsMap[nsKey] = make(map[string]struct {
						CPU    int64
						Memory int64
					})
				}
				var totalCPU, totalMemory int64
				for _, container := range metric.Containers {
					totalCPU += container.Usage.Cpu().MilliValue()
					totalMemory += container.Usage.Memory().Value()
				}
				podMetricsMap[nsKey][metric.Name] = struct {
					CPU    int64
					Memory int64
				}{
					CPU:    totalCPU,
					Memory: totalMemory,
				}
			}
		}
	}

	var result []PodInfo
	for _, pod := range pods.Items {
		readyCount := 0
		restarts := 0
		nsKey := pod.Namespace
		// Only process pods for specified namespaces
		if !namespaceMap[nsKey] {
			continue
		}
		for _, status := range pod.Status.ContainerStatuses {
			if status.Ready {
				readyCount++
			}
			restarts += int(status.RestartCount)
		}
		podInfo := PodInfo{
			Name:            pod.Name,
			Namespace:       pod.Namespace,
			Status:          string(pod.Status.Phase),
			ContainerCount:  len(pod.Spec.Containers),
			ReadyContainers: readyCount,
			RestartCount:    restarts,
			Age:             time.Since(pod.CreationTimestamp.Time),
			//IP:              pod.Status.PodIP,
			IsStarting: time.Since(pod.CreationTimestamp.Time) < 1*time.Minute || isPodStartingOrRestarting(pod),
			Node:       pod.Spec.NodeName,
		}
		if metrics, ok := podMetricsMap[pod.Namespace][pod.Name]; ok {
			podInfo.CPUUsage = fmt.Sprintf("%dm", metrics.CPU)
			podInfo.MemoryUsage = humanize.Bytes(uint64(metrics.Memory))
		} else {
			podInfo.CPUUsage = "N/A"
			podInfo.MemoryUsage = "N/A"
		}
		result = append(result, podInfo)
	}
	sort.Slice(result, func(i, j int) bool {
		if result[i].Namespace == result[j].Namespace {
			return result[i].Name < result[j].Name
		}
		return result[i].Namespace < result[j].Namespace
	})
	return result, nil
}

// Helper function to check if pod is starting or restarting
func isPodStartingOrRestarting(pod v1.Pod) bool {
	// Check if pod is in Pending state
	if pod.Status.Phase == corev1.PodPending {
		return true
	}

	// Check if any container is in waiting state with specific reasons
	for _, containerStatus := range pod.Status.ContainerStatuses {
		if containerStatus.State.Waiting != nil {
			reason := containerStatus.State.Waiting.Reason
			if reason == "ContainerCreating" ||
				reason == "PodInitializing" ||
				reason == "CrashLoopBackOff" {
				return true
			}
		}
		// Check if container recently restarted (within last minute)
		if containerStatus.LastTerminationState.Terminated != nil {
			terminatedTime := containerStatus.LastTerminationState.Terminated.FinishedAt.Time
			if time.Since(terminatedTime) < 60*time.Second {
				return true
			}
		}
	}
	return false
}