/*
   Copyright The containerd Authors.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package main

import (
	"flag"
	"fmt"
	"strings"
	"unicode"

	"google.golang.org/protobuf/compiler/protogen"
	"google.golang.org/protobuf/reflect/protoreflect"
)

const (
	version       = "0.0.1"
	licenseHeader = `/*
   Copyright The containerd Authors.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/`
)

var stripGoFilename string

// The root messages for which Strip() functions should be generated, including their children.
var rootMessages = map[string]bool{
	"ContainerAdjustment": true,
	"ContainerUpdate":     true,
}

// findReachableMessages performs a BFS traversal of the message graph to find
// all messages reachable from the root messages.
func findReachableMessages(messages []*protogen.Message) map[string]bool {
	reachable := make(map[string]bool)
	queue := []*protogen.Message{}

	messageMap := make(map[string]*protogen.Message)
	for _, m := range messages {
		messageMap[string(m.Desc.Name())] = m
	}

	for name := range rootMessages {
		if m, ok := messageMap[name]; ok {
			queue = append(queue, m)
		}
	}

	for len(queue) > 0 {
		m := queue[0]
		queue = queue[1:]

		if reachable[string(m.Desc.Name())] {
			continue
		}
		reachable[string(m.Desc.Name())] = true

		for _, field := range m.Fields {
			// Only traverse direct, non-repeated, non-map message fields.
			if field.Message != nil && !field.Desc.IsList() && !field.Desc.IsMap() {
				if !isSimpleOptional(field.Message) {
					queue = append(queue, field.Message)
				}
			}
		}
	}
	return reachable
}

func isSimpleOptional(m *protogen.Message) bool {
	return len(m.Fields) == 1 && m.Fields[0].GoName == "Value" && m.Fields[0].Message == nil
}

func main() {
	var flags flag.FlagSet
	flags.StringVar(&stripGoFilename, "file", "strip.go", "output filename for strip functions")

	protogen.Options{
		ParamFunc: flags.Set,
	}.Run(func(gen *protogen.Plugin) error {
		for _, f := range gen.Files {
			if !f.Generate {
				continue
			}
			generateFile(gen, f)
		}
		return nil
	})
}

// generateFile generates a .go file containing strip functions.
func generateFile(gen *protogen.Plugin, file *protogen.File) *protogen.GeneratedFile {
	if len(file.Messages) == 0 {
		return nil
	}

	// Determine output filename: use `stripGoFilename` in the `file.GoImportPath` directory.
	// For example, if file.GoImportPath is "github.com/containerd/nri/pkg/api" and stripGoFilename is "strip.go",
	// the output path will be "github.com/containerd/nri/pkg/api/strip.go".
	// The `protoc` command will need to specify the output directory.
	filename := stripGoFilename
	g := gen.NewGeneratedFile(filename, file.GoImportPath)

	for _, line := range strings.Split(licenseHeader, "\n") {
		g.P(line)
	}
	g.P()
	g.P("// Code generated by tools/protoc-gen-strip. DO NOT EDIT.")
	g.P("// versions:")
	g.P(fmt.Sprintf("// 	protoc-gen-strip %s", version))
	g.P(fmt.Sprintf("// 	protoc           %s", protocVersion(gen)))
	g.P()
	g.P("package ", file.GoPackageName)
	g.P()

	// Collect all non-map-entry messages from the file.
	var allMessages []*protogen.Message
	for _, m := range file.Messages {
		if !m.Desc.IsMapEntry() {
			allMessages = append(allMessages, m)
		}
	}

	// Find all messages reachable from the root messages.
	reachable := findReachableMessages(allMessages)

	// Generate Strip() functions only for the reachable messages.
	for _, m := range allMessages {
		if reachable[string(m.Desc.Name())] {
			generateStripFunc(g, m)
		}
	}

	return g
}

func generateStripFunc(g *protogen.GeneratedFile, m *protogen.Message) {
	receiverName := generateReceiverName(m.GoIdent.GoName)

	g.P("// Strip empty fields from ", m.GoIdent.GoName, ", reducing a fully empty")
	g.P("// one to nil. Strip allows comparison of two ", m.GoIdent.GoName, " structs for semantic")
	g.P("// equality using go-cmp.")
	g.P("func (", receiverName, " *", m.GoIdent.GoName, ") Strip() *", m.GoIdent.GoName, " {")
	g.P("	if ", receiverName, " == nil {")
	g.P("		return nil")
	g.P("	}")
	g.P()
	g.P("	empty := true")
	g.P()

	for _, field := range m.Fields {
		goName := field.GoName

		if field.Desc.IsMap() {
			g.P("	if len(", receiverName, ".", goName, ") == 0 {")
			g.P("			", receiverName, ".", goName, " = nil")
			g.P("		} else {")
			g.P("			empty = false")
			g.P("		}")
		} else if field.Desc.IsList() {
			g.P("	if len(", receiverName, ".", goName, ") == 0 {")
			g.P("			", receiverName, ".", goName, " = nil")
			g.P("		} else {")
			g.P("			empty = false")
			g.P("		}")
		} else if field.Message != nil {
			if isSimpleOptional(field.Message) {
				g.P("	if ", receiverName, ".", goName, " != nil { empty = false }")
			} else {
				g.P("	if ", receiverName, ".", goName, " = ", receiverName, ".", goName, ".Strip(); ", receiverName, ".", goName, " != nil {")
				g.P("			empty = false")
				g.P("		}")
			}
		} else {
			switch field.Desc.Kind() {
			case protoreflect.BoolKind:
				g.P("	if ", receiverName, ".", goName, " != false { empty = false }")
			case protoreflect.StringKind:
				g.P("	if ", receiverName, ".", goName, " != \"\" { empty = false }")
			case protoreflect.Int32Kind, protoreflect.Uint32Kind,
				protoreflect.Int64Kind, protoreflect.Uint64Kind,
				protoreflect.Sint32Kind, protoreflect.Sint64Kind,
				protoreflect.Fixed32Kind, protoreflect.Fixed64Kind,
				protoreflect.Sfixed32Kind, protoreflect.Sfixed64Kind,
				protoreflect.EnumKind:
				g.P("	if ", receiverName, ".", goName, " != 0 { empty = false }")
			case protoreflect.BytesKind:
				g.P("	if len(", receiverName, ".", goName, ") != 0 { empty = false }")
			case protoreflect.FloatKind, protoreflect.DoubleKind:
				g.P("	if ", receiverName, ".", goName, " != 0.0 { empty = false }")
			default:
				g.P("	// TODO(protoc-gen-strip): Handle field type ", field.Desc.Kind().String(), " for ", goName, " (unhandled primitive)")
				g.P("	// Fallback check assuming it's a non-nil type if non-zero")
				g.P("	if ", receiverName, ".", goName, " != ", goZeroValueString(field.Desc.Kind()), " { empty = false }")
			}
		}
	}

	g.P()
	g.P("	if empty {")
	g.P("		return nil")
	g.P("	}")
	g.P()
	g.P("	return ", receiverName)
	g.P("}")
	g.P()
}

func goZeroValueString(kind protoreflect.Kind) string {
	switch kind {
	case protoreflect.BoolKind:
		return "false"
	case protoreflect.StringKind:
		return `""`
	case protoreflect.BytesKind: // []byte is nil as zero value
		return "nil"
	case protoreflect.MessageKind: // *MyMessage is nil as zero value
		return "nil"
	case protoreflect.EnumKind,
		protoreflect.Int32Kind, protoreflect.Uint32Kind,
		protoreflect.Int64Kind, protoreflect.Uint64Kind,
		protoreflect.Sint32Kind, protoreflect.Sint64Kind,
		protoreflect.Fixed32Kind, protoreflect.Fixed64Kind,
		protoreflect.Sfixed32Kind, protoreflect.Sfixed64Kind:
		return "0"
	case protoreflect.FloatKind, protoreflect.DoubleKind:
		return "0.0"
	default:
		return "nil" // Fallback, should not happen for standard types
	}
}

func protocVersion(gen *protogen.Plugin) string {
	if v := gen.Request.GetCompilerVersion(); v != nil {
		return fmt.Sprintf("v%d.%d.%d-%s", v.GetMajor(), v.GetMinor(), v.GetPatch(), v.GetSuffix())
	}
	return "(unknown)"
}

// generateReceiverName generates a receiver name for a given type.
func generateReceiverName(typeName string) string {
	if typeName == "" {
		return "x"
	}

	// Store the first letter, lowercased.
	firstLetter := strings.ToLower(string(typeName[0]))

	// Scan forward for an uppercase letter after the first character.
	for i := 1; i < len(typeName); i++ {
		if unicode.IsUpper(rune(typeName[i])) {
			return strings.ToLower(string(typeName[i]))
		}
	}

	// If no other uppercase letter is found, return the lowercased first letter.
	return firstLetter
}