package main

import (
	"context"
	"fmt"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/klog/v2"
	"strings"
)

func canonicalizeImageName(image string) string {
	// Check if image has a tag
	var name, tag string
	if parts := strings.Split(image, ":"); len(parts) > 1 {
		name = parts[0]
		tag = parts[1]
	} else {
		name = image
		tag = "latest" // Default tag
	}

	// Check if image has a host
	var host, remainder string
	if strings.Contains(name, "/") {
		parts := strings.SplitN(name, "/", 2)

		// Determine if the first part is a host
		// A host must contain a dot or colon (for registries with ports)
		if strings.Contains(parts[0], ".") || strings.Contains(parts[0], ":") {
			host = parts[0]
			remainder = parts[1]
		} else {
			// No host specified, use default docker.io
			host = "docker.io"
			remainder = name
		}
	} else {
		// No host and no /, use default docker.io and library
		host = "docker.io"
		remainder = "library/" + name
	}

	// Handle the case when remainder doesn't specify library but it's not a docker.io official image
	if host == "docker.io" && !strings.Contains(remainder, "/") {
		remainder = "library/" + remainder
	} else if host == "docker.io" && !strings.HasPrefix(remainder, "library/") {
		// Check if it's not already in the library
		if parts := strings.SplitN(remainder, "/", 2); len(parts) == 2 && parts[0] != "library" {
			// Not a library image and not already properly formatted
			// This is a docker.io user image (e.g., user/image)
			// We keep it as is
		}
	}

	return fmt.Sprintf("%s/%s:%s", host, remainder, tag)
}

func getContainers(clientset *kubernetes.Clientset, kubernetesNamespace string) map[string]bool {
	pods, err := clientset.CoreV1().Pods(kubernetesNamespace).List(context.Background(),
		metav1.ListOptions{})
	if err != nil {
		panic(err)
	}

	containers := make(map[string]bool)

	for _, pod := range pods.Items {
		klog.V(3).Infof("%s/%s\n", pod.Namespace, pod.Name)
		for _, container := range pod.Spec.InitContainers {
			klog.V(3).Infof("  %s\n", container.Image)
			containers[canonicalizeImageName(container.Image)] = true
		}
		for _, container := range pod.Spec.Containers {
			klog.V(3).Infof("  %s\n", container.Image)
			containers[canonicalizeImageName(container.Image)] = true
		}
	}
	return containers
}

func pullAndPin(kubernetesNamespace, socketPath, containerdNamespace string) error {

	// Create the image manager
	containerd, err := NewContainerd(socketPath, containerdNamespace)
	if err != nil {
		klog.Fatalf("Failed to create image manager: %v", err)
	}

	imgs, err := containerd.List()
	if err != nil {
		return err
	}

	clientset, _ := GetKubernetesConnection()

	containers := getContainers(clientset, kubernetesNamespace)
	for container, _ := range containers {
		klog.V(3).Infof("Found container %s\n", container)
	}

	// unpin images that are not used
	for container, pinned := range imgs {
		if !containers[container] && pinned {
			klog.Infof("Unpinning %s\n", container)
			err := containerd.Unpin(container)
			if err != nil {
				klog.Warningf("  error: %v", err)
			}
		}
	}

	// Pull images that are used
	for container := range containers {
		if _, found := imgs[container]; !found {
			klog.Infof("Pulling %s\n", container)
			err := containerd.Pull(container)
			if err != nil {
				klog.Warningf("error: %v", err)
			}
		}
	}

	imgs, err = containerd.List()
	if err != nil {
		return err
	}

	// Pin images that are used and present
	for container := range containers {
		if pinned, found := imgs[container]; found && !pinned {
			klog.Infof("Pinning %s\n", container)
			err := containerd.Pin(container)
			if err != nil {
				klog.Warningf("  error: %v", err)
			}
		}
	}
	return nil
}