package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"go/format"
	"os"
	"path/filepath"
	"strings"
	"unicode"
)

type lifecycleQuery struct {
	name string
	sql  string
}

func main() {
	input := flag.String("input", "", "lifecycle SQL input file")
	output := flag.String("output", "", "generated Go output file")
	packageName := flag.String("package", "", "generated Go package name")
	flag.Parse()

	if strings.TrimSpace(*input) == "" {
		exitWithError(errors.New("input is required"))
	}
	if strings.TrimSpace(*output) == "" {
		exitWithError(errors.New("output is required"))
	}
	if strings.TrimSpace(*packageName) == "" {
		exitWithError(errors.New("package is required"))
	}
	source, err := os.ReadFile(*input)
	if err != nil {
		exitWithError(fmt.Errorf("read input: %w", err))
	}
	generated, err := generateLifecycleGo(source, *packageName)
	if err != nil {
		exitWithError(err)
	}
	if err := os.MkdirAll(filepath.Dir(*output), 0o755); err != nil {
		exitWithError(fmt.Errorf("create output dir: %w", err))
	}
	if err := os.WriteFile(*output, generated, 0o644); err != nil {
		exitWithError(fmt.Errorf("write output: %w", err))
	}
}

func exitWithError(err error) {
	_, _ = fmt.Fprintf(os.Stderr, "sqlitelifecyclegen: %v\n", err)
	os.Exit(1)
}

func generateLifecycleGo(source []byte, packageName string) ([]byte, error) {
	queries, err := parseLifecycleQueries(string(source))
	if err != nil {
		return nil, err
	}
	if len(queries) != 1 {
		return nil, fmt.Errorf("expected exactly one lifecycle query, got %d", len(queries))
	}
	query := queries[0]
	if query.name != "DeferForeignKeys" {
		return nil, fmt.Errorf("unsupported lifecycle query %q", query.name)
	}
	var out bytes.Buffer
	fmt.Fprintf(&out, "// Code generated by server/metadata/lifecyclegen. DO NOT EDIT.\n\n")
	fmt.Fprintf(&out, "package %s\n\n", packageName)
	fmt.Fprintf(&out, "import (\n")
	fmt.Fprintf(&out, "\t%q\n", "context")
	fmt.Fprintf(&out, "\t%q\n", "database/sql")
	fmt.Fprintf(&out, ")\n\n")
	fmt.Fprintf(&out, "type DBTX interface {\n")
	fmt.Fprintf(&out, "\tExecContext(context.Context, string, ...interface{}) (sql.Result, error)\n")
	fmt.Fprintf(&out, "}\n\n")
	fmt.Fprintf(&out, "func New(db DBTX) *Queries {\n")
	fmt.Fprintf(&out, "\treturn &Queries{db: db}\n")
	fmt.Fprintf(&out, "}\n\n")
	fmt.Fprintf(&out, "type Queries struct {\n")
	fmt.Fprintf(&out, "\tdb DBTX\n")
	fmt.Fprintf(&out, "}\n\n")
	fmt.Fprintf(&out, "func (q *Queries) WithTx(tx *sql.Tx) *Queries {\n")
	fmt.Fprintf(&out, "\treturn &Queries{db: tx}\n")
	fmt.Fprintf(&out, "}\n\n")
	fmt.Fprintf(&out, "const deferForeignKeys = %q\n\n", query.sql)
	fmt.Fprintf(&out, "func (q *Queries) DeferForeignKeys(ctx context.Context) error {\n")
	fmt.Fprintf(&out, "\t_, err := q.db.ExecContext(ctx, deferForeignKeys)\n")
	fmt.Fprintf(&out, "\treturn err\n")
	fmt.Fprintf(&out, "}\n")
	formatted, err := format.Source(out.Bytes())
	if err != nil {
		return nil, fmt.Errorf("format generated source: %w", err)
	}
	return formatted, nil
}

func parseLifecycleQueries(source string) ([]lifecycleQuery, error) {
	lines := strings.Split(source, "\n")
	queries := make([]lifecycleQuery, 0)
	var activeName string
	activeSQL := make([]string, 0)
	flush := func() error {
		if activeName == "" {
			return nil
		}
		sql := strings.TrimSpace(strings.Join(activeSQL, "\n"))
		if sql == "" {
			return fmt.Errorf("lifecycle query %q has empty body", activeName)
		}
		queries = append(queries, lifecycleQuery{name: activeName, sql: sql})
		activeName = ""
		activeSQL = activeSQL[:0]
		return nil
	}
	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "-- name:") {
			if err := flush(); err != nil {
				return nil, err
			}
			name, execOnly, err := parseLifecycleName(trimmed)
			if err != nil {
				return nil, err
			}
			if !execOnly {
				return nil, fmt.Errorf("lifecycle query %q must use exec annotation", name)
			}
			activeName = name
			continue
		}
		if activeName != "" {
			activeSQL = append(activeSQL, line)
		}
	}
	if err := flush(); err != nil {
		return nil, err
	}
	return queries, nil
}

func parseLifecycleName(line string) (string, bool, error) {
	fields := strings.Fields(strings.TrimPrefix(line, "-- name:"))
	if len(fields) != 2 {
		return "", false, fmt.Errorf("invalid lifecycle query declaration %q", line)
	}
	if !isExportedIdentifier(fields[0]) {
		return "", false, fmt.Errorf("invalid lifecycle query name %q", fields[0])
	}
	return fields[0], strings.EqualFold(fields[1], ":exec"), nil
}

func isExportedIdentifier(name string) bool {
	if name == "" {
		return false
	}
	for i, r := range name {
		switch {
		case i == 0:
			if !unicode.IsUpper(r) {
				return false
			}
		case !(r == '_' || unicode.IsLetter(r) || unicode.IsDigit(r)):
			return false
		}
	}
	return true
}
