package main

import (
	"context"
	"embed"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"

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

//
// ─── EMBEDDED WEB UI ────────────────────────────────────────────────────────────
//

//go:embed web/*
var webFS embed.FS

//
// ─── CONFIG ─────────────────────────────────────────────────────────────────────
//

var (
	CPUOverloadThreshold     int64
	MemoryOverloadThreshold  int64
	CPUUnderloadThreshold    int64
	MemoryUnderloadThreshold int64
)

func getEnvInt64(name string, defaultVal int64) int64 {
	valStr := os.Getenv(name)
	if valStr == "" {
		return defaultVal
	}
	val, err := strconv.ParseInt(valStr, 10, 64)
	if err != nil {
		log.Printf("Invalid %s: %s, using default %d", name, valStr, defaultVal)
		return defaultVal
	}
	return val
}

func init() {
	CPUOverloadThreshold = getEnvInt64("CPU_OVERLOAD_THRESHOLD", 80)
	MemoryOverloadThreshold = getEnvInt64("MEMORY_OVERLOAD_THRESHOLD", 80)
	CPUUnderloadThreshold = getEnvInt64("CPU_UNDERLOAD_THRESHOLD", 50)
	MemoryUnderloadThreshold = getEnvInt64("MEMORY_UNDERLOAD_THRESHOLD", 50)
}

//
// ─── EVENT MODEL + BUS ──────────────────────────────────────────────────────────
//

type PodMoveEvent struct {
	Time      time.Time `json:"time"`
	Namespace string    `json:"namespace"`
	Pod       string    `json:"pod"`
	FromNode  string    `json:"fromNode"`
	ToNode    string    `json:"toNode"`
	Reason    string    `json:"reason"`
}

type EventBus struct {
	mu      sync.Mutex
	history []PodMoveEvent
	clients map[chan PodMoveEvent]struct{}
}

func NewEventBus() *EventBus {
	return &EventBus{
		history: make([]PodMoveEvent, 0, 100),
		clients: make(map[chan PodMoveEvent]struct{}),
	}
}

func (b *EventBus) Emit(evt PodMoveEvent) {
	b.mu.Lock()
	defer b.mu.Unlock()

	if len(b.history) == 100 {
		b.history = b.history[1:]
	}
	b.history = append(b.history, evt)

	for ch := range b.clients {
		select {
		case ch <- evt:
		default:
		}
	}
}

func (b *EventBus) History() []PodMoveEvent {
	b.mu.Lock()
	defer b.mu.Unlock()
	out := make([]PodMoveEvent, len(b.history))
	copy(out, b.history)
	return out
}

func (b *EventBus) Subscribe() chan PodMoveEvent {
	b.mu.Lock()
	defer b.mu.Unlock()
	ch := make(chan PodMoveEvent, 10)
	b.clients[ch] = struct{}{}
	return ch
}

func (b *EventBus) Unsubscribe(ch chan PodMoveEvent) {
	b.mu.Lock()
	defer b.mu.Unlock()
	delete(b.clients, ch)
	close(ch)
}

//
// ─── MAIN ───────────────────────────────────────────────────────────────────────
//

func main() {
	config, err := rest.InClusterConfig()
	if err != nil {
		log.Fatal(err)
	}

	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		log.Fatal(err)
	}

	metricsClient, err := metricsv.NewForConfig(config)
	if err != nil {
		log.Fatal(err)
	}

	eventBus := NewEventBus()
        go startWebUI(eventBus, clientset)
	for {
		balanceNodes(clientset, metricsClient, eventBus)
		time.Sleep(1 * time.Minute)
	}
}

//
// ─── WEB UI SERVER (SAME PROCESS) ───────────────────────────────────────────────
//
func startWebUI(bus *EventBus, clientset *kubernetes.Clientset) {
	mux := http.NewServeMux()

	mux.Handle("/", http.FileServer(http.FS(webFS)))

	mux.HandleFunc("/api/events", func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(bus.History())
	})

	mux.HandleFunc("/api/stream", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")

		flusher, ok := w.(http.Flusher)
		if !ok {
			http.Error(w, "streaming unsupported", 500)
			return
		}

		ch := bus.Subscribe()
		defer bus.Unsubscribe(ch)

		for {
			select {
			case evt := <-ch:
				data, _ := json.Marshal(evt)
				fmt.Fprintf(w, "data: %s\n\n", data)
				flusher.Flush()
			case <-r.Context().Done():
				return
			}
		}
	})

	log.Println("WebUI listening on :8080")
        mux.HandleFunc("/api/nodes", func(w http.ResponseWriter, r *http.Request) {
            nodes, err := clientset.CoreV1().Nodes().List(context.Background(), metav1.ListOptions{})
            if err != nil {
                http.Error(w, "failed to list nodes", http.StatusInternalServerError)
                return
            }
            names := []string{}
            for _, n := range nodes.Items {
                names = append(names, n.Name)
            }
            w.Header().Set("Content-Type", "application/json")
            _ = json.NewEncoder(w).Encode(names)
        })
	_ = http.ListenAndServe(":8080", mux)
}

//
// ─── BALANCER (MINIMALLY MODIFIED) ──────────────────────────────────────────────
//

func balanceNodes(clientset *kubernetes.Clientset, metricsClient *metricsv.Clientset, bus *EventBus) {
	ctx := context.Background()

	nodeMetricsList, err := metricsClient.MetricsV1beta1().NodeMetricses().List(ctx, metav1.ListOptions{})
	if err != nil {
		log.Println("Failed to get node metrics:", err)
		return
	}

	nodeUsage := make(map[string]map[string]int64)
	for _, m := range nodeMetricsList.Items {
		nodeUsage[m.Name] = map[string]int64{
			"cpu": m.Usage.Cpu().MilliValue(),
			"mem": m.Usage.Memory().Value() / (1024 * 1024),
		}
	}

	nodes, err := clientset.CoreV1().Nodes().List(ctx, metav1.ListOptions{})
	if err != nil {
		log.Println("Failed to list nodes:", err)
		return
	}

	nodeCapacity := make(map[string]map[string]int64)
	for _, n := range nodes.Items {
		nodeCapacity[n.Name] = map[string]int64{
			"cpu": n.Status.Capacity.Cpu().MilliValue(),
			"mem": n.Status.Capacity.Memory().Value() / (1024 * 1024),
		}
	}

	var overloaded, underloaded []string
	for node, usage := range nodeUsage {
		cpuPercent := usage["cpu"] * 100 / nodeCapacity[node]["cpu"]
		memPercent := usage["mem"] * 100 / nodeCapacity[node]["mem"]

		if cpuPercent > CPUOverloadThreshold || memPercent > MemoryOverloadThreshold {
			overloaded = append(overloaded, node)
		} else if cpuPercent < CPUUnderloadThreshold && memPercent < MemoryUnderloadThreshold {
			underloaded = append(underloaded, node)
		}
	}

	fmt.Println("Overloaded nodes:", overloaded)
	fmt.Println("Underloaded nodes:", underloaded)

	if len(overloaded) == 0 || len(underloaded) == 0 {
		return
	}

	for _, node := range overloaded {
		pods, err := clientset.CoreV1().Pods("").List(ctx, metav1.ListOptions{
			FieldSelector: fmt.Sprintf("spec.nodeName=%s", node),
		})
		if err != nil {
			log.Println("Failed to list pods for node", node, err)
			continue
		}

		for _, pod := range pods.Items {
			if pod.Namespace == "kube-system" || isDaemonSet(&pod) {
				continue
			}

			targetNode := pickTargetNode(underloaded, nodeUsage, nodeCapacity)
			if targetNode == "" {
				return
			}

			fmt.Printf(
				"Rescheduling pod %s/%s from %s to %s\n",
				pod.Namespace,
				pod.Name,
				node,
				targetNode,
			)

			// 🔴 EMIT EVENT FOR WEB UI
			bus.Emit(PodMoveEvent{
				Time:      time.Now(),
				Namespace: pod.Namespace,
				Pod:       pod.Name,
				FromNode:  node,
				ToNode:    targetNode,
				Reason:    "node-utilization-imbalance",
			})

			grace := int64(0)
			err = clientset.CoreV1().Pods(pod.Namespace).Delete(ctx, pod.Name, metav1.DeleteOptions{
				GracePeriodSeconds: &grace,
			})
			if err != nil {
				log.Println("Failed to delete pod", pod.Name, err)
				return
			}

			nodeUsage[node]["cpu"] -= estimatePodCPU(&pod)
			nodeUsage[node]["mem"] -= estimatePodMem(&pod)
			nodeUsage[targetNode]["cpu"] += estimatePodCPU(&pod)
			nodeUsage[targetNode]["mem"] += estimatePodMem(&pod)

			return
		}
	}
}

//
// ─── HELPERS ────────────────────────────────────────────────────────────────────
//

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

func pickTargetNode(nodes []string, usage map[string]map[string]int64, capacity map[string]map[string]int64) string {
	var bestNode string
	bestLoad := int64(1 << 62)
	for _, n := range nodes {
		cpuPercent := usage[n]["cpu"] * 100 / capacity[n]["cpu"]
		memPercent := usage[n]["mem"] * 100 / capacity[n]["mem"]
		load := cpuPercent + memPercent
		if load < bestLoad {
			bestLoad = load
			bestNode = n
		}
	}
	return bestNode
}

func estimatePodCPU(pod *corev1.Pod) int64 {
	var cpu int64
	for _, c := range pod.Spec.Containers {
		if q, ok := c.Resources.Requests[corev1.ResourceCPU]; ok {
			cpu += q.MilliValue()
		}
	}
	return cpu
}

func estimatePodMem(pod *corev1.Pod) int64 {
	var mem int64
	for _, c := range pod.Spec.Containers {
		if q, ok := c.Resources.Requests[corev1.ResourceMemory]; ok {
			mem += q.Value() / (1024 * 1024)
		}
	}
	return mem
}