package merger

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	apiextensionsv1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1"
	apiequality "k8s.io/apimachinery/pkg/api/equality"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
)

// SchemaResolver returns the structural schema for one custom-resource GVK.
type SchemaResolver interface {
	Resolve(ctx context.Context, gvk schema.GroupVersionKind) (*apiextensionsv1.JSONSchemaProps, error)
}

// Interface merges a partial desired object into a live object in place.
type Interface interface {
	Merge(
		ctx context.Context,
		gvk schema.GroupVersionKind,
		live map[string]interface{},
		desired map[string]interface{},
	) error
}

// Strategy defines the local merge behavior for a path.
type Strategy string

const (
	// StrategyMergeMap merges the desired map into the live map by key.
	StrategyMergeMap Strategy = "MergeMap"
	// StrategyMergeListMap merges list entries using Rule.ListMapKeys.
	StrategyMergeListMap Strategy = "MergeListMap"
	// StrategyMergeSet retains distinct values from both lists.
	StrategyMergeSet Strategy = "MergeSet"
	// StrategyReplace replaces the live value with the desired value.
	StrategyReplace Strategy = "Replace"
)

// Rule overrides the CRD merge behavior for one JSON-pointer path. Empty Group
// or Kind fields match every group or kind. A "*" path segment matches an item
// inside a list-map.
type Rule struct {
	Group       string
	Kind        string
	Path        string
	Strategy    Strategy
	ListMapKeys []string
}

type compiledRule struct {
	Rule
	segments []string
}

// Merger applies explicit rules first and otherwise follows the target CRD's
// x-kubernetes-map-type and x-kubernetes-list-type markers.
type Merger struct {
	resolver SchemaResolver
	rules    []compiledRule
}

var _ Interface = (*Merger)(nil)

// New creates a schema-aware merger.
func New(resolver SchemaResolver, rules ...Rule) (*Merger, error) {
	if resolver == nil {
		return nil, fmt.Errorf("merger: schema resolver is required")
	}
	compiled := make([]compiledRule, 0, len(rules))
	for i := range rules {
		segments, err := parseJSONPointer(rules[i].Path)
		if err != nil {
			return nil, fmt.Errorf("merger: rule %d: %w", i, err)
		}
		if err := validateRule(&rules[i]); err != nil {
			return nil, fmt.Errorf("merger: rule %d path %q: %w", i, rules[i].Path, err)
		}
		compiled = append(compiled, compiledRule{Rule: rules[i], segments: segments})
	}
	return &Merger{resolver: resolver, rules: compiled}, nil
}

// Merge merges desired into live. The merge is staged on a deep copy so a
// validation error never leaves live partially modified.
func (m *Merger) Merge(
	ctx context.Context,
	gvk schema.GroupVersionKind,
	live map[string]interface{},
	desired map[string]interface{},
) error {
	if m == nil || m.resolver == nil {
		return fmt.Errorf("merger: nil merger")
	}
	if live == nil {
		return fmt.Errorf("merger: live object is required")
	}
	if desired == nil {
		return fmt.Errorf("merger: desired object is required")
	}
	if gvk.Empty() {
		return fmt.Errorf("merger: target GVK is required")
	}

	rootSchema, err := m.resolver.Resolve(ctx, gvk)
	if err != nil {
		return fmt.Errorf("merger: resolve schema for %s: %w", gvk, err)
	}
	if rootSchema == nil {
		return fmt.Errorf("merger: schema resolver returned nil schema for %s", gvk)
	}

	staged := runtime.DeepCopyJSONValue(live).(map[string]interface{})
	merged, err := m.mergeValue(gvk, staged, desired, rootSchema, nil)
	if err != nil {
		return err
	}
	mergedMap, ok := merged.(map[string]interface{})
	if !ok {
		return fmt.Errorf("merger: root value for %s must remain an object", gvk)
	}

	clear(live)
	for key, value := range mergedMap {
		live[key] = value
	}
	return nil
}

func (m *Merger) mergeValue(
	gvk schema.GroupVersionKind,
	live interface{},
	desired interface{},
	fieldSchema *apiextensionsv1.JSONSchemaProps,
	path []string,
) (interface{}, error) {
	strategy, listMapKeys, err := m.strategyFor(gvk, desired, fieldSchema, path)
	if err != nil {
		return nil, err
	}

	switch strategy {
	case StrategyMergeMap:
		return m.mergeMap(gvk, live, desired, fieldSchema, path)
	case StrategyMergeListMap:
		return m.mergeListMap(gvk, live, desired, fieldSchema, path, listMapKeys)
	case StrategyMergeSet:
		return mergeSet(live, desired, path)
	case StrategyReplace:
		return runtime.DeepCopyJSONValue(desired), nil
	default:
		return nil, fmt.Errorf("merger: path %q has unsupported strategy %q", pointer(path), strategy)
	}
}

func (m *Merger) mergeMap(
	gvk schema.GroupVersionKind,
	live interface{},
	desired interface{},
	fieldSchema *apiextensionsv1.JSONSchemaProps,
	path []string,
) (interface{}, error) {
	desiredMap, ok := desired.(map[string]interface{})
	if !ok {
		return nil, fmt.Errorf("merger: path %q must be an object for strategy %s", pointer(path), StrategyMergeMap)
	}
	liveMap, ok := live.(map[string]interface{})
	if !ok {
		liveMap = map[string]interface{}{}
	}

	for key, desiredValue := range desiredMap {
		childPath := appendPath(path, key)
		merged, err := m.mergeValue(gvk, liveMap[key], desiredValue, childSchema(fieldSchema, key), childPath)
		if err != nil {
			return nil, err
		}
		liveMap[key] = merged
	}
	return liveMap, nil
}

func (m *Merger) mergeListMap(
	gvk schema.GroupVersionKind,
	live interface{},
	desired interface{},
	fieldSchema *apiextensionsv1.JSONSchemaProps,
	path []string,
	keys []string,
) (interface{}, error) {
	if len(keys) == 0 {
		return nil, fmt.Errorf("merger: list-map path %q has no keys", pointer(path))
	}
	desiredList, ok := desired.([]interface{})
	if !ok {
		return nil, fmt.Errorf("merger: path %q must be an array for strategy %s", pointer(path), StrategyMergeListMap)
	}
	liveList, ok := live.([]interface{})
	if !ok {
		liveList = []interface{}{}
	}

	index := make(map[string]int, len(liveList))
	for i := range liveList {
		identity, err := listMapIdentity(liveList[i], keys, path)
		if err != nil {
			return nil, fmt.Errorf("merger: live list item %d: %w", i, err)
		}
		if _, duplicate := index[identity]; duplicate {
			return nil, fmt.Errorf("merger: live list-map path %q contains duplicate key %s", pointer(path), identity)
		}
		index[identity] = i
	}

	seenDesired := make(map[string]struct{}, len(desiredList))
	itemSchema := listItemSchema(fieldSchema)
	itemPath := appendPath(path, "*")
	for i := range desiredList {
		identity, err := listMapIdentity(desiredList[i], keys, path)
		if err != nil {
			return nil, fmt.Errorf("merger: desired list item %d: %w", i, err)
		}
		if _, duplicate := seenDesired[identity]; duplicate {
			return nil, fmt.Errorf("merger: desired list-map path %q contains duplicate key %s", pointer(path), identity)
		}
		seenDesired[identity] = struct{}{}

		if liveIndex, found := index[identity]; found {
			merged, err := m.mergeValue(gvk, liveList[liveIndex], desiredList[i], itemSchema, itemPath)
			if err != nil {
				return nil, err
			}
			liveList[liveIndex] = merged
			continue
		}
		liveList = append(liveList, runtime.DeepCopyJSONValue(desiredList[i]))
	}
	return liveList, nil
}

func mergeSet(live, desired interface{}, path []string) (interface{}, error) {
	desiredList, ok := desired.([]interface{})
	if !ok {
		return nil, fmt.Errorf("merger: path %q must be an array for strategy %s", pointer(path), StrategyMergeSet)
	}
	liveList, ok := live.([]interface{})
	if !ok {
		liveList = []interface{}{}
	}
	for i := range desiredList {
		found := false
		for j := range liveList {
			if apiequality.Semantic.DeepEqual(liveList[j], desiredList[i]) {
				found = true
				break
			}
		}
		if !found {
			liveList = append(liveList, runtime.DeepCopyJSONValue(desiredList[i]))
		}
	}
	return liveList, nil
}

func (m *Merger) strategyFor(
	gvk schema.GroupVersionKind,
	desired interface{},
	fieldSchema *apiextensionsv1.JSONSchemaProps,
	path []string,
) (Strategy, []string, error) {
	if rule := m.ruleFor(gvk, path); rule != nil {
		return rule.Strategy, rule.ListMapKeys, nil
	}

	switch desired.(type) {
	case map[string]interface{}:
		if fieldSchema != nil && fieldSchema.XMapType != nil {
			switch *fieldSchema.XMapType {
			case "atomic":
				return StrategyReplace, nil, nil
			case "granular":
				return StrategyMergeMap, nil, nil
			default:
				return "", nil, fmt.Errorf("merger: path %q has unsupported x-kubernetes-map-type %q", pointer(path), *fieldSchema.XMapType)
			}
		}
		return StrategyMergeMap, nil, nil
	case []interface{}:
		if fieldSchema == nil || fieldSchema.XListType == nil {
			return StrategyReplace, nil, nil
		}
		switch *fieldSchema.XListType {
		case "atomic":
			return StrategyReplace, nil, nil
		case "map":
			return StrategyMergeListMap, fieldSchema.XListMapKeys, nil
		case "set":
			return StrategyMergeSet, nil, nil
		default:
			return "", nil, fmt.Errorf("merger: path %q has unsupported x-kubernetes-list-type %q", pointer(path), *fieldSchema.XListType)
		}
	default:
		return StrategyReplace, nil, nil
	}
}

func (m *Merger) ruleFor(gvk schema.GroupVersionKind, path []string) *compiledRule {
	for i := range m.rules {
		rule := &m.rules[i]
		if rule.Group != "" && rule.Group != gvk.Group {
			continue
		}
		if rule.Kind != "" && rule.Kind != gvk.Kind {
			continue
		}
		if pathMatches(rule.segments, path) {
			return rule
		}
	}
	return nil
}

func childSchema(parent *apiextensionsv1.JSONSchemaProps, key string) *apiextensionsv1.JSONSchemaProps {
	if parent == nil {
		return nil
	}
	if property, ok := parent.Properties[key]; ok {
		copy := property
		return &copy
	}
	if parent.AdditionalProperties != nil && parent.AdditionalProperties.Schema != nil {
		return parent.AdditionalProperties.Schema
	}
	return nil
}

func listItemSchema(parent *apiextensionsv1.JSONSchemaProps) *apiextensionsv1.JSONSchemaProps {
	if parent == nil || parent.Items == nil {
		return nil
	}
	return parent.Items.Schema
}

func listMapIdentity(value interface{}, keys []string, path []string) (string, error) {
	object, ok := value.(map[string]interface{})
	if !ok {
		return "", fmt.Errorf("list-map path %q item must be an object", pointer(path))
	}
	identity := make([]interface{}, len(keys))
	for i := range keys {
		value, exists := object[keys[i]]
		if !exists {
			return "", fmt.Errorf("list-map path %q item is missing key %q", pointer(path), keys[i])
		}
		identity[i] = value
	}
	data, err := json.Marshal(identity)
	if err != nil {
		return "", fmt.Errorf("encode list-map identity: %w", err)
	}
	return string(data), nil
}

func validateRule(rule *Rule) error {
	switch rule.Strategy {
	case StrategyMergeMap, StrategyMergeSet, StrategyReplace:
		if len(rule.ListMapKeys) != 0 {
			return fmt.Errorf("listMapKeys are only valid for strategy %s", StrategyMergeListMap)
		}
	case StrategyMergeListMap:
		if len(rule.ListMapKeys) == 0 {
			return fmt.Errorf("listMapKeys are required for strategy %s", StrategyMergeListMap)
		}
	default:
		return fmt.Errorf("unsupported strategy %q", rule.Strategy)
	}
	return nil
}

func parseJSONPointer(value string) ([]string, error) {
	if value == "" {
		return nil, nil
	}
	if !strings.HasPrefix(value, "/") {
		return nil, fmt.Errorf("path %q must be a JSON pointer", value)
	}
	encoded := strings.Split(strings.TrimPrefix(value, "/"), "/")
	segments := make([]string, len(encoded))
	for i := range encoded {
		segment := strings.ReplaceAll(encoded[i], "~1", "/")
		segment = strings.ReplaceAll(segment, "~0", "~")
		if strings.Contains(segment, "~") && strings.Contains(encoded[i], "~") {
			// Re-encoding catches invalid escape sequences while allowing a
			// correctly escaped literal '~'.
			reencoded := strings.ReplaceAll(strings.ReplaceAll(segment, "~", "~0"), "/", "~1")
			if reencoded != encoded[i] {
				return nil, fmt.Errorf("path %q contains an invalid JSON-pointer escape", value)
			}
		}
		segments[i] = segment
	}
	return segments, nil
}

func pathMatches(pattern, path []string) bool {
	if len(pattern) != len(path) {
		return false
	}
	for i := range pattern {
		if pattern[i] != "*" && pattern[i] != path[i] {
			return false
		}
	}
	return true
}

func appendPath(path []string, segment string) []string {
	out := make([]string, len(path), len(path)+1)
	copy(out, path)
	return append(out, segment)
}

func pointer(path []string) string {
	if len(path) == 0 {
		return ""
	}
	encoded := make([]string, len(path))
	for i := range path {
		encoded[i] = strings.ReplaceAll(strings.ReplaceAll(path[i], "~", "~0"), "/", "~1")
	}
	return "/" + strings.Join(encoded, "/")
}