package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"strconv"
	"strings"
	"time"

	_ "github.com/go-sql-driver/mysql"
)

const (
	temporaryMode  = "temporary"
	durableMode    = "durable"
	temporaryTable = "vehicle_alert_candidate_benchmark"
	durableTable   = "vehicle_alert_candidate_benchmark_durable"
	durableLock    = "lingniu.vehicle-alert-candidate-benchmark"
)

type benchmarkConfig struct {
	Rows           int
	BatchSize      int
	Mode           string
	ConfirmDurable bool
}

type benchmarkResult struct {
	Mode                  string            `json:"mode"`
	Rows                  int               `json:"rows"`
	BatchSize             int               `json:"batchSize"`
	Batches               int               `json:"batches"`
	Transactions          int               `json:"transactions"`
	DurationMS            int64             `json:"durationMs"`
	RowsPerSec            float64           `json:"rowsPerSec"`
	VerifiedRows          int               `json:"verifiedRows"`
	Temporary             bool              `json:"temporaryTable"`
	DurableWrite          bool              `json:"durableWrite"`
	CleanupMS             int64             `json:"cleanupMs"`
	CleanupVerified       bool              `json:"cleanupVerified"`
	GlobalStatusDelta     map[string]uint64 `json:"globalStatusDelta,omitempty"`
	GlobalStatusAvailable bool              `json:"globalStatusAvailable"`
}

func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, "alert benchmark failed:", err)
		os.Exit(1)
	}
}

func run() error {
	config := benchmarkConfig{}
	flag.IntVar(&config.Rows, "rows", 10_000, "candidate rows to write")
	flag.IntVar(&config.BatchSize, "batch-size", 500, "rows per insert")
	flag.StringVar(&config.Mode, "mode", temporaryMode, "storage mode: temporary or durable")
	flag.BoolVar(&config.ConfirmDurable, "confirm-durable-write", false, "required acknowledgement for durable writes and DDL")
	flag.Parse()
	if err := config.validate(); err != nil {
		return err
	}
	dsn := strings.TrimSpace(os.Getenv("MYSQL_DSN"))
	if dsn == "" {
		return fmt.Errorf("MYSQL_DSN is required")
	}
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return err
	}
	defer db.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()
	conn, err := db.Conn(ctx)
	if err != nil {
		return err
	}
	defer conn.Close()
	result, err := executeBenchmark(ctx, conn, config)
	if err != nil {
		return err
	}
	return json.NewEncoder(os.Stdout).Encode(result)
}

func (config benchmarkConfig) validate() error {
	config.Mode = strings.ToLower(strings.TrimSpace(config.Mode))
	if config.Rows < 1 || config.Rows > 100_000 || config.BatchSize < 1 || config.BatchSize > 1_000 {
		return fmt.Errorf("rows must be 1..100000 and batch-size 1..1000")
	}
	if config.Mode != temporaryMode && config.Mode != durableMode {
		return fmt.Errorf("mode must be temporary or durable")
	}
	if config.Mode == durableMode && !config.ConfirmDurable {
		return fmt.Errorf("durable mode creates, commits to, and drops a dedicated physical table; pass --confirm-durable-write")
	}
	return nil
}

func executeBenchmark(ctx context.Context, conn *sql.Conn, config benchmarkConfig) (result benchmarkResult, err error) {
	mode := strings.ToLower(strings.TrimSpace(config.Mode))
	table := temporaryTable
	createPrefix := "CREATE TEMPORARY TABLE"
	dropPrefix := "DROP TEMPORARY TABLE IF EXISTS"
	if mode == durableMode {
		table = durableTable
		createPrefix = "CREATE TABLE"
		dropPrefix = "DROP TABLE IF EXISTS"
		locked, lockErr := acquireBenchmarkLock(ctx, conn)
		if lockErr != nil {
			return result, fmt.Errorf("acquire durable benchmark lock: %w", lockErr)
		}
		if !locked {
			return result, fmt.Errorf("another durable benchmark is already running")
		}
		defer releaseBenchmarkLock(conn)
		if _, err = conn.ExecContext(ctx, dropPrefix+" "+table); err != nil {
			return result, fmt.Errorf("remove stale durable benchmark table: %w", err)
		}
	}
	if _, err = conn.ExecContext(ctx, createPrefix+" "+table+" LIKE vehicle_alert_candidate"); err != nil {
		return result, fmt.Errorf("create %s benchmark table: %w", mode, err)
	}
	cleanupStarted := time.Time{}
	cleanupComplete := false
	defer func() {
		cleanupStarted = time.Now()
		cleanupCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		_, cleanupErr := conn.ExecContext(cleanupCtx, dropPrefix+" "+table)
		if cleanupErr == nil {
			cleanupComplete = true
			if mode == durableMode {
				var remaining int
				cleanupErr = conn.QueryRowContext(cleanupCtx, `SELECT COUNT(*) FROM information_schema.tables WHERE table_schema=DATABASE() AND table_name=?`, table).Scan(&remaining)
				cleanupComplete = cleanupErr == nil && remaining == 0
			}
		}
		result.CleanupMS = time.Since(cleanupStarted).Milliseconds()
		result.CleanupVerified = cleanupComplete
		if err == nil && cleanupErr != nil {
			err = fmt.Errorf("cleanup benchmark table: %w", cleanupErr)
		} else if err == nil && !cleanupComplete {
			err = fmt.Errorf("benchmark table cleanup could not be verified")
		}
	}()

	statusBefore, statusBeforeErr := readGlobalStatus(ctx, conn)
	started := time.Now()
	tx, err := conn.BeginTx(ctx, nil)
	if err != nil {
		return result, err
	}
	now := time.Now().UTC()
	batches := 0
	for start := 0; start < config.Rows; start += config.BatchSize {
		end := min(start+config.BatchSize, config.Rows)
		query, args := buildCandidateInsert(table, start, end, now)
		if _, err = tx.ExecContext(ctx, query, args...); err != nil {
			tx.Rollback()
			return result, fmt.Errorf("insert batch %d: %w", batches+1, err)
		}
		batches++
	}
	if err = tx.Commit(); err != nil {
		return result, err
	}
	duration := time.Since(started)
	var verified int
	if err = conn.QueryRowContext(ctx, `SELECT COUNT(*) FROM `+table).Scan(&verified); err != nil {
		return result, err
	}
	if verified != config.Rows {
		return result, fmt.Errorf("row verification failed: wrote %d, found %d", config.Rows, verified)
	}
	statusAfter, statusAfterErr := readGlobalStatus(ctx, conn)
	statusDelta := map[string]uint64{}
	statusAvailable := statusBeforeErr == nil && statusAfterErr == nil
	if statusAvailable {
		for key, after := range statusAfter {
			if before, ok := statusBefore[key]; ok && after >= before {
				statusDelta[key] = after - before
			}
		}
	}
	result = benchmarkResult{
		Mode:                  mode,
		Rows:                  config.Rows,
		BatchSize:             config.BatchSize,
		Batches:               batches,
		Transactions:          1,
		DurationMS:            duration.Milliseconds(),
		RowsPerSec:            float64(config.Rows) / duration.Seconds(),
		VerifiedRows:          verified,
		Temporary:             mode == temporaryMode,
		DurableWrite:          mode == durableMode,
		GlobalStatusDelta:     statusDelta,
		GlobalStatusAvailable: statusAvailable,
	}
	return result, nil
}

func acquireBenchmarkLock(ctx context.Context, conn *sql.Conn) (bool, error) {
	var acquired sql.NullInt64
	if err := conn.QueryRowContext(ctx, `SELECT GET_LOCK(?,0)`, durableLock).Scan(&acquired); err != nil {
		return false, err
	}
	return acquired.Valid && acquired.Int64 == 1, nil
}

func releaseBenchmarkLock(conn *sql.Conn) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, _ = conn.ExecContext(ctx, `SELECT RELEASE_LOCK(?)`, durableLock)
}

func readGlobalStatus(ctx context.Context, conn *sql.Conn) (map[string]uint64, error) {
	rows, err := conn.QueryContext(ctx, `SHOW GLOBAL STATUS WHERE Variable_name IN ('Innodb_os_log_written','Binlog_cache_use','Binlog_cache_disk_use','Com_commit')`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	values := map[string]uint64{}
	for rows.Next() {
		var key, raw string
		if err := rows.Scan(&key, &raw); err != nil {
			return nil, err
		}
		value, parseErr := strconv.ParseUint(raw, 10, 64)
		if parseErr != nil {
			return nil, parseErr
		}
		values[key] = value
	}
	return values, rows.Err()
}

func buildCandidateInsert(table string, start, end int, now time.Time) (string, []any) {
	if table != temporaryTable && table != durableTable {
		panic("unsupported benchmark table")
	}
	values := make([]string, 0, end-start)
	args := make([]any, 0, (end-start)*7)
	for index := start; index < end; index++ {
		values = append(values, "(?,?,?,?,?,?,?)")
		args = append(args, "benchmark-rule", fmt.Sprintf("BENCH%017d", index), "JT808", now, now, float64(index%120), fmt.Sprintf("benchmark-event-%d", index))
	}
	return `INSERT INTO ` + table + `(rule_id,vin,protocol,first_matched_at,last_matched_at,latest_value,source_event_id) VALUES ` + strings.Join(values, ","), args
}