//go:build ignore

package main

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

const (
	sqlOwnershipInventoryPath = "guardrails/store/sql-ownership-inventory.json"
	sqlOwnershipSpecPath      = "specs/storage/sql-ownership-inventory.md"
)

type sqlOwnershipInventory struct {
	MigratedTxStorePackages          []string                     `json:"migratedTxStorePackages"`
	CanonicalProductionSQLDirs       []string                     `json:"canonicalProductionSQLDirectories"`
	TestAndHarnessSQLDirs            []string                     `json:"testAndHarnessSQLDirectories"`
	TemporaryProductionSQLExceptions []sqlOwnershipExceptionEntry `json:"temporaryProductionSQLExceptions"`
}

type sqlOwnershipExceptionEntry struct {
	Path   string `json:"path"`
	Reason string `json:"reason"`
}

type sqlOwnershipCandidate struct {
	rel       string
	dir       string
	importsDB bool
	hasSQL    bool
}

type countedDirectory struct {
	path  string
	count int
}

func main() {
	inventory, err := loadSQLOwnershipInventory()
	if err != nil {
		fail("%v", err)
	}
	if _, err := os.Stat(sqlOwnershipSpecPath); err != nil {
		fail("missing canonical spec %s: %v", sqlOwnershipSpecPath, err)
	}

	owners, err := scanSQLOwners(inventory)
	if err != nil {
		fail("%v", err)
	}
	sort.Strings(owners)

	exceptionReasons := make(map[string]string, len(inventory.TemporaryProductionSQLExceptions))
	for _, entry := range inventory.TemporaryProductionSQLExceptions {
		path := filepath.ToSlash(strings.TrimSpace(entry.Path))
		if path == "" {
			fail("inventory %s contains an empty temporaryProductionSQLExceptions path", sqlOwnershipInventoryPath)
		}
		if _, exists := exceptionReasons[path]; exists {
			fail("inventory %s duplicates temporaryProductionSQLExceptions entry %s", sqlOwnershipInventoryPath, path)
		}
		exceptionReasons[path] = strings.TrimSpace(entry.Reason)
	}

	directoryCounts := make(map[string]int)
	var harnessOwners []string
	var exceptionOwners []string
	var unexpectedOwners []string
	for _, owner := range owners {
		directoryCounts[filepath.ToSlash(filepath.Dir(owner))]++
		switch {
		case matchesAnyDir(owner, inventory.TestAndHarnessSQLDirs):
			harnessOwners = append(harnessOwners, owner)
		case matchesAnyDir(owner, inventory.CanonicalProductionSQLDirs):
			// canonical; reported in directory summary only
		case exceptionReasons[owner] != "":
			exceptionOwners = append(exceptionOwners, owner)
		default:
			unexpectedOwners = append(unexpectedOwners, owner)
		}
	}

	var staleExceptions []string
	for path := range exceptionReasons {
		if !contains(ownerSet(owners), path) {
			staleExceptions = append(staleExceptions, path)
		}
	}
	sort.Strings(staleExceptions)

	fmt.Printf("store sql ownership spec: %s\n", sqlOwnershipSpecPath)
	fmt.Printf("store sql ownership inventory: %s\n", sqlOwnershipInventoryPath)
	fmt.Printf("migrated txstore packages: %s\n", strings.Join(inventory.MigratedTxStorePackages, ", "))

	fmt.Println("\nSQL-owning directories:")
	for _, item := range sortDirectoryCounts(directoryCounts) {
		fmt.Printf("%4d %s\n", item.count, item.path)
	}

	fmt.Println("\nTest/harness SQL owners:")
	if len(harnessOwners) == 0 {
		fmt.Println("none")
	} else {
		for _, owner := range harnessOwners {
			fmt.Printf("- %s\n", owner)
		}
	}

	fmt.Println("\nTemporary production SQL exceptions:")
	if len(inventory.TemporaryProductionSQLExceptions) == 0 {
		fmt.Println("none")
	} else {
		for _, entry := range inventory.TemporaryProductionSQLExceptions {
			fmt.Printf("- %s — %s\n", entry.Path, entry.Reason)
		}
	}

	if len(staleExceptions) > 0 {
		fmt.Fprintln(os.Stderr, "\nstale SQL ownership exception entries:")
		for _, path := range staleExceptions {
			fmt.Fprintf(os.Stderr, "- %s\n", path)
		}
	}

	if len(unexpectedOwners) > 0 {
		fmt.Fprintln(os.Stderr, "\nunexpected production SQL owners outside canonical/test/allowlisted locations:")
		for _, owner := range unexpectedOwners {
			fmt.Fprintf(os.Stderr, "- %s\n", owner)
		}
	}

	if len(staleExceptions) > 0 || len(unexpectedOwners) > 0 {
		os.Exit(1)
	}

	fmt.Printf("\nstore sql ownership: ok (%d SQL-owning files)\n", len(owners))
}

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

func scanSQLOwners(inventory *sqlOwnershipInventory) ([]string, error) {
	var candidates []sqlOwnershipCandidate
	var sqlFiles []string

	err := filepath.Walk(".", func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		rel := filepath.ToSlash(path)
		if rel == "." {
			return nil
		}
		if info.IsDir() {
			switch rel {
			case ".git", ".worktrees", "vendor", "x":
				return filepath.SkipDir
			}
			return nil
		}
		if strings.HasSuffix(rel, "_test.go") {
			return nil
		}
		switch {
		case strings.HasSuffix(rel, ".sql"):
			sqlFiles = append(sqlFiles, rel)
			return nil
		case strings.HasSuffix(rel, ".go"):
			candidate, err := scanGoFile(rel)
			if err != nil {
				return err
			}
			candidates = append(candidates, candidate)
			return nil
		default:
			return nil
		}
	})
	if err != nil {
		return nil, err
	}

	dbImportDirs := make(map[string]bool)
	for _, candidate := range candidates {
		if candidate.importsDB {
			dbImportDirs[candidate.dir] = true
		}
	}

	var owners []string
	for _, candidate := range candidates {
		if !candidate.hasSQL {
			continue
		}
		if dbImportDirs[candidate.dir] || isKnownSQLZone(candidate.rel, inventory) {
			owners = append(owners, candidate.rel)
		}
	}
	owners = append(owners, sqlFiles...)
	return owners, nil
}

func scanGoFile(rel string) (sqlOwnershipCandidate, error) {
	path := filepath.FromSlash(rel)
	fset := token.NewFileSet()
	file, err := parser.ParseFile(fset, path, nil, parser.SkipObjectResolution)
	if err != nil {
		return sqlOwnershipCandidate{}, fmt.Errorf("parse %s: %w", rel, err)
	}

	candidate := sqlOwnershipCandidate{
		rel: rel,
		dir: filepath.ToSlash(filepath.Dir(rel)),
	}
	for _, imp := range file.Imports {
		importPath, err := strconv.Unquote(imp.Path.Value)
		if err != nil {
			return sqlOwnershipCandidate{}, fmt.Errorf("parse import path %s in %s: %w", imp.Path.Value, rel, err)
		}
		if importsDatabasePackage(importPath) {
			candidate.importsDB = true
			break
		}
	}

	ast.Inspect(file, func(node ast.Node) bool {
		lit, ok := node.(*ast.BasicLit)
		if !ok || lit.Kind != token.STRING {
			return true
		}
		value, err := strconv.Unquote(lit.Value)
		if err != nil {
			return true
		}
		if looksLikeSQL(value) {
			candidate.hasSQL = true
			return false
		}
		return true
	})

	return candidate, nil
}

func importsDatabasePackage(path string) bool {
	switch path {
	case "database/sql", "github.com/jmoiron/sqlx", "github.com/tidbcloud/AgentTeam/internal/coredb":
		return true
	default:
		return false
	}
}

func looksLikeSQL(value string) bool {
	normalized := strings.ToUpper(strings.Join(strings.Fields(value), " "))
	if normalized == "" {
		return false
	}
	switch {
	case isSelectStatement(normalized):
		return true
	case isInsertStatement(normalized):
		return true
	case strings.HasPrefix(normalized, "REPLACE INTO "):
		return true
	case strings.HasPrefix(normalized, "UPDATE "):
		return strings.Contains(normalized, " SET ") &&
			(strings.Contains(normalized, " WHERE ") ||
				strings.Contains(normalized, " RETURNING ") ||
				strings.Contains(normalized, " FROM ") ||
				strings.Contains(normalized, " LIMIT ") ||
				strings.Contains(normalized, " NOW(") ||
				strings.Contains(normalized, " CURRENT_TIMESTAMP") ||
				strings.Contains(normalized, " ON DUPLICATE KEY UPDATE") ||
				strings.Contains(normalized, "$"))
	case strings.HasPrefix(normalized, "DELETE FROM "):
		return true
	case strings.HasPrefix(normalized, "WITH "):
		return strings.Contains(normalized, " SELECT ") ||
			strings.Contains(normalized, " UPDATE ") ||
			strings.Contains(normalized, " INSERT ") ||
			strings.Contains(normalized, " DELETE ")
	case strings.HasPrefix(normalized, "ALTER TABLE "),
		strings.HasPrefix(normalized, "CREATE TABLE "),
		strings.HasPrefix(normalized, "CREATE INDEX "),
		strings.HasPrefix(normalized, "DROP TABLE "),
		strings.HasPrefix(normalized, "DROP INDEX "),
		strings.HasPrefix(normalized, "TRUNCATE TABLE "):
		return true
	default:
		return false
	}
}

func isSelectStatement(normalized string) bool {
	if !strings.HasPrefix(normalized, "SELECT ") {
		return false
	}
	if strings.Contains(normalized, " FROM ") ||
		strings.HasSuffix(normalized, " FROM") ||
		strings.Contains(normalized, " JOIN ") ||
		strings.Contains(normalized, " WHERE ") ||
		strings.Contains(normalized, " GROUP BY ") ||
		strings.Contains(normalized, " ORDER BY ") ||
		strings.Contains(normalized, " HAVING ") ||
		strings.Contains(normalized, " FOR UPDATE") ||
		strings.Contains(normalized, " GET_LOCK(") ||
		strings.Contains(normalized, " RELEASE_LOCK(") ||
		strings.Contains(normalized, " DATABASE(") ||
		strings.Contains(normalized, " COUNT(") ||
		strings.Contains(normalized, " EXISTS ") ||
		strings.Contains(normalized, " CAST(") {
		return true
	}

	fields := strings.Fields(normalized)
	if len(fields) < 2 {
		return false
	}
	if strings.Contains(normalized, "(") ||
		strings.Contains(normalized, ",") ||
		strings.Contains(normalized, " AS ") ||
		strings.HasPrefix(fields[1], "$") ||
		strings.HasPrefix(fields[1], "@@") {
		return true
	}

	switch fields[1] {
	case "CURRENT_TIMESTAMP", "CURRENT_DATE", "CURRENT_TIME", "NOW()", "NULL":
		return true
	default:
		return false
	}
}

func isInsertStatement(normalized string) bool {
	fields := strings.Fields(normalized)
	if len(fields) < 2 || fields[0] != "INSERT" {
		return false
	}

	index := 1
	for index < len(fields) {
		switch fields[index] {
		case "IGNORE", "LOW_PRIORITY", "DELAYED", "HIGH_PRIORITY":
			index++
			continue
		default:
			return fields[index] == "INTO"
		}
	}

	return false
}

func isKnownSQLZone(rel string, inventory *sqlOwnershipInventory) bool {
	if matchesAnyDir(rel, inventory.CanonicalProductionSQLDirs) || matchesAnyDir(rel, inventory.TestAndHarnessSQLDirs) {
		return true
	}
	for _, entry := range inventory.TemporaryProductionSQLExceptions {
		if filepath.ToSlash(strings.TrimSpace(entry.Path)) == rel {
			return true
		}
	}
	return false
}

func matchesAnyDir(rel string, dirs []string) bool {
	for _, dir := range dirs {
		trimmed := filepath.ToSlash(strings.TrimSpace(dir))
		if trimmed == "" {
			continue
		}
		if rel == trimmed || strings.HasPrefix(rel, trimmed+"/") {
			return true
		}
	}
	return false
}

func sortDirectoryCounts(counts map[string]int) []countedDirectory {
	items := make([]countedDirectory, 0, len(counts))
	for path, count := range counts {
		items = append(items, countedDirectory{path: path, count: count})
	}
	sort.Slice(items, func(i, j int) bool {
		if items[i].count == items[j].count {
			return items[i].path < items[j].path
		}
		return items[i].count > items[j].count
	})
	return items
}

func ownerSet(owners []string) map[string]bool {
	set := make(map[string]bool, len(owners))
	for _, owner := range owners {
		set[owner] = true
	}
	return set
}

func contains(set map[string]bool, path string) bool {
	return set[path]
}

func fail(format string, args ...any) {
	fmt.Fprintf(os.Stderr, "store sql ownership: "+format+"\n", args...)
	os.Exit(1)
}