package database_test

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"database/sql"
	"database/sql/driver"
	"errors"
	"net"
	"os"
	"sync/atomic"
	"testing"
	"time"

	"git.yuntuo.hk/k2safe/go-omni/database"
	"github.com/go-sql-driver/mysql"
)

type connectorDriverLog struct{ count atomic.Int32 }

func (l *connectorDriverLog) Print(...any) { l.count.Add(1) }

func TestMySQLConnectorDriverLoggerIsolation(t *testing.T) {
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	stopped := make(chan struct{})
	go func() {
		defer close(stopped)
		for {
			c, err := listener.Accept()
			if err != nil {
				return
			}
			c.Close()
		}
	}()
	defer func() { listener.Close(); <-stopped }()
	logs := [2]connectorDriverLog{}
	for i := range logs {
		cfg := mysql.NewConfig()
		cfg.Net = "tcp"
		cfg.Addr = listener.Addr().String()
		cfg.Timeout = time.Second
		cfg.ReadTimeout = time.Second
		cfg.Logger = &logs[i]
		connector, err := mysql.NewConnector(cfg)
		if err != nil {
			t.Fatal(err)
		}
		before := [2]int32{logs[0].count.Load(), logs[1].count.Load()}
		db, err := database.OpenConnector(context.Background(), connector, database.DefaultPoolConfig())
		if err == nil {
			db.Close()
			t.Fatal("handshake to closed socket succeeded")
		}
		if logs[i].count.Load() <= before[i] || logs[1-i].count.Load() != before[1-i] {
			t.Fatal("driver logger did not stay on its connector")
		}
	}
}

func TestMySQLConnectorTLSAndPoolIsolation(t *testing.T) {
	if os.Getenv("GO_OMNI_CONNECTOR_MYSQL_TEST") != "1" {
		t.Skip("use scripts/verifyConnector.py for isolated MySQL TLS verification")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	ca, err := os.ReadFile(os.Getenv("GO_OMNI_CONNECTOR_CA_FILE"))
	if err != nil {
		t.Fatal(err)
	}
	roots := x509.NewCertPool()
	if !roots.AppendCertsFromPEM(ca) {
		t.Fatal("invalid test CA")
	}
	cfg := mysql.NewConfig()
	cfg.Net = "tcp"
	cfg.Addr = os.Getenv("GO_OMNI_CONNECTOR_ADDR")
	cfg.User = os.Getenv("GO_OMNI_CONNECTOR_USER")
	cfg.Passwd = os.Getenv("GO_OMNI_CONNECTOR_PASSWORD")
	cfg.DBName = os.Getenv("GO_OMNI_CONNECTOR_DATABASE")
	cfg.Timeout = time.Second
	cfg.ReadTimeout = 2 * time.Second
	cfg.WriteTimeout = 2 * time.Second
	cfg.TLS = &tls.Config{MinVersion: tls.VersionTLS12, ServerName: "localhost", RootCAs: roots}
	pool := database.DefaultPoolConfig()
	pool.MaxOpenConns = 2
	pool.MaxIdleConns = 2
	pool.PingTimeout = 2 * time.Second
	open := func(cfg *mysql.Config) (*database.DB, error) {
		connector, err := mysql.NewConnector(cfg)
		if err != nil {
			return nil, err
		}
		return database.OpenConnector(ctx, connector, pool)
	}
	cfg.Params = map[string]string{"time_zone": "'+00:00'"}
	first, err := open(cfg)
	if err != nil {
		t.Fatal(err)
	}
	defer first.Close()
	other := cfg.Clone()
	other.Params["time_zone"] = "'+08:00'"
	second, err := open(other)
	if err != nil {
		t.Fatal(err)
	}
	defer second.Close()
	for i, db := range []*database.DB{first, second} {
		var name, cipher, zone string
		if err := db.QueryRowContext(ctx, "SHOW SESSION STATUS LIKE 'Ssl_cipher'").Scan(&name, &cipher); err != nil {
			t.Fatal(err)
		}
		if cipher == "" {
			t.Fatal("connection is not encrypted")
		}
		if err := db.QueryRowContext(ctx, "SELECT @@session.time_zone").Scan(&zone); err != nil {
			t.Fatal(err)
		}
		if zone != []string{"+00:00", "+08:00"}[i] {
			t.Fatalf("session config leaked: %s", zone)
		}
		if db.Stats().MaxOpenConnections != 2 {
			t.Fatal("pool configuration ignored")
		}
	}
	for _, bad := range []string{"untrustedCA", "wrongHostname"} {
		t.Run(bad, func(t *testing.T) {
			invalid := cfg.Clone()
			if bad == "untrustedCA" {
				invalid.TLS.RootCAs = x509.NewCertPool()
			} else {
				invalid.TLS.ServerName = "wrong.invalid"
			}
			db, err := open(invalid)
			if err == nil {
				db.Close()
				t.Fatal("invalid server certificate accepted")
			}
			var verification *tls.CertificateVerificationError
			if !errors.As(err, &verification) {
				t.Fatalf("not a TLS verification failure: %v", err)
			}
			if err := first.Health(ctx); err != nil {
				t.Fatalf("failed connector affected existing pool: %v", err)
			}
		})
	}
	if _, err := first.ExecContext(ctx, "CREATE TABLE connector_check (id INT PRIMARY KEY)"); err != nil {
		t.Fatal(err)
	}
	if err := first.InTx(ctx, nil, func(tx *sql.Tx) error {
		_, err := tx.ExecContext(ctx, "INSERT INTO connector_check VALUES (1)")
		return err
	}); err != nil {
		t.Fatal(err)
	}
	rollbackErr := errors.New("rollback requested")
	if err := first.InTx(ctx, nil, func(tx *sql.Tx) error {
		if _, err := tx.ExecContext(ctx, "INSERT INTO connector_check VALUES (2)"); err != nil {
			return err
		}
		return rollbackErr
	}); !errors.Is(err, rollbackErr) {
		t.Fatal(err)
	} else {
		var txErr *database.TxError
		if !errors.As(err, &txErr) || txErr.Stage != database.TxStageWork || txErr.Outcome != database.TxOutcomeRolledBack {
			t.Fatal("real MySQL rollback was not classified", err)
		}
	}
	var count int
	if err := second.QueryRowContext(ctx, "SELECT COUNT(*) FROM connector_check").Scan(&count); err != nil || count != 1 {
		t.Fatalf("transaction count=%d err=%v", count, err)
	}
	// Commit/Rollback run against real MySQL, then a driver-boundary wrapper
	// suppresses the acknowledgement. This is deterministic fault injection,
	// not a claim to reproduce every wire-level disconnect timing.
	for _, loseCommit := range []bool{true, false} {
		name, id := "commitAcknowledgementLost", 3
		if !loseCommit {
			name, id = "rollbackAcknowledgementLost", 4
		}
		t.Run(name, func(t *testing.T) {
			connector, err := mysql.NewConnector(cfg.Clone())
			if err != nil {
				t.Fatal(err)
			}
			ackErr, workErr := errors.New("injected lost acknowledgement"), errors.New("work rejected")
			fault := &ackLossConnector{Connector: connector, loseCommit: loseCommit, ackErr: ackErr}
			faultPool, err := database.OpenConnector(ctx, fault, pool)
			if err != nil {
				t.Fatal(err)
			}
			defer faultPool.Close()
			err = faultPool.InTx(ctx, nil, func(tx *sql.Tx) error {
				if _, err := tx.ExecContext(ctx, "INSERT INTO connector_check VALUES (?)", id); err != nil {
					return err
				}
				if !loseCommit {
					return workErr
				}
				return nil
			})
			var txErr *database.TxError
			if !errors.As(err, &txErr) || txErr.Outcome != database.TxOutcomeUnknown || !errors.Is(err, ackErr) {
				t.Fatal("unknown acknowledgement was treated as definite outcome", err)
			}
			wantCount := 0
			if loseCommit {
				wantCount = 1
				if txErr.Stage != database.TxStageCommit || txErr.Err != ackErr || txErr.RollbackErr != sql.ErrTxDone || fault.commits.Load() != 1 || fault.rollbacks.Load() != 0 {
					t.Fatal("commit failure was retried or misclassified", err)
				}
			} else if txErr.Stage != database.TxStageWork || !errors.Is(err, workErr) || txErr.RollbackErr != ackErr || fault.commits.Load() != 0 || fault.rollbacks.Load() != 1 {
				t.Fatal("rollback failure lost its primary/secondary causes", err)
			}
			var stored int
			if err := second.QueryRowContext(ctx, "SELECT COUNT(*) FROM connector_check WHERE id = ?", id).Scan(&stored); err != nil || stored != wantCount {
				t.Fatal("independent connection observed unexpected persistence", stored, err)
			}
		})
	}
	if err := first.Close(); err != nil {
		t.Fatal(err)
	}
	if err := second.Health(ctx); err != nil {
		t.Fatal("closing first pool affected second:", err)
	}
}

type ackLossConnector struct {
	driver.Connector
	loseCommit         bool
	ackErr             error
	commits, rollbacks atomic.Int32
}

func (c *ackLossConnector) Connect(ctx context.Context) (driver.Conn, error) {
	conn, err := c.Connector.Connect(ctx)
	if err != nil {
		return nil, err
	}
	return &ackLossConn{Conn: conn, owner: c}, nil
}

type ackLossConn struct {
	driver.Conn
	owner *ackLossConnector
}

func (c *ackLossConn) Ping(ctx context.Context) error { return c.Conn.(driver.Pinger).Ping(ctx) }
func (c *ackLossConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
	tx, err := c.Conn.(driver.ConnBeginTx).BeginTx(ctx, opts)
	if err != nil {
		return nil, err
	}
	return &ackLossTx{Tx: tx, owner: c.owner}, nil
}

type ackLossTx struct {
	driver.Tx
	owner *ackLossConnector
}

func (tx *ackLossTx) Commit() error {
	tx.owner.commits.Add(1)
	if err := tx.Tx.Commit(); err != nil {
		return err
	}
	if tx.owner.loseCommit {
		return tx.owner.ackErr
	}
	return nil
}
func (tx *ackLossTx) Rollback() error {
	tx.owner.rollbacks.Add(1)
	if err := tx.Tx.Rollback(); err != nil {
		return err
	}
	if !tx.owner.loseCommit {
		return tx.owner.ackErr
	}
	return nil
}