// Copyright 2025 Florian Zenker (flo@znkr.io)
//
// 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 (
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"slices"

	"golang.org/x/tools/go/ast/astutil"
)

const header = `// generated by znkr.io/diff/internal/cmd/specializemyers
// DO NOT EDIT

`

func specialize(filename string) ([]byte, error) {
	fset := token.NewFileSet()
	file, err := parser.ParseFile(fset, filename, nil, 0)
	if err != nil {
		return nil, fmt.Errorf("parsing input: %v", err)
	}

	file = astutil.Apply(file, func(c *astutil.Cursor) bool {
		switch n := c.Node().(type) {
		case *ast.TypeSpec:
			if isMyersT(n) {
				n.Name.Name = "myersInt"
				n.TypeParams = nil

				for _, field := range n.Type.(*ast.StructType).Fields.List {
					field.Type = specializeT(field.Type)
				}
			}
			return false

		case *ast.FuncDecl:
			if isMyresTMethod(n) {
				if !methodsOfInterest[n.Name.Name] {
					return false
				}
				n.Recv.List[0].Type.(*ast.StarExpr).X = &ast.Ident{Name: "myersInt"}
				local := n.Recv.List[0].Names[0].Name
				n.Type.Params.List = slices.DeleteFunc(n.Type.Params.List, func(f *ast.Field) bool {
					if len(f.Names) != 1 {
						return false
					}
					return f.Names[0].Name == "eq" && isEqFunc(f.Type)
				})
				for _, param := range n.Type.Params.List {
					param.Type = specializeT(param.Type)
				}

				n.Body = astutil.Apply(n.Body, func(c *astutil.Cursor) bool {
					call, ok := c.Node().(*ast.CallExpr)
					if !ok {
						return true
					}
					switch fun := call.Fun.(type) {
					case *ast.SelectorExpr:
						if x, ok := fun.X.(*ast.Ident); !ok || x.Name != local {
							return true
						}
						if !methodsOfInterest[n.Name.Name] {
							return true
						}
						call.Args = slices.DeleteFunc(call.Args, func(arg ast.Expr) bool {
							name, ok := arg.(*ast.Ident)
							return ok && name.Name == "eq"
						})
					case *ast.Ident:
						if fun.Name != "eq" {
							return false
						}
						if len(call.Args) != 2 {
							return false
						}
						c.Replace(&ast.BinaryExpr{
							X:  call.Args[0],
							Op: token.EQL,
							Y:  call.Args[1],
						})
					}
					return true
				}, nil).(*ast.BlockStmt)
			}
			return false

		default:
			return true
		}
	}, nil).(*ast.File)

	var buf bytes.Buffer
	buf.WriteString(header)
	if err := format.Node(&buf, fset, file); err != nil {
		return nil, fmt.Errorf("formatting result: %v", err)
	}
	return buf.Bytes(), nil
}

var methodsOfInterest = map[string]bool{
	"split":   true,
	"compare": true,
	"init":    true,
}

func isMyersT(ts *ast.TypeSpec) bool {
	if ts.Name.Name != "myers" {
		return false
	}

	if _, ok := ts.Type.(*ast.StructType); !ok {
		return false
	}

	return matchFields(ts.TypeParams, []fmatcher{
		0: func(name string, typ ast.Expr) bool {
			tname, ok := typ.(*ast.Ident)
			return ok && name == "T" && tname.Name == "any"
		},
	})
}

func isMyresTMethod(fd *ast.FuncDecl) bool {
	return matchFields(fd.Recv, []fmatcher{
		0: func(name string, typ ast.Expr) bool {
			star, ok := typ.(*ast.StarExpr)
			if !ok {
				return false
			}
			x, ok := star.X.(*ast.IndexExpr)
			if !ok {
				return false
			}
			tname, ok := x.X.(*ast.Ident)
			if !ok {
				return false
			}
			index, ok := x.Index.(*ast.Ident)
			if !ok {
				return false
			}
			return tname.Name == "myers" && index.Name == "T"
		},
	})
}

func isEqFunc(expr ast.Expr) bool {
	typ, ok := expr.(*ast.FuncType)
	if !ok {
		return false
	}
	isTypT := func(_ string, typ ast.Expr) bool {
		tname, ok := typ.(*ast.Ident)
		return ok && tname.Name == "T"
	}
	params := matchFields(typ.Params, []fmatcher{
		0: isTypT,
		1: isTypT,
	})
	ret := matchFields(typ.Results, []fmatcher{
		0: func(_ string, typ ast.Expr) bool {
			tname, ok := typ.(*ast.Ident)
			return ok && tname.Name == "bool"
		},
	})
	return params && ret
}

func specializeT[N ast.Node](n N) N {
	return astutil.Apply(n, func(c *astutil.Cursor) bool {
		if ident, ok := c.Node().(*ast.Ident); ok && ident.Name == "T" {
			ident.Name = "int"
			return false
		}
		return true
	}, nil).(N)
}

type fmatcher func(name string, typ ast.Expr) bool

func matchFields(fl *ast.FieldList, matchers []fmatcher) bool {
	if len(fl.List) == 0 && len(matchers) != 0 {
		return false
	}
	i := 0
	for _, f := range fl.List {
		if len(f.Names) == 0 {
			if i >= len(matchers) {
				return false
			}
			if !matchers[i]("", f.Type) {
				return false
			}
			i++
		}
		for _, name := range f.Names {
			if i >= len(matchers) {
				return false
			}
			if !matchers[i](name.Name, f.Type) {
				return false
			}
			i++
		}
	}
	return i == len(matchers)
}
