// Package migrations applies the canonical SQL migrations before the API starts.
// SQL in sql/ is generated from database/migrations by scripts/sync-go-migrations.mjs.
package migrations

import (
	"context"
	"crypto/sha256"
	"embed"
	"encoding/hex"
	"errors"
	"fmt"
	"io/fs"
	"log"
	"sort"
	"strconv"
	"strings"
	"time"

	"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/postgres"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
)

//go:embed sql/*.sql
var sqlFiles embed.FS

const migrationLockID int64 = 7308731946202608

type migration struct {
	version  string
	checksum string
	sql      string
}

// Run is a no-op unless DATABASE_AUTO_MIGRATE=true. Its context covers
// connection, advisory-lock waiting, every migration, and role provisioning.
func Run(ctx context.Context, getenv func(string) string) error {
	if getenv == nil {
		return errors.New("environment getter is required")
	}
	switch strings.ToLower(strings.TrimSpace(getenv("DATABASE_AUTO_MIGRATE"))) {
	case "", "false":
		return nil
	case "true":
	default:
		return errors.New("DATABASE_AUTO_MIGRATE must be true or false")
	}
	if strings.ToLower(strings.TrimSpace(getenv("ZHINIAN_DATA_BACKEND"))) != "postgres" {
		return errors.New("DATABASE_AUTO_MIGRATE requires ZHINIAN_DATA_BACKEND=postgres")
	}
	timeout := 180 * time.Second
	if raw := strings.TrimSpace(getenv("DATABASE_MIGRATION_TIMEOUT_SECONDS")); raw != "" {
		seconds, err := strconv.ParseInt(raw, 10, 64)
		if err != nil || seconds < 1 || seconds > 86400 {
			return errors.New("DATABASE_MIGRATION_TIMEOUT_SECONDS must be an integer from 1 to 86400")
		}
		timeout = time.Duration(seconds) * time.Second
	}
	role := strings.TrimSpace(getenv("DATABASE_APP_ROLE"))
	if role == "" && strings.EqualFold(strings.TrimSpace(getenv("NODE_ENV")), "production") {
		return errors.New("DATABASE_APP_ROLE is required in production")
	}
	if role != "" && !rolePattern.MatchString(role) {
		return errors.New("DATABASE_APP_ROLE must match [a-z_][a-z0-9_]{0,62}")
	}
	all, err := discover(sqlFiles)
	if err != nil {
		return err
	}
	migrationURL := strings.TrimSpace(getenv("DATABASE_MIGRATION_URL"))
	appURL := strings.TrimSpace(getenv("DATABASE_URL"))
	if migrationURL == "" {
		migrationURL = appURL
	}
	if migrationURL == "" {
		return errors.New("DATABASE_MIGRATION_URL or DATABASE_URL is required")
	}
	// Reuse the application's URI validation and forced transport policy.
	config, err := postgres.ParseConfig(func(key string) string {
		if key == "DATABASE_URL" {
			return migrationURL
		}
		return getenv(key)
	}, nil)
	if err != nil {
		return fmt.Errorf("migration database configuration: %w", err)
	}
	pgconfig, err := pgx.ParseConfig(config.DatabaseURL)
	if err != nil {
		return errors.New("invalid migration database connection configuration")
	}
	if appURL != "" && strings.TrimSpace(getenv("DATABASE_MIGRATION_URL")) != "" {
		appConfig, err := postgres.ParseConfig(func(key string) string {
			if key == "DATABASE_URL" {
				return appURL
			}
			return getenv(key)
		}, nil)
		if err != nil {
			return fmt.Errorf("application database configuration: %w", err)
		}
		appPGConfig, err := pgx.ParseConfig(appConfig.DatabaseURL)
		if err != nil {
			return errors.New("invalid application database connection configuration")
		}
		if !strings.EqualFold(pgconfig.Host, appPGConfig.Host) || pgconfig.Port != appPGConfig.Port || pgconfig.Database != appPGConfig.Database {
			return errors.New("DATABASE_MIGRATION_URL must target the same host, port, and database as DATABASE_URL")
		}
	}
	pgconfig.RuntimeParams["application_name"] = "zhinian-migrate"
	pgconfig.ConnectTimeout = config.ConnectionTimeout
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	log.Printf("zhinian-api checking database migrations (%d bundled versions)", len(all))
	conn, err := pgx.ConnectConfig(ctx, pgconfig)
	if err != nil {
		return safeDBError("connect", "", err)
	}
	defer conn.Close(context.Background())
	log.Print("zhinian-api waiting for database migration lock")
	if _, err := conn.Exec(ctx, "SELECT pg_advisory_lock($1)", migrationLockID); err != nil {
		return safeDBError("acquire migration lock", "", err)
	}
	// Session-level lock is released when conn closes, including on timeout.
	if err := apply(ctx, conn, all, role); err != nil {
		return err
	}
	log.Printf("zhinian-api database migrations ready (%d versions)", len(all))
	return nil
}

func apply(ctx context.Context, conn *pgx.Conn, all []migration, role string) error {
	if _, err := conn.Exec(ctx, `CREATE TABLE IF NOT EXISTS schema_migrations (
		version text PRIMARY KEY,
		checksum text NOT NULL,
		applied_at timestamptz NOT NULL DEFAULT now()
	)`); err != nil {
		return safeDBError("create migration tracking table", "", err)
	}
	rows, err := conn.Query(ctx, "SELECT version, checksum FROM schema_migrations")
	if err != nil {
		return safeDBError("read migration history", "", err)
	}
	applied := make(map[string]string)
	for rows.Next() {
		var version, checksum string
		if err := rows.Scan(&version, &checksum); err != nil {
			rows.Close()
			return safeDBError("read migration history", "", err)
		}
		applied[version] = checksum
	}
	err = rows.Err()
	rows.Close()
	if err != nil {
		return safeDBError("read migration history", "", err)
	}
	known := make(map[string]string, len(all))
	for _, item := range all {
		known[item.version] = item.checksum
	}
	for version, checksum := range applied {
		expected, exists := known[version]
		if !exists {
			return fmt.Errorf("applied migration %q is absent from bundle", version)
		}
		if checksum != expected {
			return fmt.Errorf("applied migration %q checksum differs from bundle", version)
		}
	}
	for _, item := range all {
		if _, exists := applied[item.version]; exists {
			continue
		}
		body, err := withoutOuterTransaction(item.sql)
		if err != nil {
			return fmt.Errorf("migration %s: %w", item.version, err)
		}
		tx, err := conn.Begin(ctx)
		if err != nil {
			return safeDBError("begin migration", item.version, err)
		}
		// PostgreSQL's simple protocol permits multi-statement migration files.
		_, err = tx.Exec(ctx, body, pgx.QueryExecModeSimpleProtocol)
		if err == nil {
			_, err = tx.Exec(ctx, "INSERT INTO schema_migrations(version, checksum) VALUES ($1, $2)", item.version, item.checksum)
		}
		if err == nil {
			err = tx.Commit(ctx)
		}
		if err != nil {
			_ = tx.Rollback(context.Background())
			return safeDBError("apply migration", item.version, err)
		}
		log.Printf("zhinian-api applied database migration %s", item.version)
	}
	if role != "" {
		if err := provisionRole(ctx, conn, role); err != nil {
			return err
		}
	}
	return nil
}

func discover(fsys fs.FS) ([]migration, error) {
	entries, err := fs.ReadDir(fsys, "sql")
	if err != nil {
		return nil, errors.New("read embedded migration bundle")
	}
	var all []migration
	for _, entry := range entries {
		name := entry.Name()
		if entry.IsDir() || !strings.HasSuffix(name, ".sql") {
			continue
		}
		data, err := fs.ReadFile(fsys, "sql/"+name)
		if err != nil {
			return nil, fmt.Errorf("read embedded migration %s", name)
		}
		checksum := sha256.Sum256(data)
		all = append(all, migration{strings.TrimSuffix(name, ".sql"), hex.EncodeToString(checksum[:]), string(data)})
	}
	sort.Slice(all, func(i, j int) bool { return all[i].version < all[j].version })
	if len(all) == 0 {
		return nil, errors.New("embedded migration bundle is empty")
	}
	return all, nil
}

// The canonical SQL may contain an explicit outer BEGIN/COMMIT. Strip only
// when both commands are the first and last non-comment statements; otherwise
// our transaction would commit before recording the checksum.
func withoutOuterTransaction(sql string) (string, error) {
	first, last := -1, -1
	lines := strings.SplitAfter(sql, "\n")
	for i, line := range lines {
		trimmed := strings.TrimSpace(line)
		if trimmed != "" && !strings.HasPrefix(trimmed, "--") {
			if first < 0 {
				first = i
			}
			last = i
		}
	}
	if first < 0 {
		return "", errors.New("empty migration SQL")
	}
	begin := strings.EqualFold(strings.TrimSpace(lines[first]), "begin;")
	commit := strings.EqualFold(strings.TrimSpace(lines[last]), "commit;")
	if begin != commit {
		return "", errors.New("outer transaction must contain both BEGIN and COMMIT")
	}
	if !begin {
		return sql, nil
	}
	lines[first], lines[last] = "", ""
	return strings.Join(lines, ""), nil
}

func safeDBError(operation, version string, err error) error {
	if version != "" {
		operation += " " + version
	}
	var pgErr *pgconn.PgError
	if errors.As(err, &pgErr) {
		return fmt.Errorf("%s failed (SQLSTATE %s)", operation, pgErr.Code)
	}
	if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
		return fmt.Errorf("%s timed out or was canceled", operation)
	}
	return fmt.Errorf("%s failed (%T)", operation, err)
}