//go:build ignore

package main

import (
	"encoding/json"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

const inventoryPath = "guardrails/store/sql-ownership-inventory.json"

type inventory struct {
	MigratedTxStorePackages []string `json:"migratedTxStorePackages"`
}

type packageCheck struct {
	path      string
	files     []*ast.File
	store     *ast.StructType
	txStore   *ast.StructType
	storeMeth map[string]*ast.FuncDecl
	txMeth    map[string]*ast.FuncDecl
	errors    []string
}

func main() {
	inventory, err := loadInventory()
	if err != nil {
		fmt.Fprintf(os.Stderr, "store txstore pattern: %v\n", err)
		os.Exit(1)
	}
	migratedPackages := inventory.MigratedTxStorePackages
	if len(migratedPackages) == 0 {
		fmt.Fprintf(os.Stderr, "store txstore pattern: inventory %s has no migratedTxStorePackages\n", inventoryPath)
		os.Exit(1)
	}

	var failures []string
	for _, pkg := range migratedPackages {
		check, err := loadPackage(pkg)
		if err != nil {
			failures = append(failures, err.Error())
			continue
		}
		check.run()
		failures = append(failures, check.errors...)
	}
	if len(failures) > 0 {
		for _, failure := range failures {
			fmt.Fprintf(os.Stderr, "store txstore pattern: %s\n", failure)
		}
		os.Exit(1)
	}
	fmt.Printf("store txstore pattern: checked %s\n", strings.Join(migratedPackages, ", "))
}

func loadInventory() (*inventory, error) {
	data, err := os.ReadFile(inventoryPath)
	if err != nil {
		return nil, fmt.Errorf("read inventory %s: %w", inventoryPath, err)
	}
	var inventory inventory
	if err := json.Unmarshal(data, &inventory); err != nil {
		return nil, fmt.Errorf("parse inventory %s: %w", inventoryPath, err)
	}
	return &inventory, nil
}

func loadPackage(pkg string) (*packageCheck, error) {
	fset := token.NewFileSet()
	entries, err := os.ReadDir(pkg)
	if err != nil {
		return nil, fmt.Errorf("%s: read package: %w", pkg, err)
	}
	check := &packageCheck{
		path:      pkg,
		storeMeth: map[string]*ast.FuncDecl{},
		txMeth:    map[string]*ast.FuncDecl{},
	}
	for _, entry := range entries {
		name := entry.Name()
		if entry.IsDir() || !strings.HasSuffix(name, ".go") || strings.HasSuffix(name, "_test.go") {
			continue
		}
		filePath := filepath.Join(pkg, name)
		file, err := parser.ParseFile(fset, filePath, nil, 0)
		if err != nil {
			return nil, fmt.Errorf("%s: parse %s: %w", pkg, filePath, err)
		}
		check.files = append(check.files, file)
	}
	return check, nil
}

func (c *packageCheck) run() {
	c.collect()
	c.require(c.store != nil, "missing Store struct")
	c.require(c.txStore != nil, "missing TxStore struct")
	if c.txStore != nil {
		c.checkTxStoreFields()
	}
	c.checkMethodPairs()
	c.checkStoreMethods()
	c.checkTxStoreMethods()
}

func (c *packageCheck) collect() {
	for _, file := range c.files {
		for _, decl := range file.Decls {
			switch decl := decl.(type) {
			case *ast.GenDecl:
				for _, spec := range decl.Specs {
					typeSpec, ok := spec.(*ast.TypeSpec)
					if !ok {
						continue
					}
					structType, ok := typeSpec.Type.(*ast.StructType)
					if !ok {
						continue
					}
					switch typeSpec.Name.Name {
					case "Store":
						c.store = structType
					case "TxStore":
						c.txStore = structType
					}
				}
			case *ast.FuncDecl:
				receiver := receiverName(decl)
				if receiver == "" || !decl.Name.IsExported() {
					continue
				}
				switch receiver {
				case "Store":
					c.storeMeth[decl.Name.Name] = decl
				case "TxStore":
					c.txMeth[decl.Name.Name] = decl
				}
			}
		}
	}
}

func receiverName(fn *ast.FuncDecl) string {
	if fn.Recv == nil || len(fn.Recv.List) != 1 {
		return ""
	}
	switch expr := fn.Recv.List[0].Type.(type) {
	case *ast.Ident:
		return expr.Name
	case *ast.StarExpr:
		if ident, ok := expr.X.(*ast.Ident); ok {
			return ident.Name
		}
	}
	return ""
}

func (c *packageCheck) checkTxStoreFields() {
	for _, field := range c.txStore.Fields.List {
		if isCoreDBType(field.Type, "DB") {
			c.add("TxStore must not hold coredb.DB")
		}
		if isStarStore(field.Type) {
			c.add("TxStore must not hold *Store")
		}
	}
}

func isCoreDBType(expr ast.Expr, name string) bool {
	selector, ok := expr.(*ast.SelectorExpr)
	if !ok {
		return false
	}
	pkg, ok := selector.X.(*ast.Ident)
	return ok && pkg.Name == "coredb" && selector.Sel.Name == name
}

func isStarStore(expr ast.Expr) bool {
	star, ok := expr.(*ast.StarExpr)
	if !ok {
		return false
	}
	ident, ok := star.X.(*ast.Ident)
	return ok && ident.Name == "Store"
}

func (c *packageCheck) checkMethodPairs() {
	storeNames := exportedMethodNames(c.storeMeth, map[string]bool{"InTx": true, "Validate": true})
	txNames := exportedMethodNames(c.txMeth, nil)
	for _, name := range storeNames {
		if _, ok := c.txMeth[name]; !ok {
			c.add("Store.%s must have paired TxStore.%s", name, name)
		}
	}
	for _, name := range txNames {
		if _, ok := c.storeMeth[name]; !ok {
			c.add("TxStore.%s must have paired Store.%s", name, name)
		}
	}
}

func exportedMethodNames(methods map[string]*ast.FuncDecl, exclude map[string]bool) []string {
	names := make([]string, 0, len(methods))
	for name := range methods {
		if exclude[name] {
			continue
		}
		names = append(names, name)
	}
	sort.Strings(names)
	return names
}

func (c *packageCheck) checkStoreMethods() {
	for name, method := range c.storeMeth {
		if name == "InTx" || name == "Validate" {
			continue
		}
		if method.Body == nil {
			c.add("Store.%s must have a body", name)
			continue
		}
		hasWithTx := false
		hasInTx := false
		hasSameNameDelegation := false
		receiver := methodReceiverIdent(method)
		ast.Inspect(method.Body, func(node ast.Node) bool {
			call, ok := node.(*ast.CallExpr)
			if !ok {
				return true
			}
			selector, ok := call.Fun.(*ast.SelectorExpr)
			if !ok {
				return true
			}
			if isCoreDBSelector(selector, "WithTx") {
				hasWithTx = true
			}
			if isForbiddenSQLSelector(selector) {
				c.add("Store.%s must not directly call SQL helper %s", name, selector.Sel.Name)
			}
			if selector.Sel.Name == "InTx" {
				hasInTx = true
			}
			if selector.Sel.Name == name && receiverCallName(selector.X) == "InTx" {
				hasSameNameDelegation = true
			}
			if ident, ok := selector.X.(*ast.Ident); ok && ident.Name == receiver {
				if _, ok := c.storeMeth[selector.Sel.Name]; ok && selector.Sel.Name != "InTx" && selector.Sel.Name != "Validate" {
					c.add("Store.%s must not call Store.%s", name, selector.Sel.Name)
				}
			}
			return true
		})
		c.require(hasWithTx, "Store.%s must use coredb.WithTx", name)
		c.require(hasInTx, "Store.%s must call InTx", name)
		c.require(hasSameNameDelegation, "Store.%s must delegate to InTx(tx).%s", name, name)
	}
}

func methodReceiverIdent(fn *ast.FuncDecl) string {
	if fn.Recv == nil || len(fn.Recv.List) != 1 || len(fn.Recv.List[0].Names) != 1 {
		return ""
	}
	return fn.Recv.List[0].Names[0].Name
}

func receiverCallName(expr ast.Expr) string {
	call, ok := expr.(*ast.CallExpr)
	if !ok {
		return ""
	}
	selector, ok := call.Fun.(*ast.SelectorExpr)
	if !ok {
		return ""
	}
	return selector.Sel.Name
}

func isForbiddenSQLSelector(selector *ast.SelectorExpr) bool {
	forbidden := map[string]bool{
		"Exec":      true,
		"Query":     true,
		"QueryRow":  true,
		"Select":    true,
		"Get":       true,
		"NamedExec": true,
	}
	return forbidden[selector.Sel.Name] && !isCoreDBSelector(selector, "WithTx")
}

func (c *packageCheck) checkTxStoreMethods() {
	for name, method := range c.txMeth {
		if method.Body == nil {
			c.add("TxStore.%s must have a body", name)
			continue
		}
		ast.Inspect(method.Body, func(node ast.Node) bool {
			call, ok := node.(*ast.CallExpr)
			if !ok {
				return true
			}
			selector, ok := call.Fun.(*ast.SelectorExpr)
			if !ok {
				return true
			}
			if isCoreDBSelector(selector, "WithTx") || selector.Sel.Name == "BeginTx" || selector.Sel.Name == "Commit" || selector.Sel.Name == "Rollback" {
				c.add("TxStore.%s must not manage transaction lifecycle with %s", name, selector.Sel.Name)
			}
			if ident, ok := selector.X.(*ast.Ident); ok && (ident.Name == "store" || ident.Name == "s") {
				if _, ok := c.storeMeth[selector.Sel.Name]; ok {
					c.add("TxStore.%s must not call Store.%s", name, selector.Sel.Name)
				}
			}
			return true
		})
	}
}

func isCoreDBSelector(selector *ast.SelectorExpr, name string) bool {
	pkg, ok := selector.X.(*ast.Ident)
	return ok && pkg.Name == "coredb" && selector.Sel.Name == name
}

func (c *packageCheck) require(ok bool, format string, args ...any) {
	if !ok {
		c.add(format, args...)
	}
}

func (c *packageCheck) add(format string, args ...any) {
	c.errors = append(c.errors, fmt.Sprintf("%s: %s", c.path, fmt.Sprintf(format, args...)))
}