package database

import (
	"context"
	"database/sql"
	"errors"
	"io"
	"net"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

// Exercise the real mysql-tracing driver against a silent TCP peer. Transaction
// finalization has no context in database/sql/driver; query-only cancellation
// tests would miss the stalled COMMIT/ROLLBACK failure class.
func TestPoolNetworkTimeoutDiscardsStalledConnection(t *testing.T) {
	for _, command := range []string{"SELECT 1", "DO 1", "COMMIT", "ROLLBACK", "SET SESSION tidb_txn_mode = DEFAULT"} {
		t.Run(command, func(t *testing.T) {
			address, connections := silentMySQLPeer(t, command)
			host, port, err := net.SplitHostPort(address)
			if err != nil {
				t.Fatal(err)
			}
			t.Setenv("TIDB_HOST", host)
			t.Setenv("TIDB_PORT", port)
			t.Setenv("TIDB_USER", "probe")
			t.Setenv("TIDB_PASSWORD", "")
			t.Setenv("TIDB_DATABASE", "probe")
			t.Setenv("TIDB_TLS_DISABLE", "1")
			ctx, cancel := context.WithTimeout(t.Context(), 4*time.Second)
			defer cancel()
			connection, err := NewConnectionContext(ctx, PoolOptions{
				MaxOpenConns: 1, MaxIdleConns: 1, NetworkTimeout: 150 * time.Millisecond,
			})
			if err != nil {
				t.Fatal(err)
			}
			db := connection.DB()
			t.Cleanup(func() { _ = db.Close() })
			var tx *sql.Tx
			if command == "COMMIT" || command == "ROLLBACK" {
				tx, err = db.BeginTx(ctx, nil)
				if err != nil {
					t.Fatal(err)
				}
			}
			started := time.Now()
			switch command {
			case "SELECT 1":
				var value int
				err = db.QueryRowContext(ctx, command).Scan(&value)
			case "COMMIT":
				err = tx.Commit()
			case "ROLLBACK":
				err = tx.Rollback()
			default:
				_, err = db.ExecContext(ctx, command)
			}
			elapsed := time.Since(started)
			if err == nil || errors.Is(err, context.DeadlineExceeded) || elapsed > time.Second {
				t.Fatalf("silent peer: err=%v elapsed=%s; want driver timeout before context deadline", err, elapsed)
			}
			if err := db.PingContext(ctx); err != nil {
				t.Fatalf("pool did not recover: %v", err)
			}
			if got := connections.Load(); got != 2 {
				t.Fatalf("physical connections = %d; want stalled connection replaced once", got)
			}
			t.Logf("%s stopped after %s; next operation used a healthy connection", command, elapsed)
		})
	}
}

// This peer supplies only the MySQL handshake and OK packets. It makes no
// claims about SQL semantics or transaction durability; those require TiDB.
func silentMySQLPeer(t *testing.T, stallCommand string) (string, *atomic.Int32) {
	t.Helper()
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	var connections atomic.Int32
	var stalled atomic.Bool
	var workers sync.WaitGroup
	done := make(chan struct{})
	workers.Add(1)
	go func() {
		defer workers.Done()
		for {
			conn, err := listener.Accept()
			if err != nil {
				return
			}
			connections.Add(1)
			workers.Add(1)
			go func() {
				defer workers.Done()
				defer conn.Close()
				_ = conn.SetDeadline(time.Now().Add(3 * time.Second))
				capabilities := uint32(1 | 4 | 8 | 512 | 8192 | 32768 | 524288)
				hello := append([]byte{10}, []byte("5.7.0-test\x00")...)
				hello = append(hello, 1, 0, 0, 0)
				hello = append(hello, []byte("12345678\x00")...)
				hello = append(hello, byte(capabilities), byte(capabilities>>8), 33, 2, 0,
					byte(capabilities>>16), byte(capabilities>>24), 21)
				hello = append(hello, make([]byte, 10)...)
				hello = append(hello, []byte("abcdefghijkl\x00mysql_native_password\x00")...)
				if writeMySQLTestPacket(conn, 0, hello) != nil {
					return
				}
				if _, err := readMySQLTestPacket(conn); err != nil {
					return
				}
				ok := []byte{0, 0, 0, 2, 0, 0, 0}
				if writeMySQLTestPacket(conn, 2, ok) != nil {
					return
				}
				for {
					packet, err := readMySQLTestPacket(conn)
					if err != nil || len(packet) == 0 || packet[0] == 1 { // COM_QUIT
						return
					}
					if packet[0] == 3 && strings.EqualFold(strings.TrimSpace(string(packet[1:])), stallCommand) && stalled.CompareAndSwap(false, true) {
						<-done
						return
					}
					if writeMySQLTestPacket(conn, 1, ok) != nil {
						return
					}
				}
			}()
		}
	}()
	t.Cleanup(func() {
		close(done)
		_ = listener.Close()
		workers.Wait()
	})
	return listener.Addr().String(), &connections
}

func readMySQLTestPacket(conn net.Conn) ([]byte, error) {
	header := make([]byte, 4)
	if _, err := io.ReadFull(conn, header); err != nil {
		return nil, err
	}
	body := make([]byte, int(header[0])|int(header[1])<<8|int(header[2])<<16)
	_, err := io.ReadFull(conn, body)
	return body, err
}

func writeMySQLTestPacket(conn net.Conn, sequence byte, body []byte) error {
	_, err := conn.Write(append([]byte{byte(len(body)), byte(len(body) >> 8), byte(len(body) >> 16), sequence}, body...))
	return err
}