// Copyright 2021 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 thriftgo

import (
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"runtime"
	"strconv"
	"strings"
	"text/template"

	"github.com/cloudwego/thriftgo/generator/golang"
	"github.com/cloudwego/thriftgo/generator/golang/templates"
	"github.com/cloudwego/thriftgo/generator/golang/templates/slim"
	"github.com/cloudwego/thriftgo/parser"
	"github.com/cloudwego/thriftgo/plugin"

	"github.com/cloudwego/kitex/tool/internal_pkg/generator"
	"github.com/cloudwego/kitex/tool/internal_pkg/util"
)

var extraTemplates []string

type PatchFunc func(req *plugin.Request, p *Patcher) ([]*plugin.Generated, error)

var extraPatchFunc PatchFunc

// AppendToTemplate string
func AppendToTemplate(text string) {
	extraTemplates = append(extraTemplates, text)
}

func AppendExtensionPatches(f PatchFunc) {
	extraPatchFunc = f
}

const kitexUnusedProtection = `
// KitexUnusedProtection is used to prevent 'imported and not used' error.
var KitexUnusedProtection = struct{}{}
`

//lint:ignore U1000 until protectionInsertionPoint is used
var protectionInsertionPoint = "KitexUnusedProtection"

type Patcher struct {
	noFastAPI             bool
	utils                 *golang.CodeUtils
	module                string
	copyIDL               bool
	version               string
	record                bool
	recordCmd             []string
	deepCopyAPI           bool
	protocol              string
	handlerReturnKeepResp bool

	frugalStruct []string

	fileTpl *template.Template
	libs    map[string]string
}

func (p *Patcher) GetModule() string {
	return p.module
}

func (p *Patcher) GetUtils() *golang.CodeUtils {
	return p.utils
}

// UseFrugalForStruct judge if we need to replace fastCodec to frugal implementation. It's related to the argument '-frugal-struct'.
func (p *Patcher) UseFrugalForStruct(st *golang.StructLike) bool {
	// '@all' matches all structs
	// '@auto' matches those with annotation (go.codec="frugal")
	// otherwise, check if the given name matches the struct's name
	for _, structName := range p.frugalStruct {
		if structName == "@all" || st.GetName() == structName {
			return true
		}
		if structName == "@auto" {
			// find annotation go.codec="frugal"
			for _, anno := range st.Annotations {
				if anno.GetKey() == "go.codec" && len(anno.GetValues()) > 0 && anno.GetValues()[0] == "frugal" {
					return true
				}
			}
		}
	}
	return false
}

func (p *Patcher) UseLib(path, alias string) string {
	if p.libs == nil {
		p.libs = make(map[string]string)
	}
	p.libs[path] = alias
	return ""
}

func (p *Patcher) buildTemplates() (err error) {
	m := p.utils.BuildFuncMap()
	m["PrintImports"] = util.PrintlImports
	m["UseLib"] = p.UseLib
	m["UseFrugalForStruct"] = p.UseFrugalForStruct
	m["ZeroWriter"] = ZeroWriter
	m["ZeroBLength"] = ZeroBLength
	m["ReorderStructFields"] = p.reorderStructFields
	m["TypeIDToGoType"] = func(t string) string { return typeIDToGoType[t] }
	m["IsBinaryOrStringType"] = p.isBinaryOrStringType
	m["Version"] = func() string { return p.version }
	m["GenerateFastAPIs"] = func() bool { return !p.noFastAPI && p.utils.Template() != "slim" }
	m["GenerateDeepCopyAPIs"] = func() bool { return p.deepCopyAPI }
	m["GenerateArgsResultTypes"] = func() bool { return p.utils.Template() == "slim" }
	m["ImportPathTo"] = generator.ImportPathTo
	m["Str"] = func(id int32) string {
		if id < 0 {
			return "_" + strconv.Itoa(-int(id))
		}
		return strconv.Itoa(int(id))
	}
	m["IsNil"] = func(i interface{}) bool {
		return i == nil || reflect.ValueOf(i).IsNil()
	}
	m["SourceTarget"] = func(s string) string {
		// p.XXX
		if strings.HasPrefix(s, "p.") {
			return "src." + s[2:]
		}
		// _key, _val
		return s[1:]
	}
	m["FieldName"] = func(s string) string {
		// p.XXX
		return strings.ToLower(s[2:3]) + s[3:]
	}
	m["IsHessian"] = func() bool {
		return p.IsHessian2()
	}
	m["IsGoStringType"] = func(typeName golang.TypeName) bool {
		return typeName == "string" || typeName == "*string"
	}

	tpl := template.New("kitex").Funcs(m)
	allTemplates := basicTemplates
	if p.utils.Template() == "slim" {
		allTemplates = append(allTemplates, slim.StructLike,
			templates.StructLikeDefault,
			templates.FieldGetOrSet,
			templates.FieldIsSet,
			StructLikeDeepEqualEmpty,
			StructLikeDeepCopy,
			FieldDeepCopy,
			FieldDeepCopyStructLike,
			FieldDeepCopyContainer,
			FieldDeepCopyMap,
			FieldDeepCopyList,
			FieldDeepCopySet,
			FieldDeepCopyBaseType,
			StructLikeCodec,
			StructLikeProtocol,
			JavaClassName,
			Processor,
		)
	} else {
		allTemplates = append(allTemplates, StructLikeCodec,
			StructLikeFastReadField,
			StructLikeDeepCopy,
			StructLikeFastWrite,
			StructLikeFastRead,
			StructLikeFastWriteNocopy,
			StructLikeLength,
			StructLikeFastWriteField,
			StructLikeFieldLength,
			StructLikeProtocol,
			JavaClassName,
			FieldFastRead,
			FieldFastReadStructLike,
			FieldFastReadBaseType,
			FieldFastReadContainer,
			FieldFastReadMap,
			FieldFastReadSet,
			FieldFastReadList,
			FieldDeepCopy,
			FieldDeepCopyStructLike,
			FieldDeepCopyContainer,
			FieldDeepCopyMap,
			FieldDeepCopyList,
			FieldDeepCopySet,
			FieldDeepCopyBaseType,
			FieldFastWrite,
			FieldLength,
			FieldFastWriteStructLike,
			FieldStructLikeLength,
			FieldFastWriteBaseType,
			FieldBaseTypeLength,
			FieldFixedLengthTypeLength,
			FieldFastWriteContainer,
			FieldContainerLength,
			FieldFastWriteMap,
			FieldMapLength,
			FieldFastWriteSet,
			FieldSetLength,
			FieldFastWriteList,
			FieldListLength,
			templates.FieldDeepEqual,
			templates.FieldDeepEqualBase,
			templates.FieldDeepEqualStructLike,
			templates.FieldDeepEqualContainer,
			ValidateSet,
			Processor,
		)
	}
	for i, txt := range allTemplates {
		tpl, err = tpl.Parse(txt)
		if err != nil {
			name := strconv.Itoa(i)
			if ix := strings.Index(txt, "{{define "); ix >= 0 {
				ex := strings.Index(txt, "}}")
				if ex >= 0 {
					name = txt[ix+len("{{define ") : ex]
				}
			}
			return fmt.Errorf("parse template %s failed: %v", name, err)
		}
	}

	ext := `{{define "ExtraTemplates"}}{{end}}`
	if len(extraTemplates) > 0 {
		ext = fmt.Sprintf("{{define \"ExtraTemplates\"}}\n%s\n{{end}}",
			strings.Join(extraTemplates, "\n"))
	}
	tpl, err = tpl.Parse(ext)
	if err != nil {
		return fmt.Errorf("failed to parse extra templates: %w: %q", err, ext)
	}

	if p.IsHessian2() {
		tpl, err = tpl.Parse(RegisterHessian)
		if err != nil {
			return fmt.Errorf("failed to parse hessian2 templates: %w: %q", err, RegisterHessian)
		}
	}

	p.fileTpl = tpl
	return nil
}

const ImportInsertPoint = "// imports insert-point"

func (p *Patcher) patch(req *plugin.Request) (patches []*plugin.Generated, err error) {
	if err := p.buildTemplates(); err != nil {
		return nil, err
	}

	var buf strings.Builder

	// fd, _ := os.OpenFile("dump.txt", os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
	// defer fd.Close()

	protection := make(map[string]*plugin.Generated)

	var trees chan *parser.Thrift
	if req.Recursive {
		trees = req.AST.DepthFirstSearch()
	} else {
		trees = make(chan *parser.Thrift, 1)
		trees <- req.AST
		close(trees)
	}

	for ast := range trees {

		// Reset libs to empty. When executing next AST, the dependencies left by previous AST should not be retained.
		p.libs = nil

		// fd.WriteString(p.utils.GetFilename(ast) + "\n")
		// scope, err := golang.BuildScope(p.utils, ast)
		scope, _, err := golang.BuildRefScope(p.utils, ast)
		if err != nil {
			return nil, fmt.Errorf("build scope for ast %q: %w", ast.Filename, err)
		}
		p.utils.SetRootScope(scope)
		pkgName := golang.GetImportPackage(golang.GetImportPath(p.utils, ast))

		path := p.utils.CombineOutputPath(req.OutputPath, ast)
		base := p.utils.GetFilename(ast)
		target := util.JoinPath(path, "k-"+base)

		// Define KitexUnusedProtection in k-consts.go .
		// Add k-consts.go before target to force the k-consts.go generated by consts.thrift to be renamed.
		consts := util.JoinPath(path, "k-consts.go")
		if protection[consts] == nil {
			patch := &plugin.Generated{
				Content: "package " + pkgName + "\n" + kitexUnusedProtection,
				Name:    &consts,
			}
			patches = append(patches, patch)
			protection[consts] = patch
		}

		buf.Reset()

		// if all scopes are ref, don't generate k-xxx
		if scope == nil {
			continue
		}

		if p.IsHessian2() {
			register := util.JoinPath(path, fmt.Sprintf("hessian2-register-%s", base))
			patch, err := p.patchHessian(path, scope, pkgName, base)
			if err != nil {
				return nil, fmt.Errorf("patch hessian fail for %q: %w", ast.Filename, err)
			}

			patches = append(patches, patch)
			protection[register] = patch
		}

		data := &struct {
			Scope    *golang.Scope
			PkgName  string
			Imports  []util.Import
			Includes []util.Import
		}{Scope: scope, PkgName: pkgName}

		if err = p.fileTpl.ExecuteTemplate(&buf, "file", data); err != nil {
			return nil, fmt.Errorf("%q: %w", ast.Filename, err)
		}
		content := buf.String()

		imps, err := scope.ResolveImports()
		if err != nil {
			return nil, fmt.Errorf("resolve imports failed for %q: %w", ast.Filename, err)
		}
		for path, alias := range p.libs {
			imps[path] = alias
		}
		data.Imports = util.SortImports(imps, p.module)
		data.Includes = p.extractLocalLibs(data.Imports)

		// replace imports insert-pointer with newly rendered output
		buf.Reset()
		if err = p.fileTpl.ExecuteTemplate(&buf, "imports", data); err != nil {
			return nil, fmt.Errorf("%q: %w", ast.Filename, err)
		}
		imports := buf.String()
		if i := strings.Index(content, ImportInsertPoint); i >= 0 {
			content = strings.Replace(content, ImportInsertPoint, imports, 1)
		} else {
			return nil, fmt.Errorf("replace imports failed")
		}

		patches = append(patches, &plugin.Generated{
			Content: content,
			Name:    &target,
		})
		// fd.WriteString("patched: " + target + "\n")
		// fd.WriteString("content: " + content + "\nend\n")

		if p.copyIDL {
			content, err := os.ReadFile(ast.Filename)
			if err != nil {
				return nil, fmt.Errorf("read %q: %w", ast.Filename, err)
			}
			path := util.JoinPath(path, filepath.Base(ast.Filename))
			patches = append(patches, &plugin.Generated{
				Content: string(content),
				Name:    &path,
			})
		}

		if p.record {
			content := doRecord(p.recordCmd)
			bashPath := util.JoinPath(getBashPath())
			patches = append(patches, &plugin.Generated{
				Content: content,
				Name:    &bashPath,
			})
		}

	}

	if extraPatchFunc != nil {
		extPatches, err := extraPatchFunc(req, p)
		if err != nil {
			return nil, err
		}
		patches = append(patches, extPatches...)
	}

	return
}

func (p *Patcher) patchHessian(path string, scope *golang.Scope, pkgName, base string) (patch *plugin.Generated, err error) {
	buf := strings.Builder{}
	resigterIDLName := fmt.Sprintf("hessian2-register-%s", base)
	register := util.JoinPath(path, resigterIDLName)
	data := &struct {
		Scope   *golang.Scope
		PkgName string
		Imports map[string]string
		GoName  string
		IDLName string
	}{Scope: scope, PkgName: pkgName, IDLName: util.UpperFirst(strings.Replace(base, ".go", "", -1))}
	data.Imports, err = scope.ResolveImports()
	if err != nil {
		return nil, err
	}

	if err = p.fileTpl.ExecuteTemplate(&buf, "register", data); err != nil {
		return nil, err
	}
	patch = &plugin.Generated{
		Content: buf.String(),
		Name:    &register,
	}
	return patch, nil
}

func getBashPath() string {
	if runtime.GOOS == "windows" {
		return "kitex-all.bat"
	}
	return "kitex-all.sh"
}

func (p *Patcher) extractLocalLibs(imports []util.Import) []util.Import {
	ret := make([]util.Import, 0)
	prefix := p.module + "/"
	kitexPkgPath := generator.ImportPathTo("pkg")
	// remove std libs and thrift to prevent duplicate import.
	for _, v := range imports {
		if strings.HasPrefix(v.Path, kitexPkgPath) {
			// fix bad the case like: `undefined: bthrift.KitexUnusedProtection`
			// when we generate code in kitex repo.
			// we may never ref to other generate code in kitex repo, if do fix me.
			continue
		}
		// local packages
		if strings.HasPrefix(v.Path, prefix) {
			ret = append(ret, v)
			continue
		}
	}
	return ret
}

// DoRecord records current cmd into kitex-all.sh
func doRecord(recordCmd []string) string {
	bytes, err := os.ReadFile(getBashPath())
	content := string(bytes)
	if err != nil {
		content = "#! /usr/bin/env bash\n"
	}
	var input, currentIdl string
	for _, s := range recordCmd {
		if s != "-record" {
			input += s + " "
		}
		if strings.HasSuffix(s, ".thrift") || strings.HasSuffix(s, ".proto") {
			currentIdl = s
		}
	}
	if input != "" && currentIdl != "" {
		find := false
		lines := strings.Split(content, "\n")
		for i, line := range lines {
			if strings.Contains(input, "-service") && strings.Contains(line, "-service") {
				lines[i] = input
				find = true
				break
			}
			if strings.Contains(line, currentIdl) && !strings.Contains(line, "-service") {
				lines[i] = input
				find = true
				break
			}
		}
		if !find {
			content += "\n" + input
		} else {
			content = strings.Join(lines, "\n")
		}
	}
	return content
}

func (p *Patcher) reorderStructFields(fields []*golang.Field) ([]*golang.Field, error) {
	fixedLengthFields := make(map[*golang.Field]bool, len(fields))
	for _, field := range fields {
		fixedLengthFields[field] = golang.IsFixedLengthType(field.Type)
	}

	sortedFields := make([]*golang.Field, 0, len(fields))
	for _, v := range fields {
		if fixedLengthFields[v] {
			sortedFields = append(sortedFields, v)
		}
	}
	for _, v := range fields {
		if !fixedLengthFields[v] {
			sortedFields = append(sortedFields, v)
		}
	}

	return sortedFields, nil
}

func (p *Patcher) isBinaryOrStringType(t *parser.Type) bool {
	return t.Category.IsBinary() || t.Category.IsString()
}

func (p *Patcher) IsHessian2() bool {
	return strings.EqualFold(p.protocol, "hessian2")
}

var typeIDToGoType = map[string]string{
	"Bool":   "bool",
	"Byte":   "int8",
	"I16":    "int16",
	"I32":    "int32",
	"I64":    "int64",
	"Double": "float64",
	"String": "string",
	"Binary": "[]byte",
}
