package database_test

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
	"io"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"git.yuntuo.hk/k2safe/go-omni/database"
)

type fixtureConnector struct {
	identity                                                         string
	connects, closedConnections, closedConnector, commits, rollbacks atomic.Int32
	connect                                                          func(context.Context) error
	ping                                                             func(context.Context) error
	query                                                            func(context.Context) error
	closeErr                                                         error
}

func (c *fixtureConnector) Connect(ctx context.Context) (driver.Conn, error) {
	c.connects.Add(1)
	if c.connect != nil {
		if err := c.connect(ctx); err != nil {
			return nil, err
		}
	}
	return &fixtureConn{owner: c}, nil
}
func (c *fixtureConnector) Driver() driver.Driver { return fixtureDriver{} }
func (c *fixtureConnector) Close() error          { c.closedConnector.Add(1); return c.closeErr }

type fixtureDriver struct{}

func (fixtureDriver) Open(string) (driver.Conn, error) {
	return nil, errors.New("unexpected driver name/DSN path")
}

type fixtureConn struct{ owner *fixtureConnector }

func (c *fixtureConn) Prepare(string) (driver.Stmt, error) { return nil, errors.New("not used") }
func (c *fixtureConn) Close() error                        { c.owner.closedConnections.Add(1); return nil }
func (c *fixtureConn) Begin() (driver.Tx, error)           { return fixtureTx{c.owner}, nil }
func (c *fixtureConn) Ping(ctx context.Context) error {
	if c.owner.ping != nil {
		return c.owner.ping(ctx)
	}
	return nil
}
func (c *fixtureConn) QueryContext(ctx context.Context, _ string, _ []driver.NamedValue) (driver.Rows, error) {
	if c.owner.query != nil {
		if err := c.owner.query(ctx); err != nil {
			return nil, err
		}
	}
	return &fixtureRows{identity: c.owner.identity}, nil
}

type fixtureTx struct{ owner *fixtureConnector }

func (tx fixtureTx) Commit() error   { tx.owner.commits.Add(1); return nil }
func (tx fixtureTx) Rollback() error { tx.owner.rollbacks.Add(1); return nil }

type fixtureRows struct {
	identity string
	sent     bool
}

func (r *fixtureRows) Columns() []string { return []string{"identity"} }
func (r *fixtureRows) Close() error      { return nil }
func (r *fixtureRows) Next(values []driver.Value) error {
	if r.sent {
		return io.EOF
	}
	r.sent = true
	values[0] = r.identity
	return nil
}

func TestConnectorIndependentPoolsAndContext(t *testing.T) {
	type key struct{}
	ctx := context.WithValue(context.Background(), key{}, "startup")
	cfg := database.DefaultPoolConfig()
	cfg.MaxOpenConns = 2
	cfg.MaxIdleConns = 2
	var owners []*fixtureConnector
	var pools []*database.DB
	for _, identity := range []string{"alpha", "beta"} {
		owner := &fixtureConnector{identity: identity}
		// Startup Connect/Ping and each query preserve their caller context.
		// database/sql may also open replacement connections in the background.
		owner.connect = func(ctx context.Context) error {
			if owner.connects.Load() == 1 && ctx.Value(key{}) != "startup" {
				return errors.New("lost request context")
			}
			return nil
		}
		owner.ping = func(ctx context.Context) error {
			if ctx.Value(key{}) != "startup" {
				return errors.New("lost ping context")
			}
			return nil
		}
		owner.query = func(ctx context.Context) error {
			if ctx.Value(key{}) != "startup" {
				return errors.New("lost query context")
			}
			return nil
		}
		pool, err := database.OpenConnector(ctx, owner, cfg)
		if err != nil {
			t.Fatal(err)
		}
		t.Cleanup(func() { pool.Close() })
		if pool.PoolConfig() != cfg || pool.Config().Driver != "" || pool.Config().DSN != "" || pool.Stats().MaxOpenConnections != 2 {
			t.Fatal("pool settings or DSN boundary")
		}
		owners = append(owners, owner)
		pools = append(pools, pool)
	}
	var wg sync.WaitGroup
	for i, pool := range pools {
		for n := 0; n < 16; n++ {
			wg.Add(1)
			go func(i int, pool *database.DB) {
				defer wg.Done()
				var identity string
				if err := pool.QueryRowContext(ctx, "identity").Scan(&identity); err != nil {
					t.Error(err)
				} else if identity != owners[i].identity {
					t.Error("crossed connectors")
				}
			}(i, pool)
		}
	}
	wg.Wait()
	for i, pool := range pools {
		if pool.Stats().OpenConnections > cfg.MaxOpenConns {
			t.Fatal("pool cap not applied")
		}
		if err := pool.Close(); err != nil {
			t.Fatal(err)
		}
		if err := pool.Close(); err != nil {
			t.Fatal(err)
		}
		if owners[i].closedConnector.Load() != 1 || owners[i].closedConnections.Load() != owners[i].connects.Load() {
			t.Fatal("close did not release owned resources exactly once")
		}
		if pool.Health(ctx) == nil {
			t.Fatal("closed pool is healthy")
		}
	}
}

func TestConnectorRejectsInvalidInputWithoutTakingOwnership(t *testing.T) {
	var typedNil *fixtureConnector
	for _, connector := range []driver.Connector{nil, typedNil} {
		if db, err := database.OpenConnector(nil, connector, database.DefaultPoolConfig()); err == nil || db != nil {
			t.Fatal("nil connector accepted")
		}
	}
	for _, mutate := range []func(*database.PoolConfig){
		func(c *database.PoolConfig) { c.MaxOpenConns = -1 }, func(c *database.PoolConfig) { c.MaxIdleConns = -1 },
		func(c *database.PoolConfig) { c.MaxIdleConns = c.MaxOpenConns + 1 }, func(c *database.PoolConfig) { c.PingTimeout = 0 },
		func(c *database.PoolConfig) { c.ConnMaxLifetime = -1 }, func(c *database.PoolConfig) { c.ConnMaxIdleTime = -1 },
	} {
		cfg := database.DefaultPoolConfig()
		mutate(&cfg)
		owner := &fixtureConnector{}
		if db, err := database.OpenConnector(nil, owner, cfg); err == nil || db != nil {
			t.Fatal("invalid config accepted")
		}
		if owner.connects.Load() != 0 || owner.closedConnector.Load() != 0 {
			t.Fatal("validation consumed connector")
		}
	}
	owner := &fixtureConnector{}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if db, err := database.OpenConnector(ctx, owner, database.DefaultPoolConfig()); !errors.Is(err, context.Canceled) || db != nil {
		t.Fatalf("pre-canceled context: %v", err)
	}
	if owner.connects.Load() != 0 || owner.closedConnector.Load() != 0 {
		t.Fatal("pre-cancellation consumed connector")
	}
	if err := database.DefaultConfig().Validate(); err == nil {
		t.Fatal("legacy DSN requirements changed")
	}
}

func TestConnectorStartupFailuresReleaseResources(t *testing.T) {
	failure := errors.New("fixture unavailable")
	closeFailure := errors.New("fixture close failed")
	for _, phase := range []string{"connect", "ping"} {
		t.Run(phase, func(t *testing.T) {
			owner := &fixtureConnector{closeErr: closeFailure}
			if phase == "connect" {
				owner.connect = func(context.Context) error { return failure }
			} else {
				owner.ping = func(context.Context) error { return failure }
			}
			pool, err := database.OpenConnector(nil, owner, database.DefaultPoolConfig())
			if pool != nil || !errors.Is(err, failure) || !errors.Is(err, closeFailure) {
				t.Fatalf("lost failure or cleanup error: %v", err)
			}
			expectedConnections := int32(0)
			if phase == "ping" {
				expectedConnections = 1
			}
			if owner.closedConnector.Load() != 1 || owner.closedConnections.Load() != expectedConnections {
				t.Fatal("startup failure leaked resources")
			}
		})
	}
}

func TestConnectorTimeoutBoundsConnectAndPing(t *testing.T) {
	for _, phase := range []string{"connect", "ping"} {
		t.Run(phase, func(t *testing.T) {
			owner := &fixtureConnector{}
			cfg := database.DefaultPoolConfig()
			cfg.PingTimeout = 20 * time.Millisecond
			wait := func(ctx context.Context) error {
				if _, ok := ctx.Deadline(); !ok {
					return errors.New("missing deadline")
				}
				<-ctx.Done()
				return ctx.Err()
			}
			if phase == "connect" {
				owner.connect = wait
			} else {
				owner.ping = wait
			}
			pool, err := database.OpenConnector(context.Background(), owner, cfg)
			if pool != nil || !errors.Is(err, context.DeadlineExceeded) {
				t.Fatalf("%s timeout: %v", phase, err)
			}
			if owner.closedConnector.Load() != 1 {
				t.Fatal("timeout leaked connector")
			}
		})
	}
}

func TestConnectorHealthBoundsPoolWaitAndRecovers(t *testing.T) {
	cfg := database.DefaultPoolConfig()
	cfg.MaxOpenConns = 1
	cfg.MaxIdleConns = 1
	cfg.PingTimeout = 20 * time.Millisecond
	owner := &fixtureConnector{}
	pool, err := database.OpenConnector(nil, owner, cfg)
	if err != nil {
		t.Fatal(err)
	}
	defer pool.Close()
	conn, err := pool.Conn(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	if err = pool.Health(nil); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("waiting for busy pool was not bounded: %v", err)
	}
	if err = conn.Close(); err != nil {
		t.Fatal(err)
	}
	if err = pool.Health(nil); err != nil {
		t.Fatal(err)
	}
}

func TestConnectorIdleZeroAndLifetime(t *testing.T) {
	for _, mode := range []string{"noIdle", "expired"} {
		t.Run(mode, func(t *testing.T) {
			cfg := database.DefaultPoolConfig()
			if mode == "noIdle" {
				cfg.MaxIdleConns = 0
			} else {
				cfg.ConnMaxLifetime = time.Nanosecond
			}
			owner := &fixtureConnector{}
			pool, err := database.OpenConnector(nil, owner, cfg)
			if err != nil {
				t.Fatal(err)
			}
			defer pool.Close()
			if err = pool.Health(nil); err != nil {
				t.Fatal(err)
			}
			if owner.connects.Load() < 2 || owner.closedConnections.Load() == 0 {
				t.Fatal("pool did not retire connections")
			}
			if mode == "noIdle" && pool.Stats().Idle != 0 {
				t.Fatal("zero idle not honored")
			}
			if mode == "expired" && pool.Stats().MaxLifetimeClosed == 0 {
				t.Fatal("lifetime not honored")
			}
		})
	}
}

func TestConnectorTransactionsKeepCommitAndRollbackSemantics(t *testing.T) {
	owner := &fixtureConnector{}
	pool, err := database.OpenConnector(nil, owner, database.DefaultPoolConfig())
	if err != nil {
		t.Fatal(err)
	}
	defer pool.Close()
	ctx := context.Background()
	if err = pool.InTx(ctx, nil, func(*sql.Tx) error { return nil }); err != nil {
		t.Fatal(err)
	}
	failed := errors.New("business failed")
	if err = pool.InTx(ctx, nil, func(*sql.Tx) error { return failed }); !errors.Is(err, failed) {
		t.Fatal(err)
	}
	func() {
		defer func() {
			if recover() != failed {
				t.Error("panic not preserved")
			}
		}()
		_ = pool.InTx(ctx, nil, func(*sql.Tx) error { panic(failed) })
	}()
	if owner.commits.Load() != 1 || owner.rollbacks.Load() != 2 {
		t.Fatal("transaction semantics changed")
	}
}