/*
 * Copyright 2025 CloudWeGo 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 prutalgen

import (
	"go/format"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
)

type MarshalerType int

const (
	MarshalerKitexProtobuf = 1 << iota
)

type GoCodeGen struct {
	Format bool

	// Generates GetXXX funcs
	Getter bool

	Marshaler MarshalerType
}

func NewGoCodeGen() *GoCodeGen {
	return &GoCodeGen{
		Format: true,
	}
}

var shellsafeRE = regexp.MustCompile(`^[a-zA-Z0-9\-_.=/]+$`)

func argsQuote(args []string) string {
	args = append([]string(nil), args...)
	for i, a := range args {
		if !shellsafeRE.MatchString(a) {
			args[i] = "'" + strings.ReplaceAll(a, `'`, `\'`) + "'"
		}
	}
	return strings.Join(args, " ")
}

func (g *GoCodeGen) SourcePath(p *Proto, gt GenPathType, out string, suffix string) string {
	fn := ""
	switch gt {
	case GenByImport:
		// same name as proto file
		fn = strings.TrimSuffix(filepath.Base(p.ProtoFile), ".proto") + suffix
		// file path uses import path
		fn = filepath.Join(out, fullFilename(p.GoImport, fn))

	case GenBySourceRelative:
		// same path as proto file
		fn = strings.TrimSuffix(p.ProtoFile, ".proto") + suffix

	default:
		panic("unknown GenPathType: " + gt)
	}
	return fn
}

func (g *GoCodeGen) Gen(p *Proto, gt GenPathType, out string) error {
	fn := g.SourcePath(p, gt, out, ".pb.go")
	p.Infof("generating %s", refPath(fn))

	header := "// Code generated by prutalgen. DO NOT EDIT.\n" +
		"// " + argsQuote(os.Args)
	w := NewCodeWriter(header, p.GoPackage)
	g.ProtoGen(p, w)
	b := w.Bytes()
	err := os.MkdirAll(filepath.Dir(fn), 0755)
	if err != nil {
		return err
	}
	err = os.WriteFile(fn, b, 0644)
	if err != nil {
		return err
	}
	if g.Format {
		newb, err := format.Source(b)
		if err != nil {
			p.Warnf("code format err: %s.  Possibly a BUG and please report to author.", err)
		} else {
			err := os.WriteFile(fn, newb, 0644)
			if err != nil {
				return err
			}
		}
	}
	p.Infof("generated %s", refPath(fn))
	return nil
}

func (g *GoCodeGen) ProtoGen(p *Proto, w *CodeWriter) {
	for _, e := range p.Enums {
		w.F("")
		g.EnumGen(e, w)
	}

	// always generate enums before messages
	// same order as protobuf
	for _, m := range p.Messages {
		for _, e := range m.Enums {
			w.F("")
			g.EnumGen(e, w)
		}
	}

	for _, m := range p.Messages {
		w.F("")
		g.MessageGen(m, w)
	}
}

func (g *GoCodeGen) EnumGen(e *Enum, w *CodeWriter) {
	if e.HeadComment != "" {
		w.F("%s", e.HeadComment)
	}
	w.F("type %s int32"+e.InlineComment, e.GoName)
	w.F("")
	if len(e.Fields) > 0 {
		w.F("const (")
		for i, f := range e.Fields {
			if f.HeadComment != "" {
				if i != 0 {
					w.F("") // empty line between last field and the comment of new field
				}
				w.F("%s", f.HeadComment)
			}
			w.F("%s %s = %d %s", f.GoName, e.GoName, f.Value, f.InlineComment)
		}
		w.F(")")
	}

	if e.genMapping() {
		w.F("\n// Enum value maps for %s.", e.GoName)
		w.F("var %s_name = map[int32]string {", e.GoName)
		for _, f := range e.Fields {
			w.F("%d: %q,", f.Value, f.Name)
		}
		w.F("}")
		w.F("")
		w.F("var %s_value = map[string]int32 {", e.GoName)
		for _, f := range e.Fields {
			w.F("%q: %d,", f.Name, f.Value)
		}
		w.F("}")

		// String func
		w.UsePkg("strconv", "")
		w.F("func (x %s) String() string {", e.GoName)
		w.F("s, ok := %s_name[int32(x)]", e.GoName)
		w.F("if ok { return s }")
		w.F("return strconv.Itoa(int(x))")
		w.F("}")
	}
}

func (g *GoCodeGen) MessageGen(m *Message, w *CodeWriter) {
	if m.HeadComment != "" {
		w.F("%s", m.HeadComment)
	}
	generatedOneOfField := map[*Oneof]bool{}

	w.F("type %s struct {", m.GoName)
	for i, f := range m.Fields {
		if f.Oneof != nil {
			if generatedOneOfField[f.Oneof] {
				continue
			}
			// The Oneof fields will be generated in m.Oneofs loop
			// Only need placeholder here for Oneof fields
			o := f.Oneof
			w.F("// Types that are assignable to %s:", o.FieldName())
			w.F("//")
			for _, f := range o.Fields {
				w.F("//\t*%s_%s", m.GoName, f.GoName)
			}
			// field and protobuf_oneof tag
			w.F("%s %s `protobuf_oneof:\"%s\"`", o.FieldName(), o.FieldType(), o.Name)
			generatedOneOfField[f.Oneof] = true
			continue
		}
		if i != 0 && f.HeadComment != "" {
			w.F("") // empty line between last field and the comment of new field
		}
		g.FieldGen(f, w)
	}
	if m.genUnknownFields() {
		w.F("unknownFields []byte `json:\"-\"`")
	}
	w.F("} %s", m.InlineComment)

	// func Reset
	w.F("\nfunc (x *%s) Reset() { *x = %s{} }", m.GoName, m.GoName)

	if g.Marshaler&MarshalerKitexProtobuf != 0 {
		w.UsePkg("github.com/cloudwego/prutal", "")
		w.F("\nfunc (x *%s) Marshal(in []byte) ([]byte, error) { return prutal.MarshalAppend(in, x) }", m.GoName)
		w.F("\nfunc (x *%s) Unmarshal(in []byte) error { return prutal.Unmarshal(in, x) }", m.GoName)
	}

	for k := range generatedOneOfField {
		delete(generatedOneOfField, k)
	}

	if !g.Getter {
		goto SkipGetter
	}
	// GetXXX
	for _, f := range m.Fields {
		w.F("")
		if f.Oneof != nil {
			o := f.Oneof
			if !generatedOneOfField[f.Oneof] {
				w.F("func (x *%s) Get%s() %s {", m.GoName, o.FieldName(), o.FieldType())
				w.F("if x != nil { return x.%s }", o.FieldName())
				w.F("return nil")
				w.F("}")
			}
			generatedOneOfField[f.Oneof] = true

			w.F("func (x *%s) Get%s() %s {", m.GoName, f.GoName, f.GoTypeName())
			w.F("if p, ok := x.Get%s().(*%s); ok {", o.FieldName(), f.OneofStructName())
			w.F("return p.%s", f.GoName)
			w.F("}")
			w.F("return %s", f.GoZero())
			w.F("}")
			continue
		}

		typename := f.GoTypeName()
		if !f.IsMessage() && typename[0] == '*' { // optional field with pointer
			w.F("func (x *%s) Get%s() %s {", m.GoName, f.GoName, typename[1:])
			w.F("if x != nil && x.%s != nil { return *x.%s }", f.GoName, f.GoName)
		} else {
			w.F("func (x *%s) Get%s() %s {", m.GoName, f.GoName, typename)
			w.F("if x != nil { return x.%s }", f.GoName)
		}
		w.F("return %s", f.GoZero())
		w.F("}")
	}
SkipGetter:

	if len(m.Oneofs) > 0 {
		w.F("// XXX_OneofWrappers is for the internal use of the prutal package.")
		w.F("func (*%s) XXX_OneofWrappers() []interface{} {", m.GoName)
		w.F("return []interface{}{")
		for _, o := range m.Oneofs {
			for _, f := range o.Fields {
				w.F("(*%s)(nil),", f.OneofStructName())
			}
		}
		w.F("}}")
	}

	for _, o := range m.Oneofs {
		w.F("")
		g.OneofGen(o, w)
	}

	// always generate embedded messasges after main message
	// same order as protobuf
	for _, x := range m.Messages {
		w.F("")
		g.MessageGen(x, w)
	}
}

func (g *GoCodeGen) FieldStructTag(f *Field) []byte {
	b := make([]byte, 0, 100)

	// protobuf
	b = append(b, `protobuf:"`...)

	// wiretype
	if f.IsMap() {
		b = append(b, "bytes"...)
	} else {
		b = append(b, f.Type.EncodingType()...)
	}

	// field number
	b = append(b, ',')
	b = strconv.AppendInt(b, int64(f.FieldNumber), 10)

	// field label
	b = append(b, ',')
	if f.Repeated || f.IsMap() {
		b = append(b, "rep"...)
		if f.IsPackedEncoding() {
			b = append(b, ",packed"...)
		}
	} else if f.Required {
		b = append(b, "req"...)
	} else {
		b = append(b, "opt"...)
	}

	// field name
	b = append(b, ',')
	b = append(b, "name="...)
	b = append(b, f.Name...)
	b = append(b, '"') // end of protobuf tag

	// json
	b = append(b, ` json:"`...)
	b = append(b, f.Name...)
	b = append(b, ",omitempty"...)
	b = append(b, '"')

	if f.IsMap() {
		// protobuf_key
		b = append(b, ` protobuf_key:"`...)
		b = append(b, f.Key.EncodingType()...)
		b = append(b, ",1,opt,name=key"...)
		b = append(b, '"')

		// protobuf_val
		b = append(b, ` protobuf_val:"`...)
		b = append(b, f.Type.EncodingType()...)
		b = append(b, ",2,opt,name=value"...)
		b = append(b, '"')
	}
	return b
}

func (g *GoCodeGen) FieldGen(f *Field, w *CodeWriter) {
	if f.HeadComment != "" {
		w.F("%s", f.HeadComment)
	}
	w.F("%s %s `%s` %s", f.GoName, f.GoTypeName(), g.FieldStructTag(f), f.InlineComment)
	if f.Type.IsExternalType() {
		w.UsePkg(f.Type.GoImport(), "")
	}
	if f.Key != nil && f.Key.IsExternalType() {
		w.UsePkg(f.Key.GoImport(), "")
	}
}

func (g *GoCodeGen) OneofGen(o *Oneof, w *CodeWriter) {
	// all oneof fields share the same private interface
	typename := o.FieldType()
	w.F("type %s interface {", typename)
	w.F("%s()", typename)
	w.F("}")

	for _, f := range o.Fields {
		structname := f.OneofStructName()
		w.F("")
		w.F("type %s struct {", structname)
		g.FieldGen(f, w)
		w.F("}")
		w.F("")

		// implements the oneof private interface
		w.F("func (*%s) %s() {}", structname, typename)
	}
}
