package main

import (
	"context"
	"crypto/md5"
	"encoding/base64"
	"fmt"
	"github.com/containerd/containerd"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/images"
	"github.com/containerd/containerd/leases"
	"github.com/containerd/containerd/namespaces"
	"k8s.io/klog/v2"
)

// Containerd implements ContainerRuntime interface
type Containerd struct {
	client    *containerd.Client
	ctx       context.Context
	namespace string
}

// NewContainerd creates a new Containerd
func NewContainerd(socketPath, namespace string) (*Containerd, error) {
	client, err := containerd.New(socketPath)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to containerd at %s: %w", socketPath, err)
	}
	client.LeasesService()

	ctx := namespaces.WithNamespace(context.Background(), namespace)

	return &Containerd{
		client:    client,
		ctx:       ctx,
		namespace: namespace,
	}, nil
}

// List returns all images with their pinned status
func (m *Containerd) List() (map[string]bool, error) {
	// Get all images
	images, err := m.client.ImageService().List(m.ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to list images: %w", err)
	}

	for _, image := range images {
		klog.V(3).Infof("Image '%s' digest '%s'", image.Name,
			image.Target.Digest.String())
	}

	// Get all leases
	leases, err := m.client.LeasesService().List(m.ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to list leases: %w", err)
	}

	// Create a map of image references that are pinned
	pinnedImages := make(map[string]bool)
	for _, lease := range leases {
		// Check if lease has labels referencing an image
		if label, ok := lease.Labels["containerd.io/gc.ref.content.image"]; ok {
			pinnedImages[label] = true
		}
	}

	// Create the result list
	var result = make(map[string]bool)
	for _, img := range images {
		result[img.Name] = pinnedImages[img.Name]
	}

	return result, nil
}

// Pin creates a lease for an image to prevent garbage collection
func (m *Containerd) Pin(imageRef string) error {
	// Create a unique lease ID based on image reference
	leaseID := fmt.Sprintf("pin-%s", generateID(imageRef))

	// Get the image to validate it exists
	_, err := m.client.ImageService().Get(m.ctx, imageRef)
	if err != nil {
		return fmt.Errorf("failed to get image %s: %w", imageRef, err)
	}

	leaseList, err := m.findLeases(imageRef)
	if err != nil {
		return fmt.Errorf("Failed to get leases for image %s: %v", imageRef, err)
	}
	if len(leaseList) > 0 {
		return nil
	}

	// Create a new lease
	opts := []leases.Opt{
		leases.WithID(leaseID),
		leases.WithLabels(map[string]string{
			"containerd.io/gc.ref.content.image": imageRef,
		}),
	}

	_, err = m.client.LeasesService().Create(m.ctx, opts...)
	if err != nil {
		return fmt.Errorf("failed to create lease: %w", err)
	}

	return nil
}

// Unpin removes a lease for an image allowing garbage collection
func (m *Containerd) Unpin(imageRef string) error {
	leases, err := m.findLeases(imageRef)
	if err != nil {
		return err
	}
	for _, lease := range leases {
		if err := m.client.LeasesService().Delete(m.ctx, lease); err != nil {
			return fmt.Errorf("failed to delete lease %s: %w", lease.ID, err)
		}
	}
	return nil
}

func (m *Containerd) findLeases(imageRef string) ([]leases.Lease, error) {
	// List all leases
	leaseList, err := m.client.LeasesService().List(m.ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to list leases: %w", err)
	}

	// Find leases that reference our image
	var leases = make([]leases.Lease, 0)
	for _, lease := range leaseList {
		// Check if this lease has a label referencing our image
		if label, ok := lease.Labels["containerd.io/gc.ref.content.image"]; ok && label == imageRef {
			leases = append(leases, lease)
		}
	}
	return leases, nil
}

// Pull pulls an image from a registry
func (m *Containerd) Pull(imageRef string) error {
	// Set up pull options
	pullOpts := []containerd.RemoteOpt{
		//containerd.WithPlatformMatcher(platforms.Default()),
		containerd.WithPullUnpack,
	}

	// Pull the image
	_, err := m.client.Pull(m.ctx, imageRef, pullOpts...)
	if err != nil {
		return fmt.Errorf("failed to pull image %s: %w", imageRef, err)
	}

	return nil
}

// Remove deletes an image
func (m *Containerd) Remove(imageRef string) error {
	// Get the image

	_, err := m.client.ImageService().Get(m.ctx, imageRef)
	if err != nil {
		if errdefs.IsNotFound(err) {
			return fmt.Errorf("image %s not found", imageRef)
		}
		return fmt.Errorf("failed to get image %s: %w", imageRef, err)
	}

	// Delete the image
	err = m.client.ImageService().Delete(m.ctx, imageRef, images.SynchronousDelete())
	if err != nil {
		return fmt.Errorf("failed to delete image %s: %w", imageRef, err)
	}

	return nil
}

// Close closes the containerd client connection
func (m *Containerd) Close() error {
	return m.client.Close()
}

// generateID creates a random unique ID
func generateID(image string) string {

	md5Hash := md5.Sum([]byte(image))
	md5Base64 := base64.StdEncoding.EncodeToString(md5Hash[:])
	return md5Base64
}