package main

import (
	"context"
	"encoding/json"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	metrics "k8s.io/metrics/pkg/client/clientset/versioned"
	"gopkg.in/yaml.v3"
)

// ---------------------------
// Structs
// ---------------------------

type Node struct {
	Name        string  `yaml:"name" json:"name"`
	IP          string  `yaml:"ip" json:"ip"`
	Status      string  `yaml:"status" json:"status"`
	CPU         int     `yaml:"cpu" json:"cpu"`
	Memory      int     `yaml:"memory" json:"memory"`
	Role        string  `yaml:"role" json:"role"`
	Cluster     string  `yaml:"cluster" json:"cluster"`
	LastActive  string  `yaml:"last_active" json:"last_active"`
	Pods        int     `yaml:"pods,omitempty" json:"pods,omitempty"`
	Temperature float64 `yaml:"temperature,omitempty" json:"temperature,omitempty"`
}

type NodePool struct {
	Nodes []Node `yaml:"nodes" json:"nodes"`
}

// ---------------------------
// Helper functions
// ---------------------------

func mustIntEnv(name string, def int) int {
	if val := os.Getenv(name); val != "" {
		if i, err := strconv.Atoi(val); err == nil {
			return i
		}
	}
	return def
}

func loadPool(file string) (*NodePool, error) {
	data, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}
	var pool NodePool
	err = yaml.Unmarshal(data, &pool)
	if err != nil {
		return nil, err
	}
	return &pool, nil
}

func savePool(file string, pool *NodePool) error {
	data, err := yaml.Marshal(pool)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(file, data, 0644)
}

func nodeIP(n *v1.Node) string {
	for _, addr := range n.Status.Addresses {
		if addr.Type == v1.NodeInternalIP {
			return addr.Address
		}
	}
	return ""
}

func isControlPlane(n *v1.Node) bool {
	if _, ok := n.Labels["node.kubernetes.io/microk8s-controlplane"]; ok {
		return true
	}
	if _, ok := n.Labels["node-role.kubernetes.io/control-plane"]; ok {
		return true
	}
	return false
}

func runSSH(host, user, pass, cmd string) (string, error) {
	config := &ssh.ClientConfig{
		User:            user,
		Auth:            []ssh.AuthMethod{ssh.Password(pass)},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         10 * time.Second,
	}
	client, err := ssh.Dial("tcp", host+":22", config)
	if err != nil {
		return "", err
	}
	defer client.Close()
	session, err := client.NewSession()
	if err != nil {
		return "", err
	}
	defer session.Close()
	out, err := session.CombinedOutput(cmd)
	return string(out), err
}

// ---------------------------
// Node pool initialization
// ---------------------------

func initNodePool(clientset *kubernetes.Clientset, poolFile string) (*NodePool, error) {
	ctx := context.Background()
	pool, err := loadPool(poolFile)
	if err != nil {
		log.Println("Failed to load existing node pool, starting fresh:", err)
		pool = &NodePool{}
	}
	nodes, err := clientset.CoreV1().Nodes().List(ctx, metav1.ListOptions{})
	if err != nil {
		return nil, err
	}
	poolMap := map[string]*Node{}
	for i := range pool.Nodes {
		poolMap[pool.Nodes[i].Name] = &pool.Nodes[i]
	}
	updatedNodes := []Node{}
	for _, n := range nodes.Items {
		ip := nodeIP(&n)
		role := "worker"
		if isControlPlane(&n) {
			role = "microk8s-controlplane"
		}
		node := Node{
			Name:       n.Name,
			IP:         ip,
			Status:     "online",
			Role:       role,
			Cluster:    os.Getenv("CLUSTER_NAME"),
			CPU:        0,
			Memory:     0,
			LastActive: time.Now().Format(time.RFC3339),
		}
		if oldNode, ok := poolMap[n.Name]; ok {
			node.Status = oldNode.Status
			node.LastActive = oldNode.LastActive
			node.CPU = oldNode.CPU
			node.Memory = oldNode.Memory
			node.Cluster = oldNode.Cluster
			node.Role = oldNode.Role
		}
		updatedNodes = append(updatedNodes, node)
	}
	pool.Nodes = updatedNodes
	savePool(poolFile, pool)
	log.Printf("Initialized node pool with %d nodes", len(pool.Nodes))
	return pool, nil
}

// ---------------------------
// Node metrics
// ---------------------------

func updatePerNodeUtilization(cs *kubernetes.Clientset, ms *metrics.Clientset, pool *NodePool, sshUser, sshPass string) error {
	ctx := context.Background()
	nodes, err := cs.CoreV1().Nodes().List(ctx, metav1.ListOptions{})
	if err != nil {
		return err
	}
	metricsList, err := ms.MetricsV1beta1().NodeMetricses().List(ctx, metav1.ListOptions{})
	if err != nil {
		return err
	}
	capCPU := map[string]int64{}
	capMem := map[string]int64{}
	nodeMap := map[string]*v1.Node{}
	for i, n := range nodes.Items {
		capCPU[n.Name] = n.Status.Capacity.Cpu().MilliValue()
		capMem[n.Name] = n.Status.Capacity.Memory().Value()
		nodeMap[n.Name] = &nodes.Items[i]
	}
	usageCPU := map[string]int64{}
	usageMem := map[string]int64{}
	for _, m := range metricsList.Items {
		usageCPU[m.Name] = m.Usage.Cpu().MilliValue()
		usageMem[m.Name] = m.Usage.Memory().Value()
	}
	pods, err := cs.CoreV1().Pods("").List(ctx, metav1.ListOptions{})
	if err != nil {
		return err
	}
	podCount := map[string]int{}
	for _, p := range pods.Items {
		if p.Spec.NodeName != "" {
			podCount[p.Spec.NodeName]++
		}
	}
	for i := range pool.Nodes {
		pn := &pool.Nodes[i]
		name := pn.Name
		if cpuCap, ok := capCPU[name]; ok {
			if memCap, ok2 := capMem[name]; ok2 {
				if cpuUse, ok3 := usageCPU[name]; ok3 {
					if memUse, ok4 := usageMem[name]; ok4 {
						if cpuCap > 0 && memCap > 0 {
							pn.CPU = int((cpuUse * 100) / cpuCap)
							pn.Memory = int((memUse * 100) / memCap)
						}
					}
				}
			}
		}
		pn.Pods = podCount[name]
		if kn, ok := nodeMap[name]; ok {
			pn.Temperature = getNodeTemp(kn, sshUser, sshPass)
		}
	}
	return nil
}

func getNodeTemp(node *v1.Node, sshUser, sshPass string) float64 {
	ip := nodeIP(node)
	out, err := runSSH(ip, sshUser, sshPass, "cat /proc/device-tree/model")
	if err != nil {
		return 0
	}
	hw := strings.ToUpper(strings.TrimSpace(out))
	var cmd string
	switch {
	case strings.Contains(hw, "RASPBERRY"):
		cmd = "vcgencmd measure_temp | egrep -o '[0-9]+\\.[0-9]+'"
	case strings.Contains(hw, "ODROID"):
		cmd = "awk '{printf \"%3.1f\", $1/1000}' /sys/class/thermal/thermal_zone0/temp"
	default:
		return 0
	}
	tempStr, err := runSSH(ip, sshUser, sshPass, cmd)
	if err != nil {
		return 0
	}
	t, err := strconv.ParseFloat(strings.TrimSpace(tempStr), 64)
	if err != nil {
		return 0
	}
	return t
}

// ---------------------------
// Cluster utilization
// ---------------------------

func clusterUtilization(cs *kubernetes.Clientset, ms *metrics.Clientset) (cpuPct int, memPct int, err error) {
	ctx := context.Background()
	nodes, err := cs.CoreV1().Nodes().List(ctx, metav1.ListOptions{})
	if err != nil {
		return 0, 0, err
	}
	metricsList, err := ms.MetricsV1beta1().NodeMetricses().List(ctx, metav1.ListOptions{})
	if err != nil {
		return 0, 0, err
	}
	var totalCPUCap, totalMemCap, totalCPUUse, totalMemUse int64
	for _, n := range nodes.Items {
		totalCPUCap += n.Status.Capacity.Cpu().MilliValue()
		totalMemCap += n.Status.Capacity.Memory().Value()
	}
	for _, m := range metricsList.Items {
		totalCPUUse += m.Usage.Cpu().MilliValue()
		totalMemUse += m.Usage.Memory().Value()
	}
	cpuPct = int((totalCPUUse * 100) / totalCPUCap)
	memPct = int((totalMemUse * 100) / totalMemCap)
	return
}

// ---------------------------
// Node scaling
// ---------------------------

func ensureControlPlanes(cs *kubernetes.Clientset, pool *NodePool, poolFile, sshUser, sshPass, clusterName string, desired int) {
	ctx := context.Background()
	nodes, err := cs.CoreV1().Nodes().List(ctx, metav1.ListOptions{})
	if err != nil {
		log.Println("failed to list nodes:", err)
		return
	}
	cpCount := 0
	var seedIP string
	for _, n := range nodes.Items {
		if isControlPlane(&n) {
			cpCount++
			if seedIP == "" {
				seedIP = nodeIP(&n)
			}
		}
	}
	if cpCount >= desired {
		return
	}
	if seedIP == "" {
		log.Println("no available control-plane node found in the node-pool")
		return
	}
	for i := range pool.Nodes {
		n := &pool.Nodes[i]
		if n.Status != "offline" {
			continue
		}
		log.Printf("Attempting to activate node %s as control-plane", n.Name)
		out, err := runSSH(seedIP, sshUser, sshPass, "microk8s add-node")
		if err != nil {
			log.Println("add-node failed:", err)
			return
		}
		var joinCmd string
		for _, line := range strings.Split(out, "\n") {
			if strings.HasPrefix(strings.TrimSpace(line), "microk8s join") {
				joinCmd = strings.TrimSpace(line)
				break
			}
		}
		if joinCmd == "" {
			log.Println("no join command found")
			return
		}
		_, err = runSSH(n.IP, sshUser, sshPass, joinCmd)
		if err != nil {
			log.Println("join failed:", err)
			return
		}
		n.Status = "online"
		n.Cluster = clusterName
		n.Role = "microk8s-controlplane"
		n.LastActive = time.Now().Format(time.RFC3339)
		log.Printf("Control-plane %s activated (%d/%d)", n.Name, cpCount+1, desired)
		return
	}
}

// ---------------------------
// Worker activation / deactivation
// ---------------------------

func activateOneWorker(cs *kubernetes.Clientset, pool *NodePool, poolFile, sshUser, sshPass, clusterName string) {
	ctx := context.Background()
	var workerNode *Node
	for i := range pool.Nodes {
		if pool.Nodes[i].Status == "offline" {
			workerNode = &pool.Nodes[i]
			break
		}
	}
	if workerNode == nil {
		log.Println("No offline nodes available — skipping activation")
		return
	}
	nodes, err := cs.CoreV1().Nodes().List(ctx, metav1.ListOptions{})
	if err != nil {
		log.Println("Failed to list nodes:", err)
		return
	}
	var seedIP string
	for _, n := range nodes.Items {
		if isControlPlane(&n) {
			seedIP = nodeIP(&n)
			break
		}
	}
	if seedIP == "" {
		log.Println("No control-plane node available to generate join command")
		return
	}
	out, err := runSSH(seedIP, sshUser, sshPass, "microk8s add-node")
	if err != nil {
		log.Println("add-node failed:", err)
		return
	}
	var joinCmd string
	for _, line := range strings.Split(out, "\n") {
		if strings.HasPrefix(strings.TrimSpace(line), "microk8s join") {
			joinCmd = strings.TrimSpace(line)
			break
		}
	}
	if joinCmd == "" {
		log.Println("No join command found from seed control-plane")
		return
	}
	joinCmd += " --worker"
	_, err = runSSH(workerNode.IP, sshUser, sshPass, joinCmd)
	if err != nil {
		log.Println("Worker join failed:", err)
		return
	}
	workerNode.Status = "online"
	workerNode.Cluster = clusterName
	workerNode.Role = "worker"
	workerNode.LastActive = time.Now().Format(time.RFC3339)
	log.Printf("Node %s successfully activated as worker node", workerNode.Name)
}

func deactivateOneWorkerSafe(cs *kubernetes.Clientset, pool *NodePool, poolFile, sshUser, sshPass, clusterName string, waitSec int) {
	ctx := context.Background()
	nodes, err := cs.CoreV1().Nodes().List(ctx, metav1.ListOptions{})
	if err != nil {
		log.Println("failed to list nodes:", err)
		return
	}
	var workerNode *v1.Node
	for _, n := range nodes.Items {
		if !isControlPlane(&n) && !n.Spec.Unschedulable {
			workerNode = &n
			break
		}
	}
	if workerNode == nil {
		log.Println("No worker nodes available for deactivation")
		return
	}
	log.Printf("Deactivating node: %s", workerNode.Name)
	ip := nodeIP(workerNode)
	// cordon
	patch := []byte(`{"spec":{"unschedulable":true}}`)
	_, err = cs.CoreV1().Nodes().Patch(ctx, workerNode.Name, metav1.StrategicMergePatchType, patch, metav1.PatchOptions{})
	if err != nil {
		log.Println("Failed to cordon node:", err)
		return
	}
	// drain
	err = drainNode(cs, workerNode.Name)
	if err != nil {
		log.Println("Drain failed, uncordoning node")
		patch = []byte(`{"spec":{"unschedulable":false}}`)
		_, _ = cs.CoreV1().Nodes().Patch(ctx, workerNode.Name, metav1.StrategicMergePatchType, patch, metav1.PatchOptions{})
		return
	}
	// leave
	_, err = runSSH(ip, sshUser, sshPass, "microk8s leave")
	if err != nil {
		log.Println("microk8s leave failed, keeping node object in cluster")
		return
	}
	// delete node
	_ = cs.CoreV1().Nodes().Delete(ctx, workerNode.Name, metav1.DeleteOptions{})
	// update pool
	for i := range pool.Nodes {
		if pool.Nodes[i].Name == workerNode.Name {
			pool.Nodes[i].Status = "offline"
			pool.Nodes[i].Cluster = "none"
			pool.Nodes[i].Role = "none"
			pool.Nodes[i].CPU = 0
			pool.Nodes[i].Memory = 0
			pool.Nodes[i].LastActive = time.Now().Format(time.RFC3339)
		}
	}
	time.Sleep(time.Duration(waitSec) * time.Second)
}

func drainNode(cs *kubernetes.Clientset, nodeName string) error {
	ctx := context.Background()
	pods, err := cs.CoreV1().Pods("").List(ctx, metav1.ListOptions{
		FieldSelector: "spec.nodeName=" + nodeName,
	})
	if err != nil {
		return err
	}
	for _, pod := range pods.Items {
		if pod.Namespace == "kube-system" {
			continue
		}
		if _, ok := pod.ObjectMeta.Annotations["kubernetes.io/config.mirror"]; ok {
			continue
		}
		grace := int64(60)
		_ = cs.CoreV1().Pods(pod.Namespace).Delete(ctx, pod.Name, metav1.DeleteOptions{
			GracePeriodSeconds: &grace,
		})
	}
	for {
		remaining, _ := cs.CoreV1().Pods("").List(ctx, metav1.ListOptions{
			FieldSelector: "spec.nodeName=" + nodeName,
		})
		active := 0
		for _, p := range remaining.Items {
			if p.Namespace != "kube-system" && p.DeletionTimestamp == nil {
				active++
			}
		}
		if active == 0 {
			break
		}
		time.Sleep(5 * time.Second)
	}
	return nil
}

// ---------------------------
// Web GUI
// ---------------------------

func startWebGUI(poolFile string) {
	http.HandleFunc("/status", func(w http.ResponseWriter, r *http.Request) {
		pool, err := loadPool(poolFile)
		if err != nil {
			log.Println("Cannot load node pool:", err)
			http.Error(w, "cannot load node pool", http.StatusInternalServerError)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(pool)
	})

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/status", http.StatusFound)
	})

	go func() {
		log.Println("Starting web GUI on :8080")
		if err := http.ListenAndServe(":8080", nil); err != nil {
			log.Fatal(err)
		}
	}()
}

// ---------------------------
// Main
// ---------------------------

func main() {
	poolFile := mustEnv("NODE_POOL_FILE", "nodepool.yaml")
	sshUser := mustEnv("SSH_USER", "pi")
	sshPass := mustEnv("SSH_PASS", "raspberry")
	clusterName := mustEnv("CLUSTER_NAME", "cluster1")
	cpDesired := mustIntEnv("CONTROL_PLANES", 1)
	scaleInterval := time.Duration(mustIntEnv("SCALE_INTERVAL_SEC", 30)) * time.Second
	workerUpThreshold := mustIntEnv("WORKER_CPU_UP", 70)
	workerDownThreshold := mustIntEnv("WORKER_CPU_DOWN", 30)
	waitBeforeDown := mustIntEnv("WORKER_DOWN_WAIT_SEC", 30)

	// In-cluster Kubernetes client
	config, err := rest.InClusterConfig()
	if err != nil {
		log.Fatal("Failed to get in-cluster config:", err)
	}
	cs, err := kubernetes.NewForConfig(config)
	if err != nil {
		log.Fatal(err)
	}
	ms, err := metrics.NewForConfig(config)
	if err != nil {
		log.Fatal(err)
	}

	// Initialize node pool
	pool, err := initNodePool(cs, poolFile)
	if err != nil {
		log.Fatal(err)
	}

	// Start web GUI
	startWebGUI(poolFile)

	// ---------------------------
	// Main loop
	// ---------------------------
	for {
		// 1. Ensure enough control-planes
		ensureControlPlanes(cs, pool, poolFile, sshUser, sshPass, clusterName, cpDesired)

		// 2. Update per-node utilization
		err := updatePerNodeUtilization(cs, ms, pool, sshUser, sshPass)
		if err != nil {
			log.Println("Failed to update node utilization:", err)
		}

		// 3. Check cluster CPU usage
		cpuPct, _, err := clusterUtilization(cs, ms)
		if err != nil {
			log.Println("Failed to get cluster utilization:", err)
			cpuPct = 0
		}

		// 4. Scale workers up or down
		if cpuPct > workerUpThreshold {
			log.Println("CPU high, activating a worker node...")
			activateOneWorker(cs, pool, poolFile, sshUser, sshPass, clusterName)
		} else if cpuPct < workerDownThreshold {
			log.Println("CPU low, deactivating a worker node...")
			deactivateOneWorkerSafe(cs, pool, poolFile, sshUser, sshPass, clusterName, waitBeforeDown)
		}

		// 5. Save pool state
		if err := savePool(poolFile, pool); err != nil {
			log.Println("Failed to save node pool:", err)
		}

		time.Sleep(scaleInterval)
	}
}

func mustEnv(name, def string) string {
	if val := os.Getenv(name); val != "" {
		return val
	}
	return def
}