/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */

package pipeline

import (
	"context"
	"regexp"
	"strings"

	"github.com/rotisserie/eris"
	kerrors "k8s.io/apimachinery/pkg/util/errors"

	"github.com/Azure/azure-service-operator/v2/tools/generator/internal/astmodel"
	"github.com/Azure/azure-service-operator/v2/tools/generator/internal/config"
)

// TransformCrossResourceReferencesStageID is the unique identifier for this pipeline stage
const TransformCrossResourceReferencesStageID = "transformCrossResourceReferences"

var idRegex = regexp.MustCompile("^(.*)I[d|D]s?$")

// TransformCrossResourceReferences replaces cross resource references with genruntime.ResourceReference.
func TransformCrossResourceReferences(configuration *config.Configuration, idFactory astmodel.IdentifierFactory) *Stage {
	return NewStage(
		TransformCrossResourceReferencesStageID,
		"Replace cross-resource references with genruntime.ResourceReference",
		func(ctx context.Context, state *State) (*State, error) {
			updatedDefs := make(astmodel.TypeDefinitionSet)

			visitor := MakeARMIDToResourceReferenceTypeVisitor(idFactory)
			for _, def := range state.Definitions() {
				// Skip Status types
				// TODO: we need flags
				if def.Name().IsStatus() {
					updatedDefs.Add(def)
					continue
				}

				conversionContext := ARMIDToReferenceTypeConverterContext{
					DefinitionName:        def.Name(),
					SelectedReferenceType: astmodel.ResourceReferenceType,
				}

				t, err := visitor.Visit(def.Type(), conversionContext)
				if err != nil {
					return nil, eris.Wrapf(err, "visiting %q", def.Name())
				}

				updatedDefs.Add(def.WithType(t))
			}

			resultDefs, err := stripARMIDPrimitiveTypes(updatedDefs)
			if err != nil {
				return nil, eris.Wrap(err, "failed to strip ARM ID primitive types")
			}

			return state.WithDefinitions(resultDefs), nil
		})
}

func makeReferencePropertyName(existing *astmodel.PropertyDefinition, isSlice bool, isMap bool) string {
	propertyNameSuffix := "Reference"
	if isSlice || isMap {
		propertyNameSuffix = "References"
	}

	var referencePropertyName string
	// Special case for "Id" and properties that end in "Id", which are quite common in the specs. This is primarily
	// because it's awkward to have a field called "Id" not just be a string and instead be a complex type describing
	// a reference.
	s := existing.PropertyName().String()

	if strings.ToLower(s) == "id" {
		referencePropertyName = propertyNameSuffix
	} else if idRegex.MatchString(s) {
		referencePropertyName = idRegex.ReplaceAllString(s, "${1}"+propertyNameSuffix)
	} else {
		referencePropertyName = s + propertyNameSuffix
	}

	return referencePropertyName
}

// makeLegacyReferencePropertyName does not correctly deal with properties with "ID" suffix, but exists
// to ensure backward compatibility with old versions.
// See https://github.com/Azure/azure-service-operator/issues/2501#issuecomment-1251650714
func makeLegacyReferencePropertyName(existing *astmodel.PropertyDefinition, isSlice bool, isMap bool) string {
	propertyNameSuffix := "Reference"
	if isSlice || isMap {
		propertyNameSuffix = "References"
	}

	var referencePropertyName string
	// Special case for "Id" and properties that end in "Id", which are quite common in the specs. This is primarily
	// because it's awkward to have a field called "Id" not just be a string and instead but a complex type describing
	// a reference.
	if existing.PropertyName() == "Id" {
		referencePropertyName = propertyNameSuffix
	} else if strings.HasSuffix(string(existing.PropertyName()), "Id") {
		referencePropertyName = strings.TrimSuffix(string(existing.PropertyName()), "Id") + propertyNameSuffix
	} else if strings.HasSuffix(string(existing.PropertyName()), "Ids") {
		referencePropertyName = strings.TrimSuffix(string(existing.PropertyName()), "Ids") + propertyNameSuffix
	} else {
		referencePropertyName = string(existing.PropertyName()) + propertyNameSuffix
	}

	return referencePropertyName
}

type ARMIDToReferenceTypeConverter struct {
	astmodel.TypeVisitor[ARMIDToReferenceTypeConverterContext]
	idFactory astmodel.IdentifierFactory
}

type ARMIDToReferenceTypeConverterContext struct {
	DefinitionName        astmodel.InternalTypeName
	SelectedReferenceType astmodel.ExternalTypeName
}

func (v *ARMIDToReferenceTypeConverter) transformARMIDToResourceReference(
	_ *astmodel.TypeVisitor[ARMIDToReferenceTypeConverterContext],
	it *astmodel.PrimitiveType,
	ctx ARMIDToReferenceTypeConverterContext,
) (astmodel.Type, error) {
	if it == astmodel.ARMIDType {
		return ctx.SelectedReferenceType, nil
	}

	return it, nil
}

func (v *ARMIDToReferenceTypeConverter) handleTransformationFlags(
	visitor *astmodel.TypeVisitor[ARMIDToReferenceTypeConverterContext],
	it *astmodel.FlaggedType,
	ctx ARMIDToReferenceTypeConverterContext,
) (astmodel.Type, error) {
	// If a Wellknown Resource Reference is required, use it and remove the flag
	if it.HasFlag(astmodel.WellKnownReferenceFlag) {
		ctx.SelectedReferenceType = astmodel.WellKnownResourceReferenceType
		elem, err := visitor.Visit(it.Element(), ctx)
		if err != nil {
			return nil, err
		}

		return it.WithElement(elem).WithoutFlag(astmodel.WellKnownReferenceFlag), nil
	}

	// If a Compatible Resource Reference is required, use it and remove the flag
	if it.HasFlag(astmodel.CompatibleReferenceFlag) {
		elem, err := visitor.Visit(it.Element(), ctx)
		if err != nil {
			return nil, err
		}

		return it.WithElement(elem).WithoutFlag(astmodel.CompatibleReferenceFlag), nil
	}

	elem, err := visitor.Visit(it.Element(), ctx)
	if err != nil {
		return nil, err
	}

	return it.WithElement(elem), nil
}

func (v *ARMIDToReferenceTypeConverter) renamePropertiesWithARMIDReferences(
	this *astmodel.TypeVisitor[ARMIDToReferenceTypeConverterContext],
	it *astmodel.ObjectType,
	ctx ARMIDToReferenceTypeConverterContext,
) (astmodel.Type, error) {
	// First, we visit this type like normal
	result, err := astmodel.IdentityVisitOfObjectType(this, it, ctx)
	if err != nil {
		return nil, err
	}

	ot, ok := result.(*astmodel.ObjectType)
	if !ok {
		return nil, eris.Errorf("result for visitObjectType of %s was not expected ObjectType, instead %T", ctx, result)
	}

	// Now, check if any properties have been updated and if they have change their names to match
	var newProps []*astmodel.PropertyDefinition
	var errs []error
	ot.Properties().ForEach(func(prop *astmodel.PropertyDefinition) {
		origProp, ok := it.Property(prop.PropertyName())
		if !ok {
			errs = append(errs, eris.Errorf("expected to find property %q on %s", prop.PropertyName(), ctx))
			return
		}

		if origProp.Equals(prop, astmodel.EqualityOverrides{}) {
			// No change to the property type
			newProps = append(newProps, prop)
			return
		}

		newProp := makeResourceReferenceProperty(ctx.DefinitionName, v.idFactory, prop)

		if astmodel.CompatibleReferenceFlag.IsOn(origProp.PropertyType()) {
			// If the original property is flagged for compatibility,
			// We keep both versions to avoid breaking existing CRs
			// but need to remove the flag first
			pr, err := v.createCompatibilityProperty(origProp)
			if err != nil {
				errs = append(errs, eris.Wrapf(err, "removing compatibility flag from property %s", origProp.PropertyName()))
			}

			newProps = append(newProps, pr)
		}

		newProps = append(newProps, newProp)
	})

	err = kerrors.NewAggregate(errs)
	if err != nil {
		return nil, err
	}

	ot = ot.WithoutProperties()
	ot = ot.WithProperties(newProps...)
	return ot, nil
}

func (v *ARMIDToReferenceTypeConverter) stripValidationForResourceReferences(
	this *astmodel.TypeVisitor[ARMIDToReferenceTypeConverterContext],
	it *astmodel.ValidatedType,
	ctx ARMIDToReferenceTypeConverterContext,
) (astmodel.Type, error) {
	result, err := astmodel.IdentityVisitOfValidatedType(this, it, ctx)
	if err != nil {
		return nil, err
	}

	if astmodel.TypeEquals(it, result) {
		return result, nil
	}

	// If they aren't equal, just return the element
	validated, ok := result.(*astmodel.ValidatedType)
	if !ok {
		return nil, eris.Errorf("expected IdentityVisitOfValidatedType to return a ValidatedType, but it instead returned %T", result)
	}
	return validated.ElementType(), nil
}

func (v *ARMIDToReferenceTypeConverter) createCompatibilityProperty(
	property *astmodel.PropertyDefinition,
) (*astmodel.PropertyDefinition, error) {
	t, err := astmodel.CompatibleReferenceFlag.RemoveFrom(property.PropertyType())
	if err != nil {
		return nil, eris.Wrapf(err, "removing compatibility flag from property %s", property.PropertyName())
	}

	return property.WithType(t), nil
}

func MakeARMIDToResourceReferenceTypeVisitor(
	idFactory astmodel.IdentifierFactory,
) ARMIDToReferenceTypeConverter {
	result := ARMIDToReferenceTypeConverter{
		idFactory: idFactory,
	}
	result.TypeVisitor = astmodel.TypeVisitorBuilder[ARMIDToReferenceTypeConverterContext]{
		VisitPrimitive:     result.transformARMIDToResourceReference,
		VisitObjectType:    result.renamePropertiesWithARMIDReferences,
		VisitValidatedType: result.stripValidationForResourceReferences,
		VisitFlaggedType:   result.handleTransformationFlags,
	}.Build()

	return result
}

func makeResourceReferenceProperty(
	typeName astmodel.InternalTypeName,
	idFactory astmodel.IdentifierFactory,
	existing *astmodel.PropertyDefinition,
) *astmodel.PropertyDefinition {
	_, isSlice := astmodel.AsArrayType(existing.PropertyType())
	_, isMap := astmodel.AsMapType(existing.PropertyType())

	var referencePropertyName string
	// This is hacky but works
	if group, version := typeName.InternalPackageReference().GroupVersion(); group == "containerservice" && strings.Contains(version, "20210501") {
		referencePropertyName = makeLegacyReferencePropertyName(existing, isSlice, isMap)
	} else {
		referencePropertyName = makeReferencePropertyName(existing, isSlice, isMap)
	}

	originalName := string(existing.PropertyName())

	newProp := astmodel.NewPropertyDefinition(
		idFactory.CreatePropertyName(referencePropertyName, astmodel.Exported),
		idFactory.CreateStringIdentifier(referencePropertyName, astmodel.NotExported),
		existing.PropertyType())
	// TODO: We could pass this information forward some other way?
	// Add tag so that we remember what this field was before
	newProp = newProp.WithTag(astmodel.ARMReferenceTag, originalName)

	newProp = newProp.WithDescription(existing.Description())
	if existing.IsRequired() {
		newProp = newProp.MakeRequired()
	} else {
		newProp = newProp.MakeOptional()
	}

	return newProp
}

func stripARMIDPrimitiveTypes(types astmodel.TypeDefinitionSet) (astmodel.TypeDefinitionSet, error) {
	// Any astmodel.ARMReference's which have escaped need to be turned into
	// strings
	armIDStrippingVisitor := astmodel.TypeVisitorBuilder[any]{
		VisitPrimitive: func(it *astmodel.PrimitiveType) (astmodel.Type, error) {
			if astmodel.TypeEquals(it, astmodel.ARMIDType) {
				return astmodel.StringType, nil
			}

			return it, nil
		},
	}.Build()

	result := make(astmodel.TypeDefinitionSet)
	for _, def := range types {
		newDef, err := armIDStrippingVisitor.VisitDefinition(def, nil)
		if err != nil {
			return nil, err
		}
		result.Add(newDef)
	}

	return result, nil
}
