// Package tier defines vertical capacity tier contracts.
package tier

import (
	"bytes"
	"fmt"
	"math"
	"strings"
	"text/template"
	"text/template/parse"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/yaml"

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

// Resource describes one effective resource capacity exposed by a tier.
type Resource struct {
	Name     string
	Capacity resource.Quantity
}

// Tier describes one vertical capacity tier for a Kubernetes target.
type Tier interface {
	// Name returns the tier name.
	Name() string

	// Resources returns all effective resource capacities defined by this tier.
	Resources() []Resource

	// Capacity returns the effective capacity for one resource.
	Capacity(resourceName string) (resource.Quantity, bool)

	// Patch returns the Kubernetes patch that applies this tier to the target.
	Patch() client.Patch

	// Cost is the cost of one replica at this tier.
	Cost() int64
}

type materializedTier struct {
	name           string
	resources      []Resource
	capacityByName map[string]resource.Quantity
	patch          client.Patch
	cost           int64
}

var _ Tier = (*materializedTier)(nil)

const maxRangeVariants = 10000

// New materializes Autoscaler tier inputs with predefined templates.
func New(templates *v1alpha1.TierTemplates, spec *v1alpha1.Tiers) ([]Tier, error) {
	compiledTemplates, err := compileTemplates(templates)
	if err != nil {
		return nil, err
	}
	if err := ValidateTiersSpec(spec); err != nil {
		return nil, err
	}

	tmpl, ok := compiledTemplates[spec.Template.Name]
	if !ok {
		return nil, fmt.Errorf("tier: template %q not found", spec.Template.Name)
	}

	generated := expandTiers(spec)
	tiers := make([]Tier, 0, len(generated))
	for i := range generated {
		t, err := materialize(spec, tmpl, &generated[i], i)
		if err != nil {
			return nil, err
		}
		tiers = append(tiers, t)
	}
	return tiers, nil
}

func (t *materializedTier) Name() string {
	return t.name
}

func (t *materializedTier) Resources() []Resource {
	out := make([]Resource, len(t.resources))
	for i := range t.resources {
		out[i] = Resource{
			Name:     t.resources[i].Name,
			Capacity: t.resources[i].Capacity.DeepCopy(),
		}
	}
	return out
}

func (t *materializedTier) Capacity(resourceName string) (resource.Quantity, bool) {
	capacity, ok := t.capacityByName[resourceName]
	if !ok {
		return resource.Quantity{}, false
	}
	return capacity.DeepCopy(), true
}

func (t *materializedTier) Patch() client.Patch {
	return t.patch
}

func (t *materializedTier) Cost() int64 {
	return t.cost
}

func templatesByName(templates *v1alpha1.TierTemplates) (map[string]v1alpha1.TierTemplate, error) {
	if templates == nil {
		return nil, fmt.Errorf("tier: nil templates")
	}

	byName := make(map[string]v1alpha1.TierTemplate, len(templates.Templates))
	for name, tmpl := range templates.Templates {
		if name == "" {
			return nil, fmt.Errorf("tier: template name is required")
		}
		byName[name] = tmpl
	}
	return byName, nil
}

type compiledTemplate struct {
	tmpl             v1alpha1.TierTemplate
	matchers         map[string]*v1alpha1.TierMatcher
	params           map[string]*v1alpha1.TierParam
	available        []*v1alpha1.TierAvailable
	costs            []*v1alpha1.TierCost
	referencedParams map[string]struct{}
}

func compileTemplates(templates *v1alpha1.TierTemplates) (map[string]*compiledTemplate, error) {
	templateByName, err := templatesByName(templates)
	if err != nil {
		return nil, err
	}

	compiled := make(map[string]*compiledTemplate, len(templateByName))
	for name, tmpl := range templateByName {
		matchers, err := matchersByName(&tmpl)
		if err != nil {
			return nil, fmt.Errorf("tier template %q: %w", name, err)
		}
		params, err := paramsByName(&tmpl, matchers)
		if err != nil {
			return nil, fmt.Errorf("tier template %q: %w", name, err)
		}
		referencedParams, err := referencedPatchTemplateParams(&tmpl, params)
		if err != nil {
			return nil, fmt.Errorf("tier template %q: %w", name, err)
		}
		available, err := compileAvailable(&tmpl, matchers)
		if err != nil {
			return nil, fmt.Errorf("tier template %q: %w", name, err)
		}
		costs, err := costsByName(&tmpl, matchers)
		if err != nil {
			return nil, fmt.Errorf("tier template %q: %w", name, err)
		}
		compiled[name] = &compiledTemplate{
			tmpl:             tmpl,
			matchers:         matchers,
			params:           params,
			available:        available,
			costs:            costs,
			referencedParams: referencedParams,
		}
	}
	return compiled, nil
}

func materialize(spec *v1alpha1.Tiers, tmpl *compiledTemplate, generated *generatedTier, index int) (*materializedTier, error) {
	values, err := materializeParams(tmpl.matchers, tmpl.params, tmpl.referencedParams, generated.values)
	if err != nil {
		return nil, err
	}
	cost, err := materializeCost(tmpl.matchers, tmpl.costs, generated.values)
	if err != nil {
		return nil, err
	}
	available, err := materializeAvailable(tmpl.matchers, tmpl.available, generated.values)
	if err != nil {
		return nil, err
	}

	data, err := renderPatchTemplate(&tmpl.tmpl, values)
	if err != nil {
		return nil, fmt.Errorf("tier %q-%d: %w", spec.Template.Name, index, err)
	}
	resources, capacityByName := materializedResources(generated.resources, available)

	return &materializedTier{
		name:           fmt.Sprintf("%s-%d", spec.Template.Name, index),
		resources:      resources,
		capacityByName: capacityByName,
		patch:          client.RawPatch(types.MergePatchType, data),
		cost:           cost,
	}, nil
}

func materializedResources(resources []Resource, available map[string]int32) ([]Resource, map[string]resource.Quantity) {
	out := make([]Resource, len(resources))
	capacityByName := make(map[string]resource.Quantity, len(resources))
	for i := range resources {
		capacity := resources[i].Capacity.DeepCopy()
		if percentage, ok := available[resources[i].Name]; ok {
			capacity = capacityAtPercentage(capacity, percentage)
		}
		out[i] = Resource{Name: resources[i].Name, Capacity: capacity}
		capacityByName[resources[i].Name] = capacity
	}
	return out, capacityByName
}

func capacityAtPercentage(capacity resource.Quantity, percentage int32) resource.Quantity {
	percentageQuantity := resource.NewScaledQuantity(int64(percentage), -2)
	value := capacity.AsDec()
	value.Mul(value, percentageQuantity.AsDec())
	return capacity
}

type generatedTier struct {
	resources []Resource
	values    map[string]resource.Quantity
}

func expandTiers(spec *v1alpha1.Tiers) []generatedTier {
	var tiers []generatedTier
	for i := range spec.Sets {
		setTiers := expandSet(&spec.Sets[i])
		tiers = append(tiers, setTiers...)
	}
	return tiers
}

func expandSet(set *v1alpha1.TierSet) []generatedTier {
	tiers := []generatedTier{{values: map[string]resource.Quantity{}}}
	for i := range set.Dimensions {
		dimensionTiers := expandDimension(&set.Dimensions[i])
		tiers = product(tiers, dimensionTiers)
	}
	return tiers
}

func expandDimension(dimension *v1alpha1.TierDimension) []generatedTier {
	resourceVariants := make([][]resource.Quantity, len(dimension.Resources))
	for i := range dimension.Resources {
		input := &dimension.Resources[i]
		resourceVariants[i] = variantsForResource(input)
	}

	variantCount := len(resourceVariants[0])
	tiers := make([]generatedTier, 0, variantCount)
	for variantIndex := 0; variantIndex < variantCount; variantIndex++ {
		tier := generatedTier{
			resources: make([]Resource, 0, len(dimension.Resources)),
			values:    make(map[string]resource.Quantity, len(dimension.Resources)),
		}
		for resourceIndex := range dimension.Resources {
			name := string(dimension.Resources[resourceIndex].Name)
			value := resourceVariants[resourceIndex][variantIndex].DeepCopy()
			tier.resources = append(tier.resources, Resource{Name: name, Capacity: value})
			tier.values[name] = value
		}
		tiers = append(tiers, tier)
	}
	return tiers
}

func variantsForResource(input *v1alpha1.TierDimensionResource) []resource.Quantity {
	if len(input.Variants) > 0 {
		variants := make([]resource.Quantity, len(input.Variants))
		for i := range input.Variants {
			variants[i] = input.Variants[i].DeepCopy()
		}
		return variants
	}
	return expandRange(input.Range)
}

func expandRange(r *v1alpha1.TierResourceRange) []resource.Quantity {
	min := r.Min.DeepCopy()
	max := r.Max.DeepCopy()
	step := r.Step.DeepCopy()

	var variants []resource.Quantity
	current := min.DeepCopy()
	for range maxRangeVariants {
		cmp := current.Cmp(max)
		variants = append(variants, current.DeepCopy())
		if cmp == 0 {
			return variants
		}
		current.Add(step)
	}
	return variants
}

func product(left, right []generatedTier) []generatedTier {
	tiers := make([]generatedTier, 0, len(left)*len(right))
	for i := range left {
		for j := range right {
			tiers = append(tiers, mergeGeneratedTiers(&left[i], &right[j]))
		}
	}
	return tiers
}

func mergeGeneratedTiers(left, right *generatedTier) generatedTier {
	merged := generatedTier{
		resources: make([]Resource, 0, len(left.resources)+len(right.resources)),
		values:    make(map[string]resource.Quantity, len(left.values)+len(right.values)),
	}
	for _, side := range []*generatedTier{left, right} {
		for i := range side.resources {
			item := side.resources[i]
			value := item.Capacity.DeepCopy()
			merged.resources = append(merged.resources, Resource{Name: item.Name, Capacity: value})
			merged.values[item.Name] = value
		}
	}
	return merged
}

func materializeParams(matchers map[string]*v1alpha1.TierMatcher, params map[string]*v1alpha1.TierParam, referencedParams map[string]struct{}, resourceValues map[string]resource.Quantity) (map[string]string, error) {
	values := make(map[string]string, len(referencedParams))
	for paramName := range referencedParams {
		param := params[paramName]
		matcher := matchers[param.Matcher]
		index, err := matchValue(matcher, resourceValues)
		if err != nil {
			return nil, fmt.Errorf("param %q: %w", param.Name, err)
		}
		values[param.Name] = param.Values[index]
	}
	return values, nil
}

func materializeCost(matchers map[string]*v1alpha1.TierMatcher, costs []*v1alpha1.TierCost, resourceValues map[string]resource.Quantity) (int64, error) {
	var total int64
	for i := range costs {
		cost := costs[i]
		matcher := matchers[cost.Matcher]
		index, err := matchValue(matcher, resourceValues)
		if err != nil {
			return 0, fmt.Errorf("cost %q: %w", cost.Name, err)
		}
		total += int64(cost.Values[index])
	}
	return total, nil
}

func materializeAvailable(matchers map[string]*v1alpha1.TierMatcher, available []*v1alpha1.TierAvailable, resourceValues map[string]resource.Quantity) (map[string]int32, error) {
	percentages := make(map[string]int32, len(available))
	for i := range available {
		item := available[i]
		resourceName := string(item.Name)
		if _, ok := resourceValues[resourceName]; !ok {
			return nil, fmt.Errorf("available %q: resource %q is missing", resourceName, resourceName)
		}

		matcher := matchers[item.Matcher]
		index, err := matchValue(matcher, resourceValues)
		if err != nil {
			return nil, fmt.Errorf("available %q: %w", resourceName, err)
		}
		percentages[resourceName] = item.Percentages[index]
	}
	return percentages, nil
}

func renderPatchTemplate(tmpl *v1alpha1.TierTemplate, params map[string]string) ([]byte, error) {
	if strings.TrimSpace(tmpl.Template) == "" {
		return nil, fmt.Errorf("patch template is required")
	}

	funcs := template.FuncMap{}
	for name, value := range params {
		funcs[name] = func(value string) func() string {
			return func() string {
				return value
			}
		}(value)
	}

	parsed, err := template.New("patch").Funcs(funcs).Option("missingkey=error").Parse(tmpl.Template)
	if err != nil {
		return nil, fmt.Errorf("parse patch template: %w", err)
	}

	allowed := stringSetFromValues(params)
	if err := collectPatchTemplateParams(parsed.Root, allowed, nil); err != nil {
		return nil, err
	}

	var buf bytes.Buffer
	if err := parsed.Execute(&buf, params); err != nil {
		return nil, fmt.Errorf("render patch template: %w", err)
	}

	data, err := yaml.YAMLToJSON(buf.Bytes())
	if err != nil {
		return nil, fmt.Errorf("convert patch template to JSON: %w", err)
	}
	return data, nil
}

func collectPatchTemplateParams(node parse.Node, allowed map[string]struct{}, referenced map[string]struct{}) error {
	switch n := node.(type) {
	case nil:
		return nil
	case *parse.ListNode:
		for _, child := range n.Nodes {
			if err := collectPatchTemplateParams(child, allowed, referenced); err != nil {
				return err
			}
		}
		return nil
	case *parse.TextNode:
		return nil
	case *parse.ActionNode:
		return collectPatchTemplateAction(n, allowed, referenced)
	default:
		return fmt.Errorf("patch template may only contain simple param substitutions")
	}
}

func collectPatchTemplateAction(node *parse.ActionNode, allowed map[string]struct{}, referenced map[string]struct{}) error {
	if node.Pipe == nil || len(node.Pipe.Cmds) != 1 {
		return fmt.Errorf("patch template action must reference exactly one param")
	}
	cmd := node.Pipe.Cmds[0]
	if len(cmd.Args) != 1 {
		return fmt.Errorf("patch template action must reference exactly one param")
	}

	switch arg := cmd.Args[0].(type) {
	case *parse.FieldNode:
		if len(arg.Ident) != 1 {
			return fmt.Errorf("patch template action must use a simple param reference")
		}
		if _, ok := allowed[arg.Ident[0]]; !ok {
			return fmt.Errorf("patch template references undefined param %q", arg.Ident[0])
		}
		if referenced != nil {
			referenced[arg.Ident[0]] = struct{}{}
		}
	case *parse.IdentifierNode:
		if _, ok := allowed[arg.Ident]; !ok {
			return fmt.Errorf("patch template references undefined param %q", arg.Ident)
		}
		if referenced != nil {
			referenced[arg.Ident] = struct{}{}
		}
	default:
		return fmt.Errorf("patch template action must use a simple param reference")
	}
	return nil
}

func matchersByName(tmpl *v1alpha1.TierTemplate) (map[string]*v1alpha1.TierMatcher, error) {
	matchers := make(map[string]*v1alpha1.TierMatcher, len(tmpl.Matchers))
	for i := range tmpl.Matchers {
		matcher := &tmpl.Matchers[i]
		if matcher.Name == "" {
			return nil, fmt.Errorf("matcher name is required")
		}
		if err := validateMatcher(matcher); err != nil {
			return nil, err
		}
		if _, exists := matchers[matcher.Name]; exists {
			return nil, fmt.Errorf("duplicate matcher %q", matcher.Name)
		}
		matchers[matcher.Name] = matcher
	}
	return matchers, nil
}

func paramsByName(tmpl *v1alpha1.TierTemplate, matchers map[string]*v1alpha1.TierMatcher) (map[string]*v1alpha1.TierParam, error) {
	params := make(map[string]*v1alpha1.TierParam, len(tmpl.Params))
	for i := range tmpl.Params {
		param := &tmpl.Params[i]
		if param.Name == "" {
			return nil, fmt.Errorf("param name is required")
		}
		if err := validateParam(param, matchers); err != nil {
			return nil, err
		}
		if _, exists := params[param.Name]; exists {
			return nil, fmt.Errorf("duplicate param %q", param.Name)
		}
		params[param.Name] = param
	}
	return params, nil
}

func compileAvailable(tmpl *v1alpha1.TierTemplate, matchers map[string]*v1alpha1.TierMatcher) ([]*v1alpha1.TierAvailable, error) {
	available := make([]*v1alpha1.TierAvailable, 0, len(tmpl.Available))
	names := make(map[corev1.ResourceName]struct{}, len(tmpl.Available))
	for i := range tmpl.Available {
		item := &tmpl.Available[i]
		if item.Name == "" {
			return nil, fmt.Errorf("available name is required")
		}
		if _, exists := names[item.Name]; exists {
			return nil, fmt.Errorf("duplicate available resource %q", item.Name)
		}
		if err := validateAvailable(item, matchers); err != nil {
			return nil, err
		}
		names[item.Name] = struct{}{}
		available = append(available, item)
	}
	return available, nil
}

func costsByName(tmpl *v1alpha1.TierTemplate, matchers map[string]*v1alpha1.TierMatcher) ([]*v1alpha1.TierCost, error) {
	if len(tmpl.Costs) == 0 {
		return nil, fmt.Errorf("costs are required")
	}
	// Cost values are non-negative int32 values. Capping the component count at
	// MaxInt32 guarantees their maximum possible sum is less than MaxInt64.
	if len(tmpl.Costs) > math.MaxInt32 {
		return nil, fmt.Errorf("costs has %d items, maximum is %d", len(tmpl.Costs), math.MaxInt32)
	}

	costs := make([]*v1alpha1.TierCost, 0, len(tmpl.Costs))
	names := make(map[string]struct{}, len(tmpl.Costs))
	for i := range tmpl.Costs {
		cost := &tmpl.Costs[i]
		if cost.Name == "" {
			return nil, fmt.Errorf("cost name is required")
		}
		if _, exists := names[cost.Name]; exists {
			return nil, fmt.Errorf("duplicate cost %q", cost.Name)
		}
		if err := validateCost(cost, matchers); err != nil {
			return nil, err
		}
		names[cost.Name] = struct{}{}
		costs = append(costs, cost)
	}
	return costs, nil
}

func referencedPatchTemplateParams(tmpl *v1alpha1.TierTemplate, params map[string]*v1alpha1.TierParam) (map[string]struct{}, error) {
	if strings.TrimSpace(tmpl.Template) == "" {
		return nil, fmt.Errorf("patch template is required")
	}

	funcs := template.FuncMap{}
	allowed := make(map[string]struct{}, len(params))
	for name := range params {
		allowed[name] = struct{}{}
		funcs[name] = func() string {
			return ""
		}
	}

	parsed, err := template.New("patch").Funcs(funcs).Option("missingkey=error").Parse(tmpl.Template)
	if err != nil {
		return nil, fmt.Errorf("parse patch template: %w", err)
	}

	referenced := map[string]struct{}{}
	if err := collectPatchTemplateParams(parsed.Root, allowed, referenced); err != nil {
		return nil, err
	}
	return referenced, nil
}

func stringSetFromValues(values map[string]string) map[string]struct{} {
	set := make(map[string]struct{}, len(values))
	for name := range values {
		set[name] = struct{}{}
	}
	return set
}

func validateMatcher(matcher *v1alpha1.TierMatcher) error {
	switch matcherType(matcher) {
	case v1alpha1.TierMatcherTypeResource:
		return validateResourceMatcher(matcher)
	default:
		return fmt.Errorf("matcher %q type %q is unsupported", matcher.Name, matcher.Type)
	}
}

func matcherType(matcher *v1alpha1.TierMatcher) v1alpha1.TierMatcherType {
	if matcher.Type == "" {
		return v1alpha1.TierMatcherTypeResource
	}
	return matcher.Type
}

func validateResourceMatcher(matcher *v1alpha1.TierMatcher) error {
	resourceMatcher := matcher.Resource
	if resourceMatcher == nil {
		return fmt.Errorf("matcher %q resource is required", matcher.Name)
	}
	if resourceMatcher.Name == "" {
		return fmt.Errorf("matcher %q resource.name is required", matcher.Name)
	}
	if len(resourceMatcher.Variants) == 0 {
		return fmt.Errorf("matcher %q resource.variants are required", matcher.Name)
	}
	for i := range resourceMatcher.Variants {
		for j := i + 1; j < len(resourceMatcher.Variants); j++ {
			if resourceMatcher.Variants[i].Cmp(resourceMatcher.Variants[j]) == 0 {
				return fmt.Errorf("matcher %q resource.variants contains duplicate value %q", matcher.Name, resourceMatcher.Variants[j].String())
			}
		}
	}
	return nil
}

func validateParam(param *v1alpha1.TierParam, matchers map[string]*v1alpha1.TierMatcher) error {
	if param.Matcher == "" {
		return fmt.Errorf("param %q matcher is required", param.Name)
	}
	matcher, ok := matchers[param.Matcher]
	if !ok {
		return fmt.Errorf("param %q references undefined matcher %q", param.Name, param.Matcher)
	}
	if len(param.Values) != matcherValueCount(matcher) {
		return fmt.Errorf("param %q has %d values, want %d", param.Name, len(param.Values), matcherValueCount(matcher))
	}
	return nil
}

func validateAvailable(available *v1alpha1.TierAvailable, matchers map[string]*v1alpha1.TierMatcher) error {
	if available.Matcher == "" {
		return fmt.Errorf("available %q matcher is required", available.Name)
	}
	matcher, ok := matchers[available.Matcher]
	if !ok {
		return fmt.Errorf("available %q references undefined matcher %q", available.Name, available.Matcher)
	}
	if len(available.Percentages) != matcherValueCount(matcher) {
		return fmt.Errorf("available %q has %d percentages, want %d", available.Name, len(available.Percentages), matcherValueCount(matcher))
	}
	for i := range available.Percentages {
		percentage := available.Percentages[i]
		if percentage <= 0 || percentage > 100 {
			return fmt.Errorf("available %q percentage at index %d must be greater than 0 and less than or equal to 100", available.Name, i)
		}
	}
	return nil
}

func validateCost(cost *v1alpha1.TierCost, matchers map[string]*v1alpha1.TierMatcher) error {
	if cost.Matcher == "" {
		return fmt.Errorf("cost %q matcher is required", cost.Name)
	}
	matcher, ok := matchers[cost.Matcher]
	if !ok {
		return fmt.Errorf("cost %q references undefined matcher %q", cost.Name, cost.Matcher)
	}
	if len(cost.Values) != matcherValueCount(matcher) {
		return fmt.Errorf("cost %q has %d values, want %d", cost.Name, len(cost.Values), matcherValueCount(matcher))
	}
	for i := range cost.Values {
		if cost.Values[i] < 0 {
			return fmt.Errorf("cost %q value at index %d must not be negative", cost.Name, i)
		}
	}
	return nil
}

func matcherValueCount(matcher *v1alpha1.TierMatcher) int {
	switch matcherType(matcher) {
	case v1alpha1.TierMatcherTypeResource:
		return len(matcher.Resource.Variants)
	default:
		return 0
	}
}

func matchValue(matcher *v1alpha1.TierMatcher, resourceValues map[string]resource.Quantity) (int, error) {
	switch matcherType(matcher) {
	case v1alpha1.TierMatcherTypeResource:
		return matchResourceValue(matcher, resourceValues)
	default:
		return 0, fmt.Errorf("matcher %q type %q is unsupported", matcher.Name, matcher.Type)
	}
}

func matchResourceValue(matcher *v1alpha1.TierMatcher, resourceValues map[string]resource.Quantity) (int, error) {
	resourceMatcher := matcher.Resource
	resourceName := string(resourceMatcher.Name)
	value, ok := resourceValues[resourceName]
	if !ok {
		return 0, fmt.Errorf("resource %q is missing", resourceName)
	}

	for i := range resourceMatcher.Variants {
		if value.Cmp(resourceMatcher.Variants[i]) == 0 {
			return i, nil
		}
	}
	return 0, fmt.Errorf("resource %q value %q does not match any variant", resourceName, value.String())
}