package controller

import (
	"context"
	"time"

	"github.com/sandbox0-ai/sandbox0/manager/pkg/apis/sandbox0/v1alpha1"
	"go.uber.org/zap"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/client-go/kubernetes"
	corelisters "k8s.io/client-go/listers/core/v1"
	"k8s.io/client-go/tools/record"
)

const staleDeletingPodForceDeleteAfter = 10 * time.Minute
const hardExpiredSandboxCleanupBatchSize = 500

// SandboxRuntimePauser defines the interface for checkpointing and releasing a sandbox runtime.
type SandboxRuntimePauser interface {
	PauseSandboxByID(ctx context.Context, sandboxID string) error
}

// SandboxTerminator defines the interface for terminating sandboxes.
type SandboxTerminator interface {
	TerminateSandboxByID(ctx context.Context, sandboxID string) error
}

// HardExpiredSandboxLister lists durable sandboxes whose hard TTL has expired.
type HardExpiredSandboxLister interface {
	ListHardExpiredSandboxIDs(ctx context.Context, now time.Time, limit int) ([]string, error)
}

// CleanupController handles cleanup of excess idle pods and expired active pods
type CleanupController struct {
	k8sClient         kubernetes.Interface
	podLister         corelisters.PodLister
	templateLister    TemplateLister
	recorder          record.EventRecorder
	clock             TimeProvider
	runtimePauser     SandboxRuntimePauser
	sandboxTerminator SandboxTerminator
	hardExpiredLister HardExpiredSandboxLister
	logger            *zap.Logger
	interval          time.Duration
}

// TimeProvider provides time functions, allowing for synchronized time across clusters
type TimeProvider interface {
	Now() time.Time
	Since(t time.Time) time.Duration
	Until(t time.Time) time.Duration
}

// systemTime is the default implementation using system time
type systemTime struct{}

func (systemTime) Now() time.Time                  { return time.Now() }
func (systemTime) Since(t time.Time) time.Duration { return time.Since(t) }
func (systemTime) Until(t time.Time) time.Duration { return time.Until(t) }

// TemplateLister interface for listing templates
type TemplateLister interface {
	List() ([]*v1alpha1.SandboxTemplate, error)
	Get(namespace, name string) (*v1alpha1.SandboxTemplate, error)
}

// NewCleanupController creates a new CleanupController
func NewCleanupController(
	k8sClient kubernetes.Interface,
	podLister corelisters.PodLister,
	templateLister TemplateLister,
	recorder record.EventRecorder,
	clock TimeProvider,
	runtimePauser SandboxRuntimePauser,
	sandboxTerminator SandboxTerminator,
	logger *zap.Logger,
	interval time.Duration,
) *CleanupController {
	// Use system time as fallback if clock is nil
	if clock == nil {
		clock = systemTime{}
	}

	hardExpiredLister, _ := sandboxTerminator.(HardExpiredSandboxLister)
	if hardExpiredLister == nil {
		hardExpiredLister, _ = runtimePauser.(HardExpiredSandboxLister)
	}
	return &CleanupController{
		k8sClient:         k8sClient,
		podLister:         podLister,
		templateLister:    templateLister,
		recorder:          recorder,
		clock:             clock,
		runtimePauser:     runtimePauser,
		sandboxTerminator: sandboxTerminator,
		hardExpiredLister: hardExpiredLister,
		logger:            logger,
		interval:          interval,
	}
}

// Start starts the cleanup controller
func (cc *CleanupController) Start(ctx context.Context) error {
	cc.logger.Info("Starting cleanup controller", zap.Duration("interval", cc.interval))

	ticker := time.NewTicker(cc.interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			cc.logger.Info("Cleanup controller stopped")
			return ctx.Err()
		case <-ticker.C:
			if err := cc.runCleanup(ctx); err != nil {
				cc.logger.Error("Cleanup failed", zap.Error(err))
			}
		}
	}
}

// runCleanup runs a cleanup cycle
func (cc *CleanupController) runCleanup(ctx context.Context) error {
	cc.logger.Debug("Running cleanup cycle")

	templates, err := cc.templateLister.List()
	if err != nil {
		return err
	}

	for _, template := range templates {
		if err := cc.cleanupExpired(ctx, template); err != nil {
			cc.logger.Error("Failed to cleanup expired sandbox",
				zap.String("template", template.Name),
				zap.Error(err),
			)
		}
	}

	if err := cc.cleanupHardExpiredDurableSandboxes(ctx); err != nil {
		cc.logger.Error("Failed to cleanup hard-expired durable sandboxes", zap.Error(err))
	}

	return nil
}

func (cc *CleanupController) cleanupHardExpiredDurableSandboxes(ctx context.Context) error {
	if cc.hardExpiredLister == nil || cc.sandboxTerminator == nil {
		return nil
	}
	now := cc.clock.Now()
	sandboxIDs, err := cc.hardExpiredLister.ListHardExpiredSandboxIDs(ctx, now, hardExpiredSandboxCleanupBatchSize)
	if err != nil {
		return err
	}
	for _, sandboxID := range sandboxIDs {
		if sandboxID == "" {
			continue
		}
		if err := cc.sandboxTerminator.TerminateSandboxByID(ctx, sandboxID); err != nil {
			cc.logger.Error("Failed to delete hard-expired durable sandbox",
				zap.String("sandboxID", sandboxID),
				zap.Error(err),
			)
			continue
		}
	}
	if len(sandboxIDs) > 0 {
		cc.logger.Info("Deleted hard-expired durable sandboxes", zap.Int("count", len(sandboxIDs)))
	}
	return nil
}

// cleanupExpired cleans up stale deleting idle pods and expired active pods.
func (cc *CleanupController) cleanupExpired(ctx context.Context, template *v1alpha1.SandboxTemplate) error {
	now := cc.clock.Now()

	if err := cc.cleanupStaleDeletingIdlePods(ctx, template, now); err != nil {
		return err
	}

	// Get all active pods for this template
	pods, err := cc.podLister.Pods(template.Namespace).List(labels.SelectorFromSet(map[string]string{
		LabelTemplateID: template.Name,
		LabelPoolType:   PoolTypeActive,
	}))
	if err != nil {
		return err
	}

	expiredCount := 0

	for _, pod := range pods {
		if pod.DeletionTimestamp != nil {
			cc.forceDeleteStaleDeletingPod(ctx, template, pod, now)
			continue
		}
		if pod.Status.Phase == corev1.PodSucceeded {
			if cc.deleteCompletedPod(ctx, template, pod) {
				expiredCount++
			}
			continue
		}

		// Hard expiry deletes the durable sandbox identity and state.
		if hardExpiresAtStr := pod.Annotations[AnnotationHardExpiresAt]; hardExpiresAtStr != "" {
			hardExpiresAt, err := time.Parse(time.RFC3339, hardExpiresAtStr)
			if err != nil {
				cc.logger.Warn("Invalid hard-expires-at annotation",
					zap.String("pod", pod.Name),
					zap.String("value", hardExpiresAtStr),
				)
			} else if now.After(hardExpiresAt) {
				cc.logger.Info("Deleting hard-expired sandbox",
					zap.String("pod", pod.Name),
					zap.Time("hardExpiresAt", hardExpiresAt),
				)
				sandboxID := cleanupSandboxIDFromPod(pod)
				if cc.sandboxTerminator != nil {
					if err := cc.sandboxTerminator.TerminateSandboxByID(ctx, sandboxID); err != nil {
						cc.logger.Error("Failed to delete hard-expired sandbox",
							zap.String("pod", pod.Name),
							zap.String("sandboxID", sandboxID),
							zap.Error(err),
						)
						continue
					}
				} else {
					cc.logger.Warn("SandboxTerminator not configured, deleting pod directly",
						zap.String("pod", pod.Name),
					)
					if err := cc.k8sClient.CoreV1().Pods(pod.Namespace).Delete(ctx, pod.Name, metav1.DeleteOptions{}); err != nil {
						cc.logger.Error("Failed to delete hard-expired pod",
							zap.String("pod", pod.Name),
							zap.Error(err),
						)
						continue
					}
				}

				cc.recorder.Eventf(template, corev1.EventTypeNormal, "HardExpiredSandboxDeleted",
					"Deleted hard-expired sandbox %s", sandboxID)
				expiredCount++
				continue
			}
		}

		// Check if pod has expiration annotation
		expiresAtStr, ok := pod.Annotations[AnnotationExpiresAt]
		if !ok {
			continue
		}

		expiresAt, err := time.Parse(time.RFC3339, expiresAtStr)
		if err != nil {
			cc.logger.Warn("Invalid expires-at annotation",
				zap.String("pod", pod.Name),
				zap.String("value", expiresAtStr),
			)
			continue
		}

		// Check if pod is expired
		if now.After(expiresAt) {
			cc.logger.Info("Pausing expired sandbox runtime",
				zap.String("pod", pod.Name),
				zap.Time("expiresAt", expiresAt),
			)

			if cc.runtimePauser != nil {
				err := cc.runtimePauser.PauseSandboxByID(ctx, cleanupSandboxIDFromPod(pod))
				if err != nil {
					cc.logger.Error("Failed to pause expired sandbox runtime",
						zap.String("pod", pod.Name),
						zap.Error(err),
					)
					continue
				}

				cc.recorder.Eventf(template, corev1.EventTypeNormal, "ExpiredSandboxRuntimePaused",
					"Paused expired sandbox runtime %s", pod.Name)
				expiredCount++
			} else {
				cc.logger.Warn("Sandbox runtime pauser not configured, skipping pause for expired pod",
					zap.String("pod", pod.Name),
				)
			}
		}
	}

	if expiredCount > 0 {
		cc.logger.Info("Processed expired pods",
			zap.String("template", template.Name),
			zap.Int("count", expiredCount),
		)
	}

	return nil
}

func (cc *CleanupController) cleanupStaleDeletingIdlePods(ctx context.Context, template *v1alpha1.SandboxTemplate, now time.Time) error {
	pods, err := cc.podLister.Pods(template.Namespace).List(labels.SelectorFromSet(map[string]string{
		LabelTemplateID: template.Name,
		LabelPoolType:   PoolTypeIdle,
	}))
	if err != nil {
		return err
	}

	for _, pod := range pods {
		if pod.DeletionTimestamp != nil {
			cc.forceDeleteStaleDeletingPod(ctx, template, pod, now)
		}
	}
	return nil
}

func (cc *CleanupController) deleteCompletedPod(ctx context.Context, template *v1alpha1.SandboxTemplate, pod *corev1.Pod) bool {
	if pod == nil {
		return false
	}

	cc.logger.Info("Deleting completed sandbox pod",
		zap.String("pod", pod.Name),
		zap.String("phase", string(pod.Status.Phase)),
	)

	if cc.sandboxTerminator != nil {
		if err := cc.sandboxTerminator.TerminateSandboxByID(ctx, cleanupSandboxIDFromPod(pod)); err != nil {
			cc.logger.Error("Failed to delete completed sandbox pod",
				zap.String("pod", pod.Name),
				zap.Error(err),
			)
			return false
		}
	} else {
		if cc.k8sClient == nil {
			cc.logger.Warn("Kubernetes client not configured, skipping completed sandbox pod delete",
				zap.String("pod", pod.Name),
			)
			return false
		}
		if err := cc.k8sClient.CoreV1().Pods(pod.Namespace).Delete(ctx, pod.Name, metav1.DeleteOptions{}); err != nil {
			cc.logger.Error("Failed to delete completed sandbox pod",
				zap.String("pod", pod.Name),
				zap.Error(err),
			)
			return false
		}
	}

	cc.recorder.Eventf(template, corev1.EventTypeNormal, "CompletedPodDeleted",
		"Deleted completed sandbox pod %s", pod.Name)
	return true
}

func cleanupSandboxIDFromPod(pod *corev1.Pod) string {
	if pod == nil {
		return ""
	}
	if pod.Labels != nil {
		if sandboxID := pod.Labels[LabelSandboxID]; sandboxID != "" {
			return sandboxID
		}
	}
	if pod.Annotations != nil {
		if sandboxID := pod.Annotations[AnnotationSandboxID]; sandboxID != "" {
			return sandboxID
		}
	}
	return pod.Name
}

func (cc *CleanupController) forceDeleteStaleDeletingPod(ctx context.Context, template *v1alpha1.SandboxTemplate, pod *corev1.Pod, now time.Time) {
	if pod == nil || pod.DeletionTimestamp == nil {
		return
	}
	if now.Sub(pod.DeletionTimestamp.Time) < staleDeletingPodForceDeleteAfter {
		return
	}
	if cc.k8sClient == nil {
		cc.logger.Warn("Kubernetes client not configured, skipping stale deleting pod force delete",
			zap.String("pod", pod.Name),
			zap.Time("deletionTimestamp", pod.DeletionTimestamp.Time),
		)
		return
	}

	gracePeriodSeconds := int64(0)
	uid := pod.UID
	deleteOptions := metav1.DeleteOptions{
		GracePeriodSeconds: &gracePeriodSeconds,
		Preconditions: &metav1.Preconditions{
			UID: &uid,
		},
	}
	cc.logger.Warn("Force deleting stale terminating pod",
		zap.String("pod", pod.Name),
		zap.Time("deletionTimestamp", pod.DeletionTimestamp.Time),
	)
	if err := cc.k8sClient.CoreV1().Pods(pod.Namespace).Delete(ctx, pod.Name, deleteOptions); err != nil {
		cc.logger.Error("Failed to force delete stale terminating pod",
			zap.String("pod", pod.Name),
			zap.Error(err),
		)
		return
	}
	cc.recorder.Eventf(template, corev1.EventTypeWarning, "StaleDeletingPodForceDeleted",
		"Force deleted stale terminating pod %s", pod.Name)
}
