/*
   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"

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

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 ownersGoFilename string

// Hardcoded map to distinguish simple vs compound fields.
var compoundFields = map[string]bool{
	"Annotations":     true,
	"Mounts":          true,
	"Devices":         true,
	"CdiDevices":      true,
	"Env":             true,
	"HugepageLimits":  true,
	"CgroupsUnified":  true,
	"Rlimits":         true,
	"Namespace":       true,
	"Sysctl":          true,
	"LinuxNetDevices": true,
}

var compoundKeys = map[string]string{
	"Annotations":     "key",
	"Mounts":          "destination",
	"Devices":         "path",
	"CdiDevices":      "name",
	"Env":             "name",
	"HugepageLimits":  "size",
	"CgroupsUnified":  "key",
	"Rlimits":         "typ",
	"Namespace":       "typ",
	"Sysctl":          "key",
	"LinuxNetDevices": "path",
}

func main() {
	var flags flag.FlagSet
	flags.StringVar(&ownersGoFilename, "file", "owners_generated.go", "output filename for owner 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 owner functions.
func generateFile(gen *protogen.Plugin, file *protogen.File) *protogen.GeneratedFile {
	var fieldEnum *protogen.Enum
	for _, e := range file.Enums {
		if e.GoIdent.GoName == "Field" {
			fieldEnum = e
			break
		}
	}
	if fieldEnum == nil {
		return nil
	}

	filename := ownersGoFilename
	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-owners. DO NOT EDIT.")
	g.P("// versions:")
	g.P(fmt.Sprintf("// 	protoc-gen-owners %s", version))
	g.P(fmt.Sprintf("// 	protoc            %s", protocVersion(gen)))
	g.P()
	g.P("package ", file.GoPackageName)
	g.P()

	generateOwnerFuncs(g, fieldEnum)

	return g
}

var skipFields = map[string]bool{
	"OciHooks": true,
}

var nameOverrides = map[string]string{
	"Annotations":     "Annotation",
	"Mounts":          "Mount",
	"Devices":         "Device",
	"HugepageLimits":  "HugepageLimit",
	"LinuxNetDevices": "LinuxNetDevice",
	"CdiDevices":      "CdiDevice",
	"IoPriority":      "IOPriority",
	"LinuxSched":      "LinuxScheduler",
	"Rlimits":         "Rlimit",
}

func generateOwnerFuncs(g *protogen.GeneratedFile, fieldEnum *protogen.Enum) {
	for _, enumValue := range fieldEnum.Values {
		enumValueName := enumValue.GoIdent.GoName
		if enumValueName == "Field_None" || enumValueName == "None" {
			continue
		}
		originalFieldName := strings.TrimPrefix(enumValueName, "Field_")
		if skipFields[originalFieldName] {
			continue
		}
		fieldName := originalFieldName
		if override, ok := nameOverrides[originalFieldName]; ok {
			fieldName = override
		}

		isCompound := compoundFields[originalFieldName]
		keyName := compoundKeys[originalFieldName]

		if isCompound {
			// OwningPlugins.Claim<Field>
			g.P("func (o *OwningPlugins) Claim", fieldName, "(id, ", keyName, ", plugin string) error {")
			g.P("    return o.mustOwnersFor(id).Claim", fieldName, "(", keyName, ", plugin)")
			g.P("}")
			g.P()

			// FieldOwners.Claim<Field>
			g.P("func (f *FieldOwners) Claim", fieldName, "(", keyName, ", plugin string) error {")
			g.P("    return f.claimCompound(Field_", originalFieldName, ".Key(), ", keyName, ", plugin)")
			g.P("}")
			g.P()

			// OwningPlugins.<Field>Owner
			g.P("func (o *OwningPlugins) ", fieldName, "Owner(id, ", keyName, " string) (string, bool) {")
			g.P("    return o.ownersFor(id).compoundOwner(Field_", originalFieldName, ".Key(), ", keyName, ")")
			g.P("}")
			g.P()

			// FieldOwners.<Field>Owner
			g.P("func (f *FieldOwners) ", fieldName, "Owner(", keyName, " string) (string, bool) {")
			g.P("    return f.compoundOwner(Field_", originalFieldName, ".Key(), ", keyName, ")")
			g.P("}")
			g.P()

			// OwningPlugins.Clear<Field>
			g.P("func (o *OwningPlugins) Clear", fieldName, "(id, ", keyName, ", plugin string) {")
			g.P("    o.mustOwnersFor(id).Clear", fieldName, "(", keyName, ", plugin)")
			g.P("}")
			g.P()

			// FieldOwners.Clear<Field>
			g.P("func (f *FieldOwners) Clear", fieldName, "(", keyName, ", plugin string) {")
			g.P("    f.clearCompound(Field_", originalFieldName, ".Key(), ", keyName, ", plugin)")
			g.P("}")
			g.P()
		} else {
			// OwningPlugins.Claim<Field>
			g.P("func (o *OwningPlugins) Claim", fieldName, "(id, plugin string) error {")
			g.P("    return o.mustOwnersFor(id).Claim", fieldName, "(plugin)")
			g.P("}")
			g.P()

			// FieldOwners.Claim<Field>
			g.P("func (f *FieldOwners) Claim", fieldName, "(plugin string) error {")
			g.P("    return f.claimSimple(Field_", originalFieldName, ".Key(), plugin)")
			g.P("}")
			g.P()

			// OwningPlugins.<Field>Owner
			g.P("func (o *OwningPlugins) ", fieldName, "Owner(id string) (string, bool) {")
			g.P("    return o.ownersFor(id).simpleOwner(Field_", originalFieldName, ".Key())")
			g.P("}")
			g.P()

			// FieldOwners.<Field>Owner
			g.P("func (f *FieldOwners) ", fieldName, "Owner() (string, bool) {")
			g.P("    return f.simpleOwner(Field_", originalFieldName, ".Key())")
			g.P("}")
			g.P()

			// OwningPlugins.Clear<Field>
			g.P("func (o *OwningPlugins) Clear", fieldName, "(id, plugin string) {")
			g.P("	o.mustOwnersFor(id).Clear", fieldName, "(plugin)")
			g.P("}")
			g.P()

			// FieldOwners.Clear<Field>
			g.P("func (f *FieldOwners) Clear", fieldName, "(plugin string) {")
			g.P("	f.clearSimple(Field_", originalFieldName, ".Key(), plugin)")
			g.P("}")
			g.P()
		}
	}
}

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