// Package discovery watches Autoscalers and caches their metrics targets by
// targetRef kind.
package discovery

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"sync"

	promlabels "github.com/prometheus/prometheus/model/labels"
	"k8s.io/apimachinery/pkg/types"
	toolscache "k8s.io/client-go/tools/cache"
	"sigs.k8s.io/controller-runtime/pkg/cache"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"

	"github.com/tidbcloud/cluster-service-ng/internal/capacity-planner/api/v1alpha1"
	"github.com/tidbcloud/cluster-service-ng/internal/capacity-planner/pkg/target"
)

// AllKinds requests every discovered Pod target regardless of targetRef kind.
const AllKinds = ""

// Discovery provides a read-only view of the target cache. An empty kind
// returns all unique Pod targets; a non-empty kind returns only targets whose
// Autoscaler targetRef has that exact kind.
type Discovery interface {
	Targets(kind string) []target.Target
}

// Runnable is a Discovery that maintains its cache by watching Autoscalers.
type Runnable interface {
	Discovery
	Start(ctx context.Context) error
	NeedLeaderElection() bool
}

// Config configures Autoscaler target discovery.
type Config struct {
	Informers cache.Informers
}

type discovery struct {
	informers cache.Informers

	mu                sync.RWMutex
	autoscalerTargets map[types.NamespacedName]targetReference
	targetsByKind     map[string]map[string]*targetRegistration
}

type targetReference struct {
	kind string
	key  string
}

type targetRegistration struct {
	target target.Target
	count  int
}

var _ Runnable = (*discovery)(nil)

// New constructs an empty discovery cache.
func New(cfg *Config) (Runnable, error) {
	if cfg == nil {
		return nil, fmt.Errorf("target discovery: config is required")
	}
	if cfg.Informers == nil {
		return nil, fmt.Errorf("target discovery: informers are required")
	}
	return &discovery{
		informers:         cfg.Informers,
		autoscalerTargets: make(map[types.NamespacedName]targetReference),
		targetsByKind:     make(map[string]map[string]*targetRegistration),
	}, nil
}

// Start maintains the target cache until ctx is cancelled.
func (d *discovery) Start(ctx context.Context) error {
	informer, err := d.informers.GetInformer(ctx, &v1alpha1.Autoscaler{})
	if err != nil {
		return fmt.Errorf("target discovery: get Autoscaler informer: %w", err)
	}
	registration, err := informer.AddEventHandler(toolscache.ResourceEventHandlerFuncs{
		AddFunc: func(obj any) {
			d.handleAutoscaler(ctx, obj)
		},
		UpdateFunc: func(_, newObj any) {
			d.handleAutoscaler(ctx, newObj)
		},
		DeleteFunc: func(obj any) {
			d.handleAutoscalerDelete(obj)
		},
	})
	if err != nil {
		return fmt.Errorf("target discovery: add Autoscaler event handler: %w", err)
	}
	defer func() {
		if err := informer.RemoveEventHandler(registration); err != nil && !errors.Is(err, context.Canceled) {
			log.FromContext(ctx).Error(err, "unable to remove target discovery event handler")
		}
	}()

	<-ctx.Done()
	return nil
}

// Targets returns a stable, deduplicated snapshot from the cache.
func (d *discovery) Targets(kind string) []target.Target {
	d.mu.RLock()
	defer d.mu.RUnlock()

	byKey := make(map[string]target.Target)
	if kind == AllKinds {
		for _, registrations := range d.targetsByKind {
			for key, registration := range registrations {
				byKey[key] = registration.target
			}
		}
	} else {
		for key, registration := range d.targetsByKind[kind] {
			byKey[key] = registration.target
		}
	}
	keys := make([]string, 0, len(byKey))
	for key := range byKey {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	result := make([]target.Target, 0, len(keys))
	for _, key := range keys {
		result = append(result, byKey[key])
	}
	return result
}

func (d *discovery) handleAutoscaler(ctx context.Context, obj any) {
	autoscaler, ok := autoscalerFromEvent(obj)
	if !ok {
		return
	}
	key := client.ObjectKeyFromObject(autoscaler)
	if !autoscaler.DeletionTimestamp.IsZero() {
		d.setAutoscalerTarget(key, "", nil)
		return
	}
	discovered, err := target.NewPodTargetForRef(autoscaler.Namespace, autoscaler.Spec.TargetRef)
	if err != nil {
		d.setAutoscalerTarget(key, "", nil)
		if ctx.Err() == nil {
			log.FromContext(ctx).Error(err, "skipping Autoscaler with invalid metrics target",
				"namespace", autoscaler.Namespace,
				"name", autoscaler.Name,
			)
		}
		return
	}
	d.setAutoscalerTarget(key, autoscaler.Spec.TargetRef.Kind, discovered)
}

func (d *discovery) handleAutoscalerDelete(obj any) {
	autoscaler, ok := autoscalerFromEvent(obj)
	if !ok {
		return
	}
	d.setAutoscalerTarget(client.ObjectKeyFromObject(autoscaler), "", nil)
}

func (d *discovery) setAutoscalerTarget(autoscalerKey types.NamespacedName, kind string, discovered target.Target) {
	d.mu.Lock()
	defer d.mu.Unlock()

	newKey := ""
	if discovered != nil {
		newKey = targetKey(discovered)
	}
	previous, existed := d.autoscalerTargets[autoscalerKey]
	if existed && discovered != nil && previous.kind == kind && previous.key == newKey {
		return
	}
	if existed {
		delete(d.autoscalerTargets, autoscalerKey)
		registrations := d.targetsByKind[previous.kind]
		registration := registrations[previous.key]
		registration.count--
		if registration.count == 0 {
			delete(registrations, previous.key)
		}
		if len(registrations) == 0 {
			delete(d.targetsByKind, previous.kind)
		}
	}
	if discovered == nil {
		return
	}

	d.autoscalerTargets[autoscalerKey] = targetReference{kind: kind, key: newKey}
	registrations := d.targetsByKind[kind]
	if registrations == nil {
		registrations = make(map[string]*targetRegistration)
		d.targetsByKind[kind] = registrations
	}
	if registration := registrations[newKey]; registration != nil {
		registration.count++
	} else {
		registrations[newKey] = &targetRegistration{target: discovered, count: 1}
	}
}

func targetKey(discovered target.Target) string {
	return discovered.Namespace() + "\x00" + promlabels.FromMap(discovered.Labels()).String()
}

func autoscalerFromEvent(obj any) (*v1alpha1.Autoscaler, bool) {
	switch typed := obj.(type) {
	case *v1alpha1.Autoscaler:
		return typed, true
	case toolscache.DeletedFinalStateUnknown:
		autoscaler, ok := typed.Obj.(*v1alpha1.Autoscaler)
		return autoscaler, ok
	default:
		return nil, false
	}
}

// NeedLeaderElection ensures only the active controller manager maintains the
// process-local target cache.
func (*discovery) NeedLeaderElection() bool {
	return true
}