// Copyright 2025 Dolthub, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"context"
	"errors"
	"fmt"
	"math/rand/v2"
	"strings"
	"testing"
	"time"

	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/sync/errgroup"

	driver "github.com/dolthub/doltgresql/integration-tests/go-sql-server-driver/driver"
)

func TestStatsGCConcurrency(t *testing.T) {
	// Doltgres' statistics provider does not implement the ExtendedStatsProvider
	// interface that dolt_stats_gc() requires ("provider does not implement
	// ExtendedStatsProvider"), so this test cannot run yet. The logic below is a
	// faithful port of the Dolt test and should be enabled once Doltgres supports
	// statistics GC.
	t.Skip("dolt_stats_gc() is unsupported in Doltgres (statistics provider does not implement ExtendedStatsProvider)")
	t.Parallel()
	// At one time, it was possible for a GC of stats to block the
	// analyzer from accessing stats for a long time. This test sets
	// up the following condition:
	//
	// * Several branches, each with 10 tables, each with several indexes.
	//   All of the branches and tables and indexes are non-empty.
	//
	// * One thread calls `select dolt_stats_gc()`
	//
	// * Other threads continue to connect to the database and run
	//   queries.
	//
	// * We expect all queries to complete without much delay. If
	//   there is a long delay on the read queries, this fails the
	//   test.
	ctx := context.Background()
	u, err := driver.NewDoltUser()
	require.NoError(t, err)
	t.Cleanup(func() {
		u.Cleanup()
	})

	rs, err := u.MakeRepoStore()
	require.NoError(t, err)
	_, err = rs.MakeRepo("stats_gc_concurrency_test")
	require.NoError(t, err)

	srvSettings := &driver.Server{
		Args:        []string{},
		DynamicPort: "server",
	}
	ports := newPorts(t)
	server := StartServer(t, rs, "stats_gc_concurrency_test", srvSettings, ports)

	// Connect to the database
	db, err := server.DB(driver.Connection{})
	require.NoError(t, err)
	t.Cleanup(func() {
		db.Close()
	})
	db.SetMaxOpenConns(32)
	db.SetMaxIdleConns(32)

	numbranches := 3

	var branches = []string{"main"}
	for i := 0; i < numbranches; i++ {
		branches = append(branches, uuid.New().String())
	}

	var tables = []string{
		"t0", "t1", "t2", "t3", "t4",
		"t5", "t6", "t7", "t8", "t9",
	}
	var columns = []string{
		"c0", "c1", "c2", "c3", "c4",
		"c5", "c6", "c7", "c8", "c9",
	}
	var columndefs = []string{}
	for _, col := range columns {
		columndefs = append(columndefs, fmt.Sprintf("%s bigint unique", col))
	}

	for _, table := range tables {
		_, err = db.ExecContext(ctx, fmt.Sprintf(strings.Join([]string{
			"CREATE TABLE %s (",
			"id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,",
			strings.Join(columndefs, ",\n"),
			");",
		}, "\n"), table))
		require.NoError(t, err)
	}

	_, err = db.ExecContext(ctx, "select dolt_commit('-Am', 'initial commit')")
	require.NoError(t, err)

	for _, branch := range branches[1:] {
		_, err = db.ExecContext(ctx, fmt.Sprintf("select dolt_checkout('-b', '%s')", branch))
		require.NoError(t, err)
		for _, table := range tables {
			_, err = db.ExecContext(ctx, fmt.Sprintf("insert into %s values %s", table, t0_values_str()))
			require.NoError(t, err)
		}
	}

	_, err = db.ExecContext(ctx, "select dolt_checkout('main')")
	require.NoError(t, err)
	for _, table := range tables {
		_, err = db.ExecContext(ctx, fmt.Sprintf("insert into %s values %s", table, t0_values_str()))
		require.NoError(t, err)
	}

	numThreads := 10
	// The workers will run until this is closed. This will be closed after the stats_gc() is finished.
	stop := make(chan struct{})
	maxQuery := make([]time.Duration, numThreads)
	eg, egCtx := errgroup.WithContext(ctx)
	for i := 0; i < numThreads; i++ {
		eg.Go(func() error {
			for {
				select {
				case <-stop:
					return nil
				default:
				}

				err := func() (err error) {
					start := time.Now()
					defer func() {
						elapsed := time.Since(start)
						if elapsed > maxQuery[i] {
							maxQuery[i] = elapsed
						}
					}()
					conn, err := db.Conn(egCtx)
					if err != nil {
						return err
					}
					defer func() {
						err = errors.Join(err, conn.Close())
					}()
					branch := branches[rand.IntN(len(branches))]
					_, err = conn.ExecContext(egCtx, fmt.Sprintf("select dolt_checkout('%s')", branch))
					if err != nil {
						return err
					}
					t1 := tables[rand.IntN(len(tables))]
					t2 := tables[rand.IntN(len(tables))]
					for t2 == t1 {
						t2 = tables[rand.IntN(len(tables))]
					}
					c1 := columns[rand.IntN(len(columns))]
					c2 := columns[rand.IntN(len(columns))]
					rows, err := conn.QueryContext(egCtx, fmt.Sprintf("select * from %s join %s on %s.%s = %s.%s",
						t1, t2, t1, c1, t2, c2))
					if err != nil {
						return err
					}
					defer func() {
						err = errors.Join(err, rows.Close())
					}()
					for rows.Next() {
					}
					return nil
				}()
				if err != nil {
					return err
				}
			}
		})
	}

	var gcDuration time.Duration
	eg.Go(func() (err error) {
		defer close(stop)
		start := time.Now()
		defer func() {
			gcDuration = time.Since(start)
		}()
		conn, err := db.Conn(egCtx)
		if err != nil {
			return err
		}
		defer func() {
			err = errors.Join(err, conn.Close())
		}()
		_, err = conn.ExecContext(egCtx, "select dolt_stats_gc()")
		if err != nil {
			return err
		}
		return nil
	})

	require.NoError(t, eg.Wait())

	for _, d := range maxQuery {
		assert.Greater(t, 3*time.Second, d)
		assert.Greater(t, gcDuration, d)
	}
}

func TestStatsAnalyzeTableSpeed(t *testing.T) {
	// Doltgres does not provide the dolt_stats_timers() procedure that this
	// test uses to slow background stats refresh, so the rate-limiting scenario
	// it asserts cannot be set up. The logic below is a faithful port and
	// should be enabled once Doltgres supports dolt_stats_timers().
	t.Skip("dolt_stats_timers() is unsupported in Doltgres")
	t.Parallel()
	// At one time, calling ANALYZE would go through the rate-limiting system
	// that background stats refresh goes through. This would make ANALYZE take
	// a longer time than necessary, and it would be very dependent on the stats
	// timers that were in effect.
	//
	// This tests that even with very large stats timers, ANALYZE behaves
	// reasonably.
	ctx := context.Background()
	u, err := driver.NewDoltUser()
	require.NoError(t, err)
	t.Cleanup(func() {
		u.Cleanup()
	})

	rs, err := u.MakeRepoStore()
	require.NoError(t, err)
	_, err = rs.MakeRepo("stats_analyze_table_test")
	require.NoError(t, err)

	srvSettings := &driver.Server{
		Args:        []string{},
		DynamicPort: "server",
	}
	ports := newPorts(t)
	server := StartServer(t, rs, "stats_analyze_table_test", srvSettings, ports)

	// Connect to the database
	db, err := server.DB(driver.Connection{})
	require.NoError(t, err)
	t.Cleanup(func() {
		db.Close()
	})

	// Task rate-limiting of one per minute will make things more
	// than slow enough that ANALYZE down below will fail on
	// the context timeout if these rates are affecting the
	// ANALYZE call.
	_, err = db.ExecContext(ctx, "select dolt_stats_timers(60000000000, 60000000000)")
	require.NoError(t, err)

	var columns = []string{
		"c0", "c1", "c2", "c3", "c4",
		"c5", "c6", "c7", "c8", "c9",
	}
	var columndefs = []string{}
	for _, col := range columns {
		columndefs = append(columndefs, fmt.Sprintf("%s bigint unique", col))
	}

	_, err = db.ExecContext(ctx, fmt.Sprintf(strings.Join([]string{
		"CREATE TABLE %s (",
		"id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,",
		strings.Join(columndefs, ",\n"),
		");",
	}, "\n"), "t0"))
	require.NoError(t, err)

	_, err = db.ExecContext(ctx, "select dolt_commit('-Am', 'initial commit')")
	require.NoError(t, err)

	_, err = db.ExecContext(ctx, "select dolt_checkout('main')")
	require.NoError(t, err)
	_, err = db.ExecContext(ctx, fmt.Sprintf("insert into %s values %s", "t0", t0_values_strs(65536)))
	require.NoError(t, err)

	timeoutCtx, cancel := context.WithTimeout(ctx, time.Second*30)
	defer cancel()
	_, err = db.ExecContext(timeoutCtx, "ANALYZE t0")
	require.NoError(t, err)
}

func t0_values_str() string {
	return fmt.Sprintf("(%d, %d, %d, %d, %d, %d, %d, %d, %d, %d, %d)",
		rand.IntN(1<<48), rand.IntN(1<<48), rand.IntN(1<<48), rand.IntN(1<<48),
		rand.IntN(1<<48), rand.IntN(1<<48), rand.IntN(1<<48), rand.IntN(1<<48),
		rand.IntN(1<<48), rand.IntN(1<<48), rand.IntN(1<<48))
}

func t0_values_strs(n int) string {
	vals := make([]string, n)
	for i := range vals {
		vals[i] = t0_values_str()
	}
	return strings.Join(vals, ",")
}