// Copyright 2024 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 logrepl_test

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
	"testing"
	"time"

	"github.com/apecloud/myduckserver/adapter"
	"github.com/apecloud/myduckserver/configuration"
	"github.com/jackc/pglogrepl"

	"github.com/apecloud/myduckserver/pgserver"
	"github.com/apecloud/myduckserver/pgserver/logrepl"
	"github.com/apecloud/myduckserver/pgtest"
	"github.com/cockroachdb/errors"
	"github.com/dolthub/go-mysql-server/sql"
	"github.com/jackc/pgx/v5"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// special pseudo-queries for orchestrating replication tests
const (
	createReplicationSlot = "createReplicationSlot"
	dropReplicationSlot   = "dropReplicationSlot"
	stopReplication       = "stopReplication"
	startReplication      = "startReplication"
	waitForCatchup        = "waitForCatchup"
	sleep                 = "sleep"
)

type ScriptTestAssertion = pgtest.ScriptTestAssertion

type ReplicationTest struct {
	// Name of the script.
	Name string
	// The database to create and use. If not provided, then it defaults to "postgres".
	Database string
	// The SQL statements to execute as setup, in order. Results are not checked, but statements must not error.
	// An initial comment can be used to Setup is always run on the primary.
	SetUpScript []string
	// The set of assertions to make after setup, in order
	Assertions []ScriptTestAssertion
	// When using RunScripts, setting this on one (or more) tests causes RunScripts to ignore all tests that have this
	// set to false (which is the default value). This allows a developer to easily "focus" on a specific test without
	// having to comment out other tests, pull it into a different function, etc. In addition, CI ensures that this is
	// false before passing, meaning this prevents the commented-out situation where the developer forgets to uncomment
	// their code.
	Focus bool
	// Skip is used to completely skip a test including setup
	Skip bool
}

var replicationTests = []ReplicationTest{
	{
		Name: "simple replication, strings and integers",
		SetUpScript: []string{
			dropReplicationSlot,
			createReplicationSlot,
			startReplication,
			"/* replica */ drop table if exists public.test",
			"/* replica */ create table public.test (id INT primary key, name varchar(100))",
			"drop table if exists public.test",
			"CREATE TABLE public.test (id INT primary key, name varchar(100))",
			"INSERT INTO public.test VALUES (1, 'one')",
			"INSERT INTO public.test VALUES (2, 'two')",
			"UPDATE public.test SET name = 'three' WHERE id = 2",
			"DELETE FROM public.test WHERE id = 1",
			"INSERT INTO public.test VALUES (3, 'one')",
			"INSERT INTO public.test VALUES (4, 'two')",
			"UPDATE public.test SET name = 'five' WHERE id = 4",
			"DELETE FROM public.test WHERE id = 3",
			"INSERT INTO public.test VALUES (5, 'one')",
			"INSERT INTO public.test VALUES (6, 'two')",
			"UPDATE public.test SET name = 'six' WHERE id = 6",
			"DELETE FROM public.test WHERE id = 5",
			waitForCatchup,
		},
		Assertions: []ScriptTestAssertion{
			{
				Query: "/* replica */ SELECT * FROM public.test order by id",
				Expected: []sql.Row{
					{int32(2), "three"},
					{int32(4), "five"},
					{int32(6), "six"},
				},
			},
		},
	},
	{
		Name: "stale start",
		SetUpScript: []string{
			// Postgres will not start tracking which WAL locations to send until the replication slot is created, so we have
			// to do that first. Customers have the same constraint: they must import any table data that existed before
			// they create the replication slot.
			dropReplicationSlot,
			createReplicationSlot,
			"/* replica */ drop table if exists public.test",
			"/* replica */ create table public.test (id INT primary key, name varchar(100))",
			"drop table if exists public.test",
			"CREATE TABLE public.test (id INT primary key, name varchar(100))",
			"INSERT INTO public.test VALUES (1, 'one')",
			"INSERT INTO public.test VALUES (2, 'two')",
			"UPDATE public.test SET name = 'three' WHERE id = 2",
			"DELETE FROM public.test WHERE id = 1",
			"INSERT INTO public.test VALUES (3, 'one')",
			"INSERT INTO public.test VALUES (4, 'two')",
			"UPDATE public.test SET name = 'five' WHERE id = 4",
			"DELETE FROM public.test WHERE id = 3",
			"INSERT INTO public.test VALUES (5, 'one')",
			"INSERT INTO public.test VALUES (6, 'two')",
			"UPDATE public.test SET name = 'six' WHERE id = 6",
			"DELETE FROM public.test WHERE id = 5",
			startReplication,
			waitForCatchup,
		},
		Assertions: []ScriptTestAssertion{
			{
				Query: "/* replica */ SELECT * FROM public.test order by id",
				Expected: []sql.Row{
					{int32(2), "three"},
					{int32(4), "five"},
					{int32(6), "six"},
				},
			},
		},
	},
	{
		Name: "stopping and resuming replication",
		SetUpScript: []string{
			dropReplicationSlot,
			createReplicationSlot,
			startReplication,
			"/* replica */ drop table if exists public.test",
			"/* replica */ create table public.test (id INT primary key, name varchar(100))",
			"drop table if exists public.test",
			"CREATE TABLE public.test (id INT primary key, name varchar(100))",
			"INSERT INTO public.test VALUES (1, 'one')",
			"INSERT INTO public.test VALUES (2, 'two')",
			waitForCatchup,
			stopReplication,
			"UPDATE public.test SET name = 'three' WHERE id = 2",
			"DELETE FROM public.test WHERE id = 1",
			"INSERT INTO public.test VALUES (3, 'one')",
			"INSERT INTO public.test VALUES (4, 'two')",
			"UPDATE public.test SET name = 'five' WHERE id = 4",
			"DELETE FROM public.test WHERE id = 3",
			startReplication,
			"INSERT INTO public.test VALUES (5, 'one')",
			"INSERT INTO public.test VALUES (6, 'two')",
			"UPDATE public.test SET name = 'six' WHERE id = 6",
			"DELETE FROM public.test WHERE id = 5",
			waitForCatchup,
		},
		Assertions: []ScriptTestAssertion{
			{
				Query: "/* replica */ SELECT * FROM public.test order by id",
				Expected: []sql.Row{
					{int32(2), "three"},
					{int32(4), "five"},
					{int32(6), "six"},
				},
			},
		},
	},
	{
		Name: "extended stop/start",
		SetUpScript: []string{
			dropReplicationSlot,
			createReplicationSlot,
			"/* replica */ drop table if exists public.test",
			"/* replica */ create table public.test (id INT primary key, name varchar(100))",
			"drop table if exists public.test",
			"CREATE TABLE public.test (id INT primary key, name varchar(100))",
			"INSERT INTO public.test VALUES (1, 'one')",
			"INSERT INTO public.test VALUES (2, 'two')",
			"UPDATE public.test SET name = 'three' WHERE id = 2",
			"DELETE FROM public.test WHERE id = 1",
			"INSERT INTO public.test VALUES (3, 'one')",
			"INSERT INTO public.test VALUES (4, 'two')",
			"UPDATE public.test SET name = 'five' WHERE id = 4",
			"DELETE FROM public.test WHERE id = 3",
			"INSERT INTO public.test VALUES (5, 'one')",
			startReplication,
			"INSERT INTO public.test VALUES (6, 'two')",
			"UPDATE public.test SET name = 'six' WHERE id = 6",
			stopReplication,
			"DELETE FROM public.test WHERE id = 5",
			"INSERT INTO public.test VALUES (7, 'one')",
			"INSERT INTO public.test VALUES (8, 'two')",
			startReplication,
			"UPDATE public.test SET name = 'nine' WHERE id = 8",
			"DELETE FROM public.test WHERE id = 7",
			"INSERT INTO public.test VALUES (9, 'one')",
			stopReplication,
			startReplication,
			"INSERT INTO public.test VALUES (10, 'two')",
			"UPDATE public.test SET name = 'eleven' WHERE id = 10",
			stopReplication,
			"DELETE FROM public.test WHERE id = 9",
			"INSERT INTO public.test VALUES (11, 'one')",
			"INSERT INTO public.test VALUES (12, 'two')",
			"UPDATE public.test SET name = 'thirteen' WHERE id = 12",
			"DELETE FROM public.test WHERE id = 11",
			startReplication,
			"INSERT INTO public.test VALUES (13, 'one')",
			"INSERT INTO public.test VALUES (14, 'two')",
			"UPDATE public.test SET name = 'fifteen' WHERE id = 14",
			"DELETE FROM public.test WHERE id = 13",
			waitForCatchup,
			// since replication lag is a heuristic, this sleep is necessary to ensure that the replica has caught
			// up in all cases before we shut it off
			sleep,
			stopReplication,
			// below this point we don't expect to find any values replicated because replication was stopped
			"INSERT INTO public.test VALUES (15, 'one')",
			"INSERT INTO public.test VALUES (16, 'two')",
			"UPDATE public.test SET name = 'seventeen' WHERE id = 16",
			"DELETE FROM public.test WHERE id = 15",
			sleep, // final sleep to make sure that any replication events that will arrive have
		},
		Assertions: []ScriptTestAssertion{
			{
				Query: "/* replica */ SELECT * FROM public.test order by id",
				Expected: []sql.Row{
					{int32(2), "three"},
					{int32(4), "five"},
					{int32(6), "six"},
					{int32(8), "nine"},
					{int32(10), "eleven"},
					{int32(12), "thirteen"},
					{int32(14), "fifteen"},
				},
			},
		},
	},
	{
		Name: "all supported types",
		SetUpScript: []string{
			dropReplicationSlot,
			createReplicationSlot,
			startReplication,
			"/* replica */ drop table if exists public.test",
			"/* replica */ create table public.test (id INT primary key, name varchar(100), u_id uuid, age INT, height FLOAT)",
			"drop table if exists public.test",
			"create table public.test (id INT primary key, name varchar(100), u_id uuid, age INT, height FLOAT)",
			"INSERT INTO public.test VALUES (1, 'one', '5ef34887-e635-4c9c-a994-97b1cb810786', 1, 1.1)",
			"INSERT INTO public.test VALUES (2, 'two', '2de55648-76ec-4f66-9fae-bd3d853fb0da', 2, 2.2)",
			"UPDATE public.test SET name = 'three' WHERE id = 2",
			"update public.test set u_id = '3232abe7-560b-4714-a020-2b1a11a1ec65' where id = 2",
			"DELETE FROM public.test WHERE id = 1",
			waitForCatchup,
		},
		Assertions: []ScriptTestAssertion{
			{
				Query: "/* replica */ SELECT * FROM public.test order by id",
				Expected: []sql.Row{
					{int32(2), "three", "3232abe7-560b-4714-a020-2b1a11a1ec65", int32(2), float32(2.2)},
				},
			},
		},
	},
	{
		Name: "concurrent writes",
		SetUpScript: []string{
			dropReplicationSlot,
			createReplicationSlot,
			startReplication,
			"/* replica */ drop table if exists public.test",
			"/* replica */ create table public.test (id INT primary key, name varchar(100))",
			"drop table if exists public.test",
			"CREATE TABLE public.test (id INT primary key, name varchar(100))",
			"/* primary a */ START TRANSACTION",
			"/* primary a */ INSERT INTO public.test VALUES (1, 'one')",
			"/* primary a */ INSERT INTO public.test VALUES (2, 'two')",
			"/* primary b */ START TRANSACTION",
			"/* primary b */ INSERT INTO public.test VALUES (3, 'one')",
			"/* primary b */ INSERT INTO public.test VALUES (4, 'two')",
			"/* primary a */ UPDATE public.test SET name = 'three' WHERE id > 0",
			"/* primary a */ DELETE FROM public.test WHERE id = 1",
			"/* primary b */ UPDATE public.test SET name = 'five' WHERE id > 0",
			"/* primary b */ DELETE FROM public.test WHERE id = 3",
			"/* primary b */ COMMIT",
			"/* primary a */ COMMIT",
			waitForCatchup,
		},
		Assertions: []ScriptTestAssertion{
			{
				Query: "/* replica */ SELECT * FROM public.test order by id",
				Expected: []sql.Row{
					{int32(2), "three"},
					{int32(4), "five"},
				},
			},
		},
	},
	{
		Name: "concurrent writes with restarts",
		SetUpScript: []string{
			dropReplicationSlot,
			createReplicationSlot,
			startReplication,
			"/* replica */ drop table if exists public.test",
			"/* replica */ create table public.test (id INT primary key, name varchar(100))",
			"drop table if exists public.test",
			"CREATE TABLE public.test (id INT primary key, name varchar(100))",
			"/* primary a */ START TRANSACTION",
			"/* primary a */ INSERT INTO public.test VALUES (1, 'one')",
			"/* primary a */ INSERT INTO public.test VALUES (2, 'two')",
			stopReplication,
			"/* primary b */ START TRANSACTION",
			"/* primary b */ INSERT INTO public.test VALUES (3, 'one')",
			"/* primary b */ INSERT INTO public.test VALUES (4, 'two')",
			"/* primary c */ START TRANSACTION",
			"/* primary c */ INSERT INTO public.test VALUES (5, 'one')",
			"/* primary c */ INSERT INTO public.test VALUES (6, 'two')",
			"/* primary a */ UPDATE public.test SET name = 'three' WHERE id > 0",
			startReplication,
			"/* primary a */ DELETE FROM public.test WHERE id = 1",
			"/* primary b */ UPDATE public.test SET name = 'five' WHERE id > 0",
			"/* primary b */ DELETE FROM public.test WHERE id = 3",
			"/* primary b */ COMMIT",
			stopReplication,
			"/* primary c */ UPDATE public.test SET name = 'seven' WHERE id > 0",
			"/* primary c */ DELETE FROM public.test WHERE id = 5",
			"/* primary a */ COMMIT",
			startReplication,
			"/* primary c */ COMMIT",
			waitForCatchup,
		},
		Assertions: []ScriptTestAssertion{
			{
				Query: "/* replica */ SELECT * FROM public.test order by id",
				Expected: []sql.Row{
					{int32(2), "three"},
					{int32(4), "seven"},
					{int32(6), "seven"},
				},
			},
		},
	},
	{
		Name: "concurrent writes with rollbacks",
		SetUpScript: []string{
			dropReplicationSlot,
			createReplicationSlot,
			startReplication,
			"/* replica */ drop table if exists public.test",
			"/* replica */ create table public.test (id INT primary key, name varchar(100))",
			"drop table if exists public.test",
			"CREATE TABLE public.test (id INT primary key, name varchar(100))",
			"/* primary a */ START TRANSACTION",
			"/* primary a */ INSERT INTO public.test VALUES (1, 'one')",
			"/* primary a */ INSERT INTO public.test VALUES (2, 'two')",
			stopReplication,
			"/* primary b */ START TRANSACTION",
			"/* primary b */ INSERT INTO public.test VALUES (3, 'one')",
			"/* primary b */ INSERT INTO public.test VALUES (4, 'two')",
			"/* primary c */ START TRANSACTION",
			"/* primary c */ INSERT INTO public.test VALUES (5, 'one')",
			"/* primary c */ INSERT INTO public.test VALUES (6, 'two')",
			"/* primary a */ UPDATE public.test SET name = 'three' WHERE id > 0",
			startReplication,
			"/* primary a */ DELETE FROM public.test WHERE id = 1",
			"/* primary b */ UPDATE public.test SET name = 'five' WHERE id > 0",
			"/* primary b */ DELETE FROM public.test WHERE id = 3",
			"/* primary b */ COMMIT",
			stopReplication,
			"/* primary c */ UPDATE public.test SET name = 'seven' WHERE id > 0",
			"/* primary c */ DELETE FROM public.test WHERE id = 5",
			"/* primary a */ ROLLBACK",
			startReplication,
			"/* primary c */ COMMIT",
			waitForCatchup,
		},
		Assertions: []ScriptTestAssertion{
			{
				Query: "/* replica */ SELECT * FROM public.test order by id",
				Expected: []sql.Row{
					{int32(4), "seven"},
					{int32(6), "seven"},
				},
			},
		},
	},
	{
		Name: "concurrent writes, stale commits",
		SetUpScript: []string{
			dropReplicationSlot,
			createReplicationSlot,
			startReplication,
			"/* replica */ drop table if exists public.test",
			"/* replica */ create table public.test (id INT primary key, name varchar(100))",
			"drop table if exists public.test",
			"CREATE TABLE public.test (id INT primary key, name varchar(100))",
			"/* primary a */ START TRANSACTION",
			"/* primary a */ INSERT INTO public.test VALUES (1, 'one')",
			"/* primary b */ START TRANSACTION",
			"/* primary b */ INSERT INTO public.test VALUES (2, 'two')",
			"/* primary b */ COMMIT",
			waitForCatchup,
			stopReplication,
			startReplication,
			// this tx includes several WAL locations before our last flush, but it must still be replicated
			"/* primary a */ COMMIT",
			waitForCatchup,
		},
		Assertions: []ScriptTestAssertion{
			{
				Query: "/* replica */ SELECT * FROM public.test order by id",
				Expected: []sql.Row{
					{int32(1), "one"},
					{int32(2), "two"},
				},
			},
		},
	},
	{
		Name: "concurrent writes, very stale commits",
		SetUpScript: []string{
			dropReplicationSlot,
			createReplicationSlot,
			startReplication,
			"/* replica */ drop table if exists public.test",
			"/* replica */ create table public.test (id INT primary key, name varchar(100))",
			"drop table if exists public.test",
			"CREATE TABLE public.test (id INT primary key, name varchar(100))",
			"/* primary a */ START TRANSACTION",
			"/* primary a */ INSERT INTO public.test VALUES (1, 'one')",
			"/* primary a */ INSERT INTO public.test VALUES (2, 'two')",
			"/* primary a */ UPDATE public.test SET name = 'three' WHERE id > 0",
			"/* primary a */ DELETE FROM public.test WHERE id = 1",
			"/* primary b */ START TRANSACTION",
			"/* primary b */ INSERT INTO public.test VALUES (3, 'one')",
			"/* primary b */ INSERT INTO public.test VALUES (4, 'two')",
			"/* primary c */ START TRANSACTION",
			"/* primary c */ INSERT INTO public.test VALUES (5, 'one')",
			"/* primary c */ INSERT INTO public.test VALUES (6, 'two')",
			"/* primary c */ UPDATE public.test SET name = 'seven' WHERE id > 0",
			"/* primary c */ DELETE FROM public.test WHERE id = 5",
			"/* primary c */ COMMIT",
			"/* primary b */ UPDATE public.test SET name = 'five' WHERE id > 0",
			"/* primary b */ DELETE FROM public.test WHERE id = 3",
			"/* primary b */ COMMIT",
			waitForCatchup,
			stopReplication,
			startReplication,
			// this tx includes several WAL locations before our last flush, but it must still be replicated
			"/* primary a */ COMMIT",
			waitForCatchup,
		},
		Assertions: []ScriptTestAssertion{
			{
				Query: "/* replica */ SELECT * FROM public.test order by id",
				Expected: []sql.Row{
					{int32(2), "three"},
					{int32(4), "five"},
					{int32(6), "five"},
				},
			},
		},
	},
	{
		Name: "all types",
		SetUpScript: []string{
			dropReplicationSlot,
			createReplicationSlot,
			startReplication,
			"/* replica */ drop table if exists public.test",
			"/* replica */ create table public.test (id INT primary key, " +
				"name varchar(100), " +
				"age SMALLINT, " +
				"is_cool BOOLEAN, " +
				"height FLOAT, " +
				"birth_date DATE, " +
				"birth_time TIME, " +
				// TIMETZ is discouraged:
				// https://www.postgresql.org/docs/current/datatype-datetime.html
				// https://github.com/jackc/pgx/issues/1940
				// "birth_timetz TIME WITH TIME ZONE, " +
				"birth_timestamp TIMESTAMP, " +
				"birth_timestamptz TIMESTAMP WITH TIME ZONE, " +
				"income NUMERIC(10,2), " +
				"binary_data BYTEA, " +
				"description TEXT, " +
				"code CHAR(3), " +
				"tags TEXT[], " +
				"scores INTEGER[], " +
				"real_nums REAL[], " +
				"decimal_nums DECIMAL(8,3)[], " +
				"small_num SMALLINT, " +
				"big_num BIGINT, " +
				"json_data JSON)",
			"drop table if exists public.test",
			"create table public.test (id INT primary key, " +
				"name varchar(100), " +
				"age SMALLINT, " +
				"is_cool BOOLEAN, " +
				"height FLOAT, " +
				"birth_date DATE, " +
				"birth_time TIME, " +
				// "birth_timetz TIME WITH TIME ZONE, " +
				"birth_timestamp TIMESTAMP, " +
				"birth_timestamptz TIMESTAMP WITH TIME ZONE, " +
				"income DECIMAL(10,2), " +
				"binary_data BYTEA, " +
				"description TEXT, " +
				"code CHAR(3), " +
				"tags TEXT[], " +
				"scores INTEGER[], " +
				"real_nums REAL[], " +
				"decimal_nums DECIMAL(8,3)[], " +
				"small_num SMALLINT, " +
				"big_num BIGINT, " +
				"json_data JSONB)",
			"INSERT INTO public.test VALUES (1, " +
				"'one', 1, true, 1.1, " +
				"'2021-01-01', '12:00:00', '2021-01-01 12:00:00', '2021-01-01 20:00:00+8', " +
				"12345678.9, " +
				`'\x0123456789ABCDEF', 'long text description', 'ABC', ` +
				"ARRAY['tag1', 'tag2'], ARRAY[1, 2, 3], ARRAY[1.1, 2.2, 3.3]::real[], ARRAY[1.1, 22.22, 333.333], " +
				`123, 9223372036854775807, '{"key": "value"}')`,
			"INSERT INTO public.test VALUES (2, " +
				"'two', 2, false, 2.2, " +
				"'2021-02-02', '13:00:00.123456', '2021-02-02 13:00:00.123456', '2021-02-02 05:00:00.123456-8', " +
				"98765432.1, " +
				`'\xDEADBEEF', 'another description', 'XYZ', ` +
				"ARRAY['tag3', 'tag4', NULL, 'NULL'], ARRAY[4, 5, 6, 7, NULL], ARRAY[NULL, 4.4, 5.5, 6.6, 7.7]::real[], ARRAY[4.4, 55.55, NULL, 66.66, 777.777], " +
				`-123, -9223372036854775808, '{"array": [1, 2, 3]}')`,
			"UPDATE public.test SET name = 'three' WHERE id = 2",
			`UPDATE public.test SET json_data = jsonb_set(json_data, '{key}', '"new_value"') WHERE id = 1`,
			waitForCatchup,
		},
		Assertions: []ScriptTestAssertion{
			{
				Query: "/* replica */ SELECT * EXCLUDE (birth_timestamptz), birth_timestamptz AT TIME ZONE 'UTC' FROM public.test order by id",
				Expected: []sql.Row{
					{int32(1), "one", int16(1), true, float32(1.1),
						"2021-01-01", "12:00:00.000000",
						// "12:00:00+00",
						"2021-01-01 12:00:00",
						pgtest.Numeric("12345678.9"),
						[]byte{0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF},
						"long text description", "ABC",
						`{tag1,tag2}`,
						`{1,2,3}`,
						`{1.1,2.2,3.3}`,
						`{1.1,22.22,333.333}`,
						int16(123),
						int64(9223372036854775807),
						`{"key": "new_value"}`,
						"2021-01-01 12:00:00",
					},
					{int32(2), "three", int16(2), false, float32(2.2),
						"2021-02-02", "13:00:00.123456",
						// "13:00:00.123456+00",
						"2021-02-02 13:00:00.123456",
						pgtest.Numeric("98765432.1"),
						[]byte{0xDE, 0xAD, 0xBE, 0xEF},
						"another description", "XYZ",
						`{tag3,tag4,NULL,"NULL"}`,
						`{4,5,6,7,NULL}`,
						`{NULL,4.4,5.5,6.6,7.7}`,
						`{4.4,55.55,NULL,66.66,777.777}`,
						int16(-123),
						int64(-9223372036854775808),
						`{"array": [1, 2, 3]}`,
						"2021-02-02 13:00:00.123456",
					},
				},
			},
		},
	},
	{
		Name: "Create table automatically",
		SetUpScript: []string{
			dropReplicationSlot,
			createReplicationSlot,
			startReplication,
			"/* replica */ drop table if exists public.test",
			"drop table if exists public.test",
			"CREATE TABLE public.test (id INT primary key, name varchar(10))",
			"INSERT INTO public.test VALUES (1, 'one')",
			"INSERT INTO public.test VALUES (2, 'two'), (3, 'six')",
			"UPDATE public.test SET name = 'three' WHERE id = 3",
			"DELETE FROM public.test WHERE id = 1 or id = 2",
			waitForCatchup,
		},
		Assertions: []ScriptTestAssertion{
			{
				Query: "/* replica */ SELECT * FROM public.test order by id",
				Expected: []sql.Row{
					{int32(3), "three"},
				},
			},
		},
	},
	{
		Name: "Truncate table",
		SetUpScript: []string{
			dropReplicationSlot,
			createReplicationSlot,
			startReplication,
			"/* replica */ drop table if exists public.test",
			"drop table if exists public.test",
			"CREATE TABLE public.test (id INT primary key, name varchar(10))",
			"INSERT INTO public.test VALUES (1, 'one'), (2, 'two'), (3, 'three')",
			"TRUNCATE TABLE public.test",
			"INSERT INTO public.test VALUES (4, 'four')",
			waitForCatchup,
		},
		Assertions: []ScriptTestAssertion{
			{
				Query: "/* replica */ SELECT * FROM public.test order by id",
				Expected: []sql.Row{
					{int32(4), "four"},
				},
			},
		},
	},
}

func TestReplication(t *testing.T) {
	// logrus.SetLevel(logrus.DebugLevel)
	RunReplicationScripts(t, replicationTests)
}

// RunScripts runs the given collection of scripts.
func RunReplicationScripts(t *testing.T, scripts []ReplicationTest) {
	// First, we'll run through the scripts to check for the Focus variable. If it's true, then append it to the new slice.
	focusScripts := make([]ReplicationTest, 0, len(scripts))
	for _, script := range scripts {
		if script.Focus {
			// If this is running in GitHub Actions, then we'll panic, because someone forgot to disable it before committing
			if _, ok := os.LookupEnv("GITHUB_ACTION"); ok {
				panic(fmt.Sprintf("The script `%s` has Focus set to `true`. GitHub Actions requires that "+
					"all tests are run, which Focus circumvents, leading to this error. Please disable Focus on "+
					"all tests.", script.Name))
			}
			focusScripts = append(focusScripts, script)
		}
	}
	// If we have scripts with Focus set, then we replace the normal script slice with the new slice.
	if len(focusScripts) > 0 {
		scripts = focusScripts
	}

	// start the docker container
	containerName, dsn, _, err := StartPostgresServer()
	require.NoError(t, err)
	defer func() {
		err := exec.Command("docker", "kill", containerName).Run()
		require.NoError(t, err)
	}()

	primaryDns := dsn + "?sslmode=disable"

	// We drop and recreate the replication slot once at the beginning of the test suite. Postgres seems to do a little
	// work in the background with a publication, so we need to wait a little bit before running any test scripts.
	require.NoError(t, logrepl.DropPublication(primaryDns, slotName))
	require.NoError(t, logrepl.CreatePublicationIfNotExists(primaryDns, slotName))
	time.Sleep(500 * time.Millisecond)

	// for i, script := range scripts {
	// 	if i == 10 {
	// 		RunReplicationScript(t, dsn, script)
	// 	}
	// }
	for _, script := range scripts {
		RunReplicationScript(t, dsn, script)
	}
}

const slotName = "myduck_slot"
const subscriptionName = "my_sub_test"

// RunReplicationScript runs the given ReplicationTest.
func RunReplicationScript(t *testing.T, dsn string, script ReplicationTest) {
	scriptDatabase := script.Database
	if len(scriptDatabase) == 0 {
		scriptDatabase = "postgres"
	}

	// primaryDns is the connection to the actual postgres database.
	// If you have postgres running on a different port, you'll need to change this.
	primaryDns := dsn + "?sslmode=disable"

	ctx, pgServer, replicaConn, close, err := pgtest.CreateTestServer(t, findFreePort())
	require.NoError(t, err)
	defer func() {
		replicaConn.Close(ctx)
		err := close()
		require.NoError(t, err)
	}()

	ctx = context.Background()
	t.Run(script.Name, func(t *testing.T) {
		runReplicationScript(ctx, t, script, pgServer, replicaConn, primaryDns)
	})
}

func newReplicator(sqlCtx *sql.Context, t *testing.T, primaryDns string) *logrepl.LogicalReplicator {
	err := logrepl.CreateSubscription(sqlCtx, subscriptionName, primaryDns, slotName, pglogrepl.LSN(0).String(), true)
	require.NoError(t, err)

	tx := adapter.TryGetTxn(sqlCtx)
	if tx != nil {
		err := tx.Commit()
		require.NoError(t, err)
		adapter.CloseTxn(sqlCtx)
	}

	r, err := logrepl.NewLogicalReplicator(subscriptionName, primaryDns)
	require.NoError(t, err)
	return r
}

// runReplicationScript runs the script given on the postgres connection provided
func runReplicationScript(
	ctx context.Context,
	t *testing.T,
	script ReplicationTest,
	server *pgserver.Server,
	replicaConn *pgx.Conn,
	primaryDns string,
) {
	r := newReplicator(server.NewInternalCtx(), t, primaryDns)
	defer r.Stop()

	if script.Skip {
		t.Skip("Skip has been set in the script")
	}

	connections := map[string]*pgx.Conn{
		"replica": replicaConn,
	}

	defer func() {
		for _, conn := range connections {
			if conn != nil {
				conn.Close(ctx)
			}
		}
	}()

	// Run the setup
	for _, query := range script.SetUpScript {
		// handle logic for special pseudo-queries
		if handlePseudoQuery(t, server, query, r) {
			continue
		}

		conn := connectionForQuery(t, query, connections, primaryDns)

		target, _ := clientSpecFromQueryComment(query)
		if target == "replica" && configuration.IsReplicationWithoutIndex() {
			// Remove the primary key index from the replica table
			query = strings.Replace(query, " primary key,", ",", 1)
		}

		log.Println("Running setup query:", query)
		_, err := conn.Exec(ctx, query)
		require.NoError(t, err)
	}

	// Run the assertions
	for _, assertion := range script.Assertions {
		t.Run(assertion.Query, func(t *testing.T) {
			if assertion.Skip {
				t.Skip("Skip has been set in the assertion")
			}

			// handle logic for special pseudo-queries
			if handlePseudoQuery(t, server, assertion.Query, r) {
				return
			}

			target, _ := clientSpecFromQueryComment(assertion.Query)
			enableRetries := target == "replica"
			conn := connectionForQuery(t, assertion.Query, connections, primaryDns)

			numRetries := 3
			for retries := 0; retries < numRetries; retries++ {
				// If we're skipping the results check, then we call Execute, as it uses a simplified message model.
				if assertion.SkipResultsCheck || assertion.ExpectedErr != "" {
					_, err := conn.Exec(ctx, assertion.Query, assertion.BindVars...)
					if assertion.ExpectedErr != "" {
						require.Error(t, err)
						assert.Contains(t, err.Error(), assertion.ExpectedErr)
					} else {
						require.NoError(t, err)
					}
				} else {
					rows, err := conn.Query(ctx, assertion.Query, assertion.BindVars...)
					require.NoError(t, err)
					readRows, err := pgtest.ReadRows(rows, true)
					require.NoError(t, err)
					normalizedRows := pgtest.NormalizeExpectedRow(rows.FieldDescriptions(), assertion.Expected)

					// For queries against the replica, whether or not replication is caught up is a heuristic that can be
					// incorrect. So we retry queries with sleeps in between to give replication a chance to catch up when this
					// happens.
					if !assert.ObjectsAreEqual(normalizedRows, readRows) {
						if enableRetries && retries < numRetries-1 {
							log.Println("Assertion failed, retrying")
							time.Sleep(500 * time.Millisecond)
							continue
						} else {
							assert.Equal(t, normalizedRows, readRows)
							break
						}
					}
				}
			}
		})
	}
}

// connectionForQuery returns the connection to use for the given query
func connectionForQuery(t *testing.T, query string, connections map[string]*pgx.Conn, primaryDns string) *pgx.Conn {
	target, client := clientSpecFromQueryComment(query)
	var conn *pgx.Conn
	switch target {
	case "primary":
		conn = connections[client]
		if conn == nil {
			var err error
			conn, err = pgx.Connect(context.Background(), primaryDns)
			require.NoError(t, err)
			connections[client] = conn
		}
	case "replica":
		conn = connections["replica"]
	default:
		require.Fail(t, "Invalid target in setup script: ", target)
	}
	return conn
}

// handlePseudoQuery handles special pseudo-queries that are used to orchestrate replication tests and returns whether
// one was handled.
func handlePseudoQuery(t *testing.T, server *pgserver.Server, query string, r *logrepl.LogicalReplicator) bool {
	switch query {
	case createReplicationSlot:
		require.NoError(t, r.CreateReplicationSlotIfNotExists(slotName))
		return true
	case dropReplicationSlot:
		require.NoError(t, r.DropReplicationSlotIfExists(slotName))
		return true
	case startReplication:
		go func() {
			require.NoError(t, r.StartReplication(server.NewInternalCtx(), slotName))
		}()
		require.NoError(t, waitForRunning(r))
		return true
	case stopReplication:
		r.Stop()
		return true
	case waitForCatchup:
		require.NoError(t, waitForCaughtUp(r))
		return true
	case sleep:
		time.Sleep(200 * time.Millisecond)
		return true
	}
	return false
}

// clientSpecFromQueryComment returns "replica" if the query is meant to be run on the replica, and "primary" if it's meant
// to be run on the primary, based on the comment in the query. If not comment, the query runs on the primary
func clientSpecFromQueryComment(query string) (string, string) {
	startCommentIdx := strings.Index(query, "/*")
	endCommentIdx := strings.Index(query, "*/")
	if startCommentIdx < 0 || endCommentIdx < 0 {
		return "primary", "a"
	}

	query = query[startCommentIdx+2 : endCommentIdx]
	if strings.Contains(query, "replica") {
		return "replica", "a"
	}

	if i := strings.Index(query, "primary "); i > 0 && i+len("primary ") < len(query) {
		return "primary", query[i+len("primary "):]
	}

	return "primary", "a"
}

func waitForRunning(r *logrepl.LogicalReplicator) error {
	start := time.Now()
	for {
		if r.Running() {
			break
		}

		if time.Since(start) > time.Second {
			return errors.New("Replication did not start")
		}
		time.Sleep(10 * time.Millisecond)
	}

	return nil
}

func waitForCaughtUp(r *logrepl.LogicalReplicator) error {
	log.Println("Waiting for replication to catch up")

	start := time.Now()
	for {
		if caughtUp, err := r.CaughtUp(150); caughtUp {
			log.Println("replication caught up")
			break
		} else if err != nil {
			return err
		}

		log.Println("replication not caught up, waiting")
		if time.Since(start) >= 30*time.Second {
			return errors.New("Replication did not catch up")
		}
		time.Sleep(1 * time.Second)
	}

	return nil
}
