// 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 generator

import (
	"bytes"
	"fmt"
	"strings"
	"text/template"

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

// File .
type File struct {
	Name    string
	Content string
}

// PackageInfo contains information to generate a package for a service.
type PackageInfo struct {
	Namespace    string            // a dot-separated string for generating service package under kitex_gen
	Dependencies map[string]string // package name => import path, used for searching imports
	*ServiceInfo                   // the target service
	Services     []*ServiceInfo    // all services defined in a IDL for multiple services scenario

	// the following fields will be filled and used by the generator
	Codec            string
	NoFastAPI        bool
	Version          string
	RealServiceName  string
	Imports          map[string]map[string]bool // import path => alias
	ExternalKitexGen string
	Features         []feature
	FrugalPretouch   bool
	Module           string
	Protocol         transport.Protocol
	IDLName          string
	ServerPkg        string
	StreamX          bool
}

// AddImport .
func (p *PackageInfo) AddImport(pkg, path string) {
	if p.Imports == nil {
		p.Imports = make(map[string]map[string]bool)
	}
	if pkg != "" {
		path = p.toExternalGenPath(path)
		if path == pkg {
			p.Imports[path] = nil
		} else {
			if p.Imports[path] == nil {
				p.Imports[path] = make(map[string]bool)
			}
			p.Imports[path][pkg] = true
		}
	}
}

// AddImports .
func (p *PackageInfo) AddImports(pkgs ...string) {
	for _, pkg := range pkgs {
		if path, ok := p.Dependencies[pkg]; ok {
			p.AddImport(pkg, path)
		} else {
			p.AddImport(pkg, pkg)
		}
	}
}

// UpdateImportPath changed the mapping between alias -> import path
// For instance:
//
//	Original import: alias "original_path"
//	Invocation:      UpdateImport("alias", "new_path")
//	New import:      alias "new_path"
//
// if pkg == newPath, then alias would be removed in import sentence:
//
//	Original import: context "path/to/custom/context"
//	Invocation:      UpdateImport("context", "context")
//	New import:      context
func (p *PackageInfo) UpdateImportPath(pkg, newPath string) {
	if p.Imports == nil || pkg == "" || newPath == "" {
		return
	}

	newPath = p.toExternalGenPath(newPath)
	var prevPath string
OutLoop:
	for path, pkgSet := range p.Imports {
		for pkgKey := range pkgSet {
			if pkgKey == pkg {
				prevPath = path
				break OutLoop
			}
		}
	}
	if prevPath == "" {
		return
	}

	delete(p.Imports, prevPath)
	if newPath == pkg { // remove the alias
		p.Imports[newPath] = nil
	} else { // change the path -> alias mapping
		p.Imports[newPath] = map[string]bool{
			pkg: true,
		}
	}
}

func (p *PackageInfo) toExternalGenPath(path string) string {
	if p.ExternalKitexGen == "" || !strings.Contains(path, KitexGenPath) {
		return path
	}
	parts := strings.Split(path, KitexGenPath)
	newPath := util.JoinPath(p.ExternalKitexGen, parts[len(parts)-1])
	return newPath
}

// PkgInfo .
type PkgInfo struct {
	PkgName    string
	PkgRefName string
	ImportPath string
}

// ServiceInfo .
type ServiceInfo struct {
	PkgInfo
	ServiceName           string
	RawServiceName        string
	ServiceTypeName       func() string
	Base                  *ServiceInfo
	Methods               []*MethodInfo
	CombineServices       []*ServiceInfo
	HasStreaming          bool
	ServiceFilePath       string
	Protocol              string
	HandlerReturnKeepResp bool
	UseThriftReflection   bool
	// for multiple services scenario, the reference name for the service
	RefName string
	// identify whether this service would generate a corresponding handler.
	GenerateHandler bool
}

// AllMethods returns all methods that the service have.
func (s *ServiceInfo) AllMethods() (ms []*MethodInfo) {
	ms = append(ms, s.Methods...)
	for base := s.Base; base != nil; base = base.Base {
		ms = append(base.Methods, ms...)
	}
	return ms
}

// FixHasStreamingForExtendedService updates the HasStreaming field for extended services.
func (s *ServiceInfo) FixHasStreamingForExtendedService() {
	if s.Base == nil {
		return
	}
	if s.HasStreaming {
		return
	}
	s.Base.FixHasStreamingForExtendedService()
	s.HasStreaming = s.Base.HasStreamingRecursive()
}

// HasStreamingRecursive recursively check if the service has streaming method
func (s *ServiceInfo) HasStreamingRecursive() bool {
	if s.HasStreaming {
		return true
	}
	if s.Base == nil {
		return false
	}
	return s.Base.HasStreamingRecursive()
}

// MethodInfo .
type MethodInfo struct {
	PkgInfo                `json:"pkg_info"`
	ServiceName            string       `json:"service_name,omitempty"`
	Name                   string       `json:"name,omitempty"`
	RawName                string       `json:"raw_name,omitempty"`
	Oneway                 bool         `json:"oneway,omitempty"`
	Void                   bool         `json:"void,omitempty"`
	Args                   []*Parameter `json:"args,omitempty"`
	ArgsLength             int          `json:"args_length,omitempty"`
	Resp                   *Parameter   `json:"resp,omitempty"`
	Exceptions             []*Parameter `json:"exceptions,omitempty"`
	ArgStructName          string       `json:"arg_struct_name,omitempty"`
	ResStructName          string       `json:"res_struct_name,omitempty"`
	IsResponseNeedRedirect bool         `json:"is_response_need_redirect,omitempty"` // int -> int*
	GenArgResultStruct     bool         `json:"gen_arg_result_struct,omitempty"`
	IsStreaming            bool         `json:"is_streaming,omitempty"`
	ClientStreaming        bool         `json:"client_streaming,omitempty"`
	ServerStreaming        bool         `json:"server_streaming,omitempty"`
}

func (m *MethodInfo) StreamingMode() string {
	if !m.IsStreaming {
		return ""
	}
	if m.ClientStreaming {
		if m.ServerStreaming {
			return "bidirectional"
		}
		return "client"
	}
	if m.ServerStreaming {
		return "server"
	}
	return "unary"
}

// Parameter .
type Parameter struct {
	Deps    []PkgInfo
	Name    string
	RawName string
	Type    string // *PkgA.StructB
}

var funcs = map[string]interface{}{
	"ToLower":       strings.ToLower,
	"LowerFirst":    util.LowerFirst,
	"UpperFirst":    util.UpperFirst,
	"NotPtr":        util.NotPtr,
	"ReplaceString": util.ReplaceString,
	"SnakeString":   util.SnakeString,
	"HasFeature":    HasFeature,
	"FilterImports": FilterImports,
	"backquoted":    BackQuoted,
}

func AddTemplateFunc(key string, f interface{}) {
	funcs[key] = f
}

var templateNames = []string{
	"@client.go-NewClient-option",
	"@client.go-NewStreamClient-option",
	"@client.go-EOF",
	"@server.go-NewServer-option",
	"@server.go-EOF",
	"@invoker.go-NewInvoker-option",
	"@invoker.go-EOF",
}

func wrapTemplate(point, content string) string {
	return fmt.Sprintf(`{{define "%s"}}%s{{end}}`, point, content)
}

var templateExtensions = (func() map[string]string {
	m := make(map[string]string)
	for _, name := range templateNames {
		// create dummy templates
		m[name] = wrapTemplate(name, "")
	}
	return m
})()

// SetTemplateExtension .
func SetTemplateExtension(name, text string) {
	if _, ok := templateExtensions[name]; ok {
		templateExtensions[name] = text
	}
}

func applyExtension(name string, x *template.Template) (*template.Template, error) {
	var err error
	for _, n := range templateNames {
		x, err = x.Parse(templateExtensions[n])
		if err != nil {
			return nil, fmt.Errorf("failed to parse extension %q for %s: %w (%#q)",
				n, name, err, templateExtensions[n])
		}
	}
	return x, nil
}

// Task .
type Task struct {
	Name string
	Path string
	Text string
	*template.Template
	Ext *APIExtension
}

// Build .
func (t *Task) Build() error {
	x, err := template.New(t.Name).Funcs(funcs).Parse(t.Text)
	if err != nil {
		return err
	}
	// old fashion
	x, err = applyExtension(t.Name, x)
	if err != nil {
		return err
	}
	// new fashion
	for _, str := range t.makeExtension() {
		if x, err = x.Parse(str); err != nil {
			return err
		}
	}

	t.Template = x
	return nil
}

func fileTemplateExtension(name string) (option, eof string) {
	for _, tn := range templateNames {
		if strings.HasPrefix(tn, "@"+name+"-") {
			if strings.HasSuffix(tn, "-option") {
				option = tn
			} else if strings.HasSuffix(tn, "-EOF") {
				eof = tn
			}
		}
	}
	return
}

func (t *Task) makeExtension() (res []string) {
	if t.Ext == nil {
		return
	}
	p1, p2 := fileTemplateExtension(t.Name)
	if t.Ext.ExtendOption != "" {
		res = append(res, wrapTemplate(p1, t.Ext.ExtendOption))
	}
	if t.Ext.ExtendFile != "" {
		res = append(res, wrapTemplate(p2, t.Ext.ExtendFile))
	}
	return
}

// Render .
func (t *Task) Render(data interface{}) (*File, error) {
	if t.Template == nil {
		err := t.Build()
		if err != nil {
			return nil, err
		}
	}

	var buf bytes.Buffer
	err := t.ExecuteTemplate(&buf, t.Name, data)
	if err != nil {
		return nil, err
	}
	return &File{t.Path, buf.String()}, nil
}

func (t *Task) RenderString(data interface{}) (string, error) {
	if t.Template == nil {
		err := t.Build()
		if err != nil {
			return "", err
		}
	}

	var buf bytes.Buffer
	err := t.ExecuteTemplate(&buf, t.Name, data)
	if err != nil {
		return "", err
	}
	return buf.String(), nil
}

func FilterImports(Imports map[string]map[string]bool, ms []*MethodInfo) map[string]map[string]bool {
	res := map[string]map[string]bool{}
	for _, m := range ms {
		if m.Resp != nil {
			for _, dep := range m.Resp.Deps {
				if _, ok := Imports[dep.ImportPath]; ok {
					res[dep.ImportPath] = Imports[dep.ImportPath]
				}
			}
		}
		for _, arg := range m.Args {
			for _, dep := range arg.Deps {
				if _, ok := Imports[dep.ImportPath]; ok {
					res[dep.ImportPath] = Imports[dep.ImportPath]
				}
			}
		}
	}
	return res
}

func BackQuoted(s string) string {
	return "`" + s + "`"
}
