package main

import (
	"context"
	"fmt"
	"iter"
	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"log"
	"maps"
	"os"
	"slices"
	"strconv"
)

type ValidationLevel int

const (
	Info ValidationLevel = iota
	Warning
	Error
)

func LogValidationMsg(level ValidationLevel, msg string, v ...any) {
	fmt.Fprintf(os.Stderr, "NOTICE: "+msg+"\n", v...)
}

func IterToSlice[K any](i iter.Seq[K]) []K {
	res := make([]K, 0)
	for v := range i {
		res = append(res, v)
	}
	return res
}

func MapKeys[K comparable, V any](m map[K]V) []K {
	return IterToSlice(maps.Keys(m))
}

func validate(files []string, options *Options) error {
	clientset, _ := GetKubernetesConnection()
	config, err := readConfig(files)
	if err != nil {
		return err
	}
	config.Infer(func(application *Application) []string {
		pods := FindPods(application, clientset)
		var res []string
		for _, pod := range pods {
			if !slices.Contains(res, pod.Spec.ServiceAccountName) {
				res = append(res, pod.Spec.ServiceAccountName)
			}
		}
		log.Printf("Inferred service accounts: %s/%s: %v", application.Namespace.Name, application.Name,
			res)
		return res
	})

	// map  applname1 -> appname2 where appname1 is in an open namespace and app2 is in a closed namespace.
	openToClosedAccess := make(map[string]string)

	applicationPods := make(map[string][]v1.Pod)
	for _, ns := range config.Namespaces {
		namespace := ns.Name
		_, err = clientset.CoreV1().Namespaces().Get(context.Background(), namespace, metav1.GetOptions{})
		if err != nil {
			LogValidationMsg(Error, "ERROR: namespace not found: %s", namespace)
			continue
		}

		// checking for service accounts shared by applications
		// map of namespace/sa -> []applicationname
		serviceAccountMap := make(map[string][]string)

		for _, application := range ns.Applications {
			pods := FindPods(application, clientset)
			applicationPods[application.Name] = pods
			//log.Printf(namespace + "/" + application.Name)
			if len(pods) == 0 {
				LogValidationMsg(Error, "application %s: no running pods found", application.Name)
			}
			ownerReferences := make(map[string]bool)
			for _, pod := range pods {
				//log.Printf("  %s %v", pod.Name, pod.OwnerReferences)
				for _, ownerReference := range pod.OwnerReferences {
					ownerReferences[ownerReference.Kind+"/"+ownerReference.Name] = true
				}
			}
			if len(ownerReferences) > 1 {
				LogValidationMsg(Error, "Application %s: multiple owners found:  %v", application.Name, MapKeys(ownerReferences))
			}
			// check ports
			for _, port := range application.Ports {
				for _, pod := range pods {
					if !HasPort(pod, port) {
						LogValidationMsg(Error, "application %s: port %v not found in pod %s/%s",
							application.Name, port, pod.Namespace, pod.Name)
					}
				}
			}
			// Check service accounts
			applicationServiceAccounts := make(map[string]bool)
			for _, sa := range application.ServiceAccounts {
				applicationServiceAccounts[sa] = true
			}
			for _, pod := range pods {
				delete(applicationServiceAccounts, pod.Spec.ServiceAccountName)
			}
			if len(applicationServiceAccounts) > 0 {
				LogValidationMsg(Error, "application %s: service accounts %v configured but not used by workloads",
					application.Name, MapKeys(applicationServiceAccounts))
			}
			for _, pod := range pods {
				sa := pod.Namespace + "/" + pod.Spec.ServiceAccountName
				serviceAccountMap[sa] = append(serviceAccountMap[sa],
					application.Name)

				if pod.Spec.ServiceAccountName == "default" {
					LogValidationMsg(Warning, "Pod %s/%s: running with default service account",
						pod.Namespace, pod.Name)
					if application.ServiceAccounts != nil && !slices.Contains(application.ServiceAccounts, "default") {
						LogValidationMsg(Warning, "application %s: Pod %s/%s: running with 'default' service account but configured with %v",
							application.Name, pod.Namespace, pod.Name, application.ServiceAccounts)
					}
				} else {
					if !slices.Contains(application.ServiceAccounts, pod.Spec.ServiceAccountName) {
						LogValidationMsg(Warning, "application %s: Pod %s/%s: running with service account '%s' but configured with %v",
							application.Name, pod.Namespace, pod.Name, pod.Spec.ServiceAccountName, application.ServiceAccounts)
					}
				}
			}
		}

		// service accounts shared by multiple applications.
		for sa, applist := range serviceAccountMap {
			if len(applist) == 1 {
				continue
			}
			LogValidationMsg(Error, "service account %s: shared by multiple applications %v", sa, applist)
		}
	}

	for _, communication := range config.Communications {
		if len(communication.Ports) == 0 {
			continue
		}
		for _, applicationName := range communication.To {
			application, _, _ := config.GetApplication(applicationName)
			if application == nil {
				continue
			}
			for _, port := range communication.Ports {
				pods := FindPods(application, clientset)
				for _, pod := range pods {
					if !HasPort(pod, port) {
						LogValidationMsg(Error, "communication %v -> %v: port %v is not configured in pod %s/%s",
							communication.From, communication.To, port, pod.Namespace, pod.Name)
					}
				}

			}
		}
		for _, communication := range config.Communications {
			for _, applicationName := range communication.To {
				application, _, _ := config.GetApplication(applicationName)
				if application == nil {
					continue
				}
				// capability linkerd must exist on target namespace
				if !slices.Contains(application.Namespace.Capabilities, "linkerd") {
					continue
				}
				// annotation linkerd.io/inject must not be disabled
				pods := applicationPods[applicationName]
				for _, pod := range pods {
					if pod.Annotations["linkerd.io/inject"] == "disabled" {
						continue
					}
				}

				if !application.Namespace.Open {
					for _, applicationNameFrom := range communication.From {
						applicationFrom, _, _ := config.GetApplication(applicationNameFrom)
						if applicationFrom != nil && !applicationFrom.Namespace.Open {
							continue
						}
						openToClosedAccess[applicationNameFrom] = applicationName
					}
				}
			}
		}
	}

	for appFrom, appTo := range openToClosedAccess {
		LogValidationMsg(Error, "Access from 'open' application '%s' to 'closed' application '%s'",
			appFrom, appTo)
	}

	return nil
}

func FindPods(application *Application, clientset *kubernetes.Clientset) []v1.Pod {
	labelSelector := &metav1.LabelSelector{
		MatchLabels:      application.MatchLabels,
		MatchExpressions: application.MatchExpressions,
	}
	selector, err := metav1.LabelSelectorAsSelector(labelSelector)
	if err != nil {
		log.Fatalf("Error creating selector: %v", err)
	}
	pods, err := clientset.CoreV1().Pods(application.Namespace.Name).List(context.TODO(), metav1.ListOptions{
		LabelSelector: selector.String(),
	})
	if err != nil {
		log.Fatalf("Error listing pods: %v", err)
	}
	return pods.Items
}

func HasPort(pod v1.Pod, port Port) bool {
	if port.Protocol == "" {
		port.Protocol = "TCP"
	}
	for _, container := range pod.Spec.Containers {
		for _, kport := range container.Ports {
			if kport.Protocol == "" {
				kport.Protocol = "TCP"
			}
			if port.Protocol == string(kport.Protocol) && (port.Port == kport.Name || port.Port == strconv.Itoa(int(kport.ContainerPort))) {
				return true
			}
		}
	}
	return false
}