// Copyright 2024 PingCAP, Inc.
//
// 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 generators

import (
	"fmt"
	"io"
	"log/slog"
	"reflect"
	"strings"

	"k8s.io/gengo/v2/generator"
	"k8s.io/gengo/v2/namer"
	"k8s.io/gengo/v2/types"
)

// overlayGenerator produces a file with autogenerated functions.
type overlayGenerator struct {
	generator.GoGenerator
	targetPackage string
	imports       namer.ImportTracker
	funcTracker   map[string]map[string]struct{}
}

func NewOverlayGenerator(outputFilename, targetPackage string) generator.Generator {
	g := &overlayGenerator{
		GoGenerator: generator.GoGenerator{
			OutputFilename: outputFilename,
		},
		targetPackage: targetPackage,
		imports:       generator.NewImportTrackerForPackage(targetPackage),
		funcTracker:   map[string]map[string]struct{}{},
	}

	g.trackFunc("overlay", types.Ref("k8s.io/apimachinery/pkg/api/resource", "Quantity"))
	g.trackFunc("overlay", types.Ref("k8s.io/apimachinery/pkg/apis/meta/v1", "ObjectMeta"))
	return g
}

// Filter returns true if this Generator cares about this type.
// This will be called for every type which made it through this Package's
// Filter method.
func (*overlayGenerator) Filter(_ *generator.Context, t *types.Type) bool {
	// We only handle exported structs.
	return t.Name.Name == "PodSpec"
}

// Namers returns a set of NameSystems which will be merged with the namers
// provided when executing this package. In case of a name collision, the
// values produced here will win.
func (g *overlayGenerator) Namers(*generator.Context) namer.NameSystems {
	return namer.NameSystems{
		// This elides package names when the name is in "this" package.
		"raw":    namer.NewRawNamer(g.targetPackage, g.imports),
		"public": namer.NewPublicNamer(0),
	}
}

// GenerateType should emit code for the specified type.  This will be called
// for every type which made it through this Generator's Filter method.
func (g *overlayGenerator) GenerateType(c *generator.Context, t *types.Type, w io.Writer) error {
	slog.Info("generating overlay", "type", t.String())

	sw := generator.NewSnippetWriter(w, c, "$", "$")

	p := NewParams(".Spec", t, nil)
	g.generateFunc(sw, p)
	return sw.Error()
}

func (g *overlayGenerator) Imports(_ *generator.Context) (imports []string) {
	importLines := []string{}
	importLines = append(importLines, g.imports.ImportLines()...)

	return importLines
}

func (g *overlayGenerator) trackFunc(ns string, t *types.Type) (hasAdded bool) {
	tracker, ok := g.funcTracker[ns]
	if !ok {
		tracker = map[string]struct{}{}
		g.funcTracker[ns] = tracker
	}

	if _, ok := tracker[t.Name.String()]; ok {
		return true
	}

	tracker[t.Name.String()] = struct{}{}

	return false
}

func (g *overlayGenerator) generateFunc(sw *generator.SnippetWriter, p Params) {
	slog.Debug("generate overlay function", "field", p.Name(), "type", p.Type())
	switch p.Type().Kind {
	case types.Struct:
		g.generateStructFunc(sw, p)
	case types.Slice, types.Array:
		g.generateSliceFunc(sw, p)
	case types.Map:
		g.generateMapFunc(sw, p)
	case types.Pointer:
		g.generateFunc(sw, p.Pointer())
	case types.Builtin:
		// builtin will be handled by inline assignment
	case types.Alias:
		g.generateAliasFunc(sw, p)
	default:
		Panic("unhandled func generation", "typeName", p.Type(), "kind", p.Type().Kind)
	}
}

func (g *overlayGenerator) generateAliasFunc(sw *generator.SnippetWriter, p Params) {
	if g.trackFunc("overlay", p.Type()) {
		return
	}
	underlying := p.Underlying()

	DoLine(sw, "", "func overlay$.|public$(dst, src *$.|raw$) {", p.Type())
	DoLine(sw, "1", "ndst := (*$.|raw$)(dst)", underlying.Type())
	DoLine(sw, "1", "nsrc := (*$.|raw$)(src)", underlying.Type())
	g.generateSafeAssignment(sw, underlying)
	DoLine(sw, "", `}`, nil)

	g.generateFunc(sw, underlying)
}

func (g *overlayGenerator) generateStructFunc(sw *generator.SnippetWriter, p Params) {
	switch p.AnnotateType() {
	case AtomicStruct:
		if g.trackFunc("overlayAtomicStruct", p.Type()) {
			return
		}
		DoLine(sw, "", "func overlayAtomic$.|public$(dst, src *$.|raw$) {", p.Type())
		DoLine(sw, "1", "*dst = *src", nil)
		DoLine(sw, "", `}`, nil)
	case GranularStruct:
		if g.trackFunc("overlay", p.Type()) {
			return
		}
		DoLine(sw, "", "func overlay$.|public$(dst, src *$.|raw$) {", p.Type())

		ms := p.Members()
		for _, m := range ms {
			g.generateAssignment(sw, m)
		}

		DoLine(sw, "", `}`, nil)
		for _, m := range ms {
			g.generateFunc(sw, m)
		}

	default:
		Panic("unhandled struct type", "type", p.Type(), "annoType", p.AnnotateType())
	}
}

func (g *overlayGenerator) generateMapFunc(sw *generator.SnippetWriter, p Params) {
	switch p.AnnotateType() {
	case AtomicMap:
		if g.trackFunc("overlayAtomicMap", p.Type()) {
			return
		}
		DoLine(sw, "", "func overlayAtomic$.|public$(dst, src *$.|raw$) {", p.Type())
		DoLine(sw, "1", "if len(*src) == 0 {", nil)
		DoLine(sw, "1", "return", nil)
		DoLine(sw, "1", "}", nil)
		DoLine(sw, "1", "*dst = *src", nil)
		DoLine(sw, "", `}`, nil)
	case GranularMap:
		if g.trackFunc("overlay", p.Type()) {
			return
		}

		elem := p.Map()

		DoLine(sw, "", "func overlay$.|public$(dst, src *$.|raw$) {", p.Type())
		DoLine(sw, "1", "for k := range *src {", nil)
		DoLine(sw, "22", "vdst, ok := (*dst)[k]", nil)
		DoLine(sw, "22", "if !ok {", nil)
		DoLine(sw, "333", "(*dst)[k] = (*src)[k]", nil)
		DoLine(sw, "22", "continue", nil)
		DoLine(sw, "22", "}", nil)
		DoLine(sw, "22", "vsrc := (*src)[k]", nil)
		g.generateAssignment(sw, elem)
		DoLine(sw, "22", "(*dst)[k] = vdst", nil)
		DoLine(sw, "1", "}", nil)
		DoLine(sw, "", "}", nil)

		g.generateFunc(sw, elem)
	default:
		Panic("unhandled map type %v for %v", p.AnnotateType(), p.Type())
	}
}

func (g *overlayGenerator) generateSliceFunc(sw *generator.SnippetWriter, p Params) {
	switch p.AnnotateType() {
	case AtomicList:
		if g.trackFunc("overlayAtomicList", p.Type()) {
			return
		}
		DoLine(sw, "", "func overlayAtomicList$.|public$(dst, src *$.|raw$) {", p.Type())
		DoLine(sw, "1", "if len(*src) == 0 {", nil)
		DoLine(sw, "1", "return", nil)
		DoLine(sw, "1", "}", nil)
		DoLine(sw, "1", "*dst = *src", nil)
		DoLine(sw, "", `}`, nil)
	case SetList:
		elem := p.Slice("di", "si")

		if g.trackFunc("overlaySetList", p.Type()) {
			return
		}
		// TODO: maybe validate elem is a scaler?
		DoLine(sw, "", "func overlaySetList$.|public$(dst, src *$.|raw$) {", p.Type())
		DoLine(sw, "1", "m := map[$.|raw$]struct{}{}", elem.Type())
		DoLine(sw, "1", "for i := range *dst {", nil)
		DoLine(sw, "22", "item := *dst[i]", nil)
		DoLine(sw, "22", "m[item] = struct{}{}", nil)
		DoLine(sw, "1", "}", nil)
		DoLine(sw, "1", "for i := range *src {", nil)
		DoLine(sw, "22", "item := *src[i]", nil)
		DoLine(sw, "22", "if _, ok := m[item]; !ok {", nil)
		DoLine(sw, "333", "*dst = append(*dst, item)", nil)
		DoLine(sw, "22", "}", nil)
		DoLine(sw, "1", "}", nil)
		DoLine(sw, "", `}`, nil)

		g.generateFunc(sw, elem)

	case MapList:
		if p.Type().Elem.Kind != types.Struct {
			Panic("elem of map list is not struct", "type", p.Type())
		}

		elem := p.Slice("di", "si")

		fields, err := getFieldsByKeys(elem.Type(), p.Keys()...)
		if err != nil {
			Panic("get field names of map list failed", "err", err, "type", p.Type())
		}
		if g.trackFunc("overlayMapList", p.Type()) {
			return
		}
		joinFunc := types.Ref("strings", "Join")
		DoLine(sw, "", "func overlayMapList$.|public$(dst, src *$.|raw$) {", p.Type())
		DoLine(sw, "1", "m := map[string]int{}", nil)
		DoLine(sw, "1", "for i := range *dst {", nil)
		DoLine(sw, "22", "item := (*dst)[i]", nil)
		DoLine(sw, "22", "keys := []string{}", nil)
		for _, field := range fields {
			DoLine(sw, "22", "keys = append(keys, $.Type|public$ToString(item.$.Name$))", &field)
		}
		DoLine(sw, "22", `key := $.|raw$(keys, ",")`, joinFunc)
		DoLine(sw, "22", "m[key] = i", nil)
		DoLine(sw, "1", "}", nil)
		DoLine(sw, "1", "for si := range *src{", nil)
		DoLine(sw, "22", "item := (*src)[si]", nil)
		DoLine(sw, "22", "keys := []string{}", nil)
		for _, field := range fields {
			DoLine(sw, "22", "keys = append(keys, $.Type|public$ToString(item.$.Name$))", &field)
		}
		DoLine(sw, "22", `key := $.|raw$(keys, ",")`, joinFunc)
		DoLine(sw, "22", "di, ok := m[key]", nil)
		DoLine(sw, "22", "if !ok {", nil)
		DoLine(sw, "333", "*dst = append(*dst, item)", nil)
		DoLine(sw, "22", "continue", nil)
		DoLine(sw, "22", "}", nil)
		g.generateAssignment(sw, elem)
		DoLine(sw, "1", "}", nil)
		DoLine(sw, "", `}`, nil)

		for _, field := range fields {
			g.generateToStringFunc(sw, field.Type)
		}

		g.generateFunc(sw, elem)
	}
}

func (g *overlayGenerator) generateToStringFunc(sw *generator.SnippetWriter, t *types.Type) {
	if g.trackFunc("toString", t) {
		return
	}
	DoLine(sw, "", "func $.|public$ToString(val $.|raw$) string {", t)
	switch t.Kind {
	case types.Builtin:
		switch t {
		case types.Int, types.Int64, types.Int32, types.Int16:
			DoLine(sw, "1", `return $.|raw$(int64(val), 10)`, types.Ref("strconv", "FormatInt"))
		case types.Uint, types.Uint64, types.Uint32, types.Uint16, types.Byte:
			DoLine(sw, "1", `return $.|raw$(uint64(val), 10)`, types.Ref("strconv", "FormatUint"))
		case types.Float, types.Float32, types.Float64:
			DoLine(sw, "1", `return $.|raw$(float64(val), 'E', -1, 64)`, types.Ref("strconv", "FormatFloat"))
		case types.String:
			DoLine(sw, "1", `return val`, nil)
		case types.Bool:
			DoLine(sw, "1", `return $.|raw$(val)`, types.Ref("strconv", "FormatBool"))
		default:
			Panic("unhandled builtin to string func", "name", t.Name)
		}
	case types.Alias:
		DoLine(sw, "1", "return $.|public$ToString($.|raw$(val))", t.Underlying)
	default:
		Panic("unhandled to string func", "name", t.Name, "kind", t.Kind)
	}
	DoLine(sw, "", "}", t)
}

type StructField struct {
	Name string
	Type *types.Type
}

func getFieldsByKeys(t *types.Type, keys ...string) ([]StructField, error) {
	fields := []StructField{}
	for _, key := range keys {
		for _, m := range t.Members {
			if m.Embedded {
				field, err := getFieldsByKeys(m.Type, key)
				if err != nil {
					continue
				}
				fields = append(fields, field...)
				break
			}
			tag := reflect.StructTag(m.Tags)
			jsonTag := tag.Get("json")
			name := strings.Split(jsonTag, ",")
			if name[0] == key {
				fields = append(fields, StructField{
					Name: m.Name,
					Type: m.Type,
				})
			}
		}
	}

	if len(fields) != len(keys) {
		return nil, fmt.Errorf("cannot get field names by keys %v, only %v are found", keys, fields)
	}

	return fields, nil
}

func (g *overlayGenerator) generateAssignment(sw *generator.SnippetWriter, p Params) {
	args := generator.Args{
		"type": p.Type(),
		"dst":  p.Dst(),
		"src":  p.Src(),
	}

	p.CheckNil()

	needNilCheck := g.needNilCheck(p.Type())
	if needNilCheck {
		DoLine(sw, "", `if $.dst$ != nil && $.src$ != nil {`, args)
	}

	g.generateSafeAssignment(sw, p)

	if needNilCheck {
		DoLine(sw, "", "} else if $.dst$ == nil {", args)
		DoLine(sw, "", "$.dst$ = $.src$", args)
		DoLine(sw, "", "}", nil)
	}
}

func (g *overlayGenerator) generateSafeAssignment(sw *generator.SnippetWriter, p Params) {
	switch p.Type().Kind {
	case types.Struct:
		g.generateStructAssignment(sw, p)
	case types.Alias:
		g.generateAliasAssignment(sw, p)
	case types.Slice, types.Array:
		g.generateSliceAssignment(sw, p)
	case types.Map:
		g.generateMapAssignment(sw, p)
	case types.Builtin:
		g.generateBuiltinAssignment(sw, p)
	case types.Pointer:
		g.generatePointerAssignment(sw, p)
	default:
		Panic("unhandled assignment generation", "dst", p.Dst(), "src", p.Src(), "typeName", p.Type())
	}
}

func (g *overlayGenerator) needNilCheck(t *types.Type) bool {
	switch t.Kind {
	case types.Alias:
		return g.needNilCheck(t.Underlying)
	case types.Slice, types.Array, types.Map, types.Pointer:
		return true
	}

	return false
}

func (g *overlayGenerator) generatePointerAssignment(sw *generator.SnippetWriter, p Params) {
	to := p.Pointer()
	g.generateSafeAssignment(sw, to)
}

func (*overlayGenerator) generateSliceAssignment(sw *generator.SnippetWriter, p Params) {
	args := generator.Args{
		"type": p.Type(),
		"dst":  p.Dst(),
		"src":  p.Src(),
	}
	switch p.AnnotateType() {
	case AtomicList:
		DoLine(sw, "1", `overlayAtomicList$.type|public$($.dst$, $.src$)`, args)
	case SetList:
		DoLine(sw, "1", `overlaySetList$.type|public$($.dst$, $.src$)`, args)
	case MapList:
		DoLine(sw, "1", `overlayMapList$.type|public$($.dst$, $.src$)`, args)
	default:
		Panic("unhandled list type", "type", p.Type(), "annoType", p.AnnotateType())
	}
}

func (*overlayGenerator) generateMapAssignment(sw *generator.SnippetWriter, p Params) {
	args := generator.Args{
		"type": p.Type(),
		"dst":  p.Dst(),
		"src":  p.Src(),
	}
	switch p.AnnotateType() {
	case AtomicMap:
		DoLine(sw, "1", `overlayAtomic$.type|public$($.dst$, $.src$)`, args)
	case GranularMap:
		DoLine(sw, "1", `overlay$.type|public$($.dst$, $.src$)`, args)
	default:
		Panic("unhandled map type", "type", p.Type(), "annoType", p.AnnotateType())
	}
}

func (*overlayGenerator) generateStructAssignment(sw *generator.SnippetWriter, p Params) {
	args := generator.Args{
		"type": p.Type(),
		"dst":  p.Dst(),
		"src":  p.Src(),
	}
	switch p.AnnotateType() {
	case AtomicStruct:
		DoLine(sw, "1", `overlayAtomic$.type|public$($.dst$, $.src$)`, args)
	case GranularStruct:
		DoLine(sw, "1", `overlay$.type|public$($.dst$, $.src$)`, args)
	default:
		Panic("unhandled struct type", "type", p.Type(), "annoType", p.AnnotateType())
	}
}

func (*overlayGenerator) generateAliasAssignment(sw *generator.SnippetWriter, p Params) {
	args := generator.Args{
		"type": p.Type(),
		"dst":  p.Dst(),
		"src":  p.Src(),
	}
	DoLine(sw, "1", `overlay$.type|public$($.dst$, $.src$)`, args)
}

func (*overlayGenerator) generateBuiltinAssignment(sw *generator.SnippetWriter, p Params) {
	switch p.Type() {
	case types.Int, types.Int64, types.Int32, types.Int16, types.Uint, types.Uint64,
		types.Uint32, types.Uint16, types.Byte, types.Float, types.Float32, types.Float64:
		DoLine(sw, "1", `if $.$ != 0 {`, p.Src())
	case types.String:
		DoLine(sw, "1", `if $.$ != "" {`, p.Src())
	case types.Bool:
		DoLine(sw, "1", `if $.$ {`, p.Src())
	default:
		Panic("unhandled builtin assignment", "dst", p.Dst(), "src", p.Src(), "typeName", p.Type())
	}
	DoLine(sw, "22", `$.dst$ = $.src$`, generator.Args{
		"dst": p.Dst(),
		"src": p.Src(),
	})
	DoLine(sw, "1", `}`, nil)
}