package telemetry

import (
	"bytes"
	"context"
	"database/sql/driver"
	"encoding/json"
	"errors"
	"io"
	"net"
	"testing"
	"time"

	"github.com/go-sql-driver/mysql"
	"github.com/nexuscrm/kernel/pkg/perftiming"
	"github.com/nexuscrm/shared/pkg/constants"
	"github.com/rs/zerolog"
	zlog "github.com/rs/zerolog/log"
)

var (
	_ driver.DriverContext     = (*TracingDriver)(nil)
	_ driver.SessionResetter   = (*TracingConn)(nil)
	_ driver.Validator         = (*TracingConn)(nil)
	_ driver.NamedValueChecker = (*TracingConn)(nil)
)

type basicTestConn struct{}

func (*basicTestConn) Prepare(string) (driver.Stmt, error) { return nil, errors.New("not implemented") }
func (*basicTestConn) Close() error                        { return nil }
func (*basicTestConn) Begin() (driver.Tx, error)           { return nil, errors.New("not implemented") }

type optionalTestConn struct {
	basicTestConn
	resetErr error
	valid    bool
	checked  bool
}

type diagnosticTestTx struct{ committed bool }

func (t *diagnosticTestTx) Commit() error { t.committed = true; return nil }
func (*diagnosticTestTx) Rollback() error { return nil }

type diagnosticTxConn struct {
	basicTestConn
	tx *diagnosticTestTx
}

func (c *diagnosticTxConn) BeginTx(context.Context, driver.TxOptions) (driver.Tx, error) {
	c.tx = &diagnosticTestTx{}
	return c.tx, nil
}

func TestSQLDiagnosticsLogsPreparedTemplateDurationWithoutArguments(t *testing.T) {
	var output bytes.Buffer
	previousLogger := zlog.Logger
	zlog.Logger = zerolog.New(&output)
	defer func() { zlog.Logger = previousLogger }()

	conn := &TracingConn{Conn: &fallbackTestConn{}}
	ctx := context.WithValue(context.Background(), constants.ContextKeyRequestID, "request-sql-1")
	ctx = perftiming.WithEnabled(ctx)
	stmt, err := conn.PrepareContext(ctx, "SELECT Name FROM Account WHERE Id = ?")
	if err != nil {
		t.Fatal(err)
	}
	rows, err := stmt.(driver.StmtQueryContext).QueryContext(ctx, []driver.NamedValue{{Ordinal: 1, Value: "secret-record-id"}})
	if err != nil {
		t.Fatal(err)
	}
	if err := rows.Close(); err != nil {
		t.Fatal(err)
	}

	var event map[string]any
	if err := json.Unmarshal(bytes.TrimSpace(output.Bytes()), &event); err != nil {
		t.Fatalf("decode SQL diagnostic %q: %v", output.String(), err)
	}
	if event["event"] != "sql_statement_timing" || event["request_id"] != "request-sql-1" ||
		event["operation"] != "query" || event["sql"] != "SELECT Name FROM Account WHERE Id = ?" {
		t.Fatalf("SQL diagnostic = %#v", event)
	}
	if bytes.Contains(output.Bytes(), []byte("secret-record-id")) {
		t.Fatalf("SQL diagnostic leaked arguments: %s", output.String())
	}
}

func TestSQLDiagnosticsLogsTransactionBeginAndOutcome(t *testing.T) {
	var output bytes.Buffer
	previousLogger := zlog.Logger
	zlog.Logger = zerolog.New(&output)
	defer func() { zlog.Logger = previousLogger }()

	underlying := &diagnosticTxConn{}
	conn := &TracingConn{Conn: underlying}
	ctx := context.WithValue(context.Background(), constants.ContextKeyRequestID, "request-tx-1")
	ctx = perftiming.WithEnabled(ctx)
	tx, err := conn.BeginTx(ctx, driver.TxOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if err := tx.Commit(); err != nil {
		t.Fatal(err)
	}
	if underlying.tx == nil || !underlying.tx.committed {
		t.Fatal("underlying transaction was not committed")
	}

	lines := bytes.Split(bytes.TrimSpace(output.Bytes()), []byte("\n"))
	if len(lines) != 2 {
		t.Fatalf("transaction diagnostic lines = %d, want 2: %s", len(lines), output.String())
	}
	stages := []string{"begin", "commit"}
	var transactionID float64
	for index, line := range lines {
		var event map[string]any
		if err := json.Unmarshal(line, &event); err != nil {
			t.Fatalf("decode transaction diagnostic %q: %v", line, err)
		}
		if event["event"] != "sql_transaction_timing" || event["request_id"] != "request-tx-1" || event["stage"] != stages[index] {
			t.Fatalf("transaction diagnostic %d = %#v", index, event)
		}
		if index == 0 {
			transactionID, _ = event["driver_transaction_id"].(float64)
		} else if event["driver_transaction_id"] != transactionID {
			t.Fatalf("transaction IDs differ: begin %v outcome %v", transactionID, event["driver_transaction_id"])
		}
	}
}

type observingTestConn struct {
	basicTestConn
	queries int
	execs   int
}

type delayedTestRows struct {
	delay time.Duration
	done  bool
}

func (*delayedTestRows) Columns() []string { return []string{"value"} }
func (*delayedTestRows) Close() error      { return nil }
func (r *delayedTestRows) Next(dest []driver.Value) error {
	if r.done {
		return io.EOF
	}
	time.Sleep(r.delay)
	dest[0] = "value"
	r.done = true
	return nil
}

type successfulObservingTestConn struct {
	basicTestConn
	delay time.Duration
}

func (c *successfulObservingTestConn) QueryContext(context.Context, string, []driver.NamedValue) (driver.Rows, error) {
	return &delayedTestRows{delay: c.delay}, nil
}

type fallbackTestConn struct {
	basicTestConn
	delay time.Duration
}

func (*fallbackTestConn) QueryContext(context.Context, string, []driver.NamedValue) (driver.Rows, error) {
	return nil, driver.ErrSkip
}

func (c *fallbackTestConn) PrepareContext(context.Context, string) (driver.Stmt, error) {
	return &fallbackTestStmt{delay: c.delay}, nil
}

type fallbackTestStmt struct {
	delay time.Duration
}

func (*fallbackTestStmt) Close() error  { return nil }
func (*fallbackTestStmt) NumInput() int { return -1 }
func (*fallbackTestStmt) Exec([]driver.Value) (driver.Result, error) {
	return nil, errors.New("legacy exec must not be used")
}
func (*fallbackTestStmt) Query([]driver.Value) (driver.Rows, error) {
	return nil, errors.New("legacy query must not be used")
}
func (s *fallbackTestStmt) QueryContext(context.Context, []driver.NamedValue) (driver.Rows, error) {
	return &delayedTestRows{delay: s.delay}, nil
}

func TestTracingDriverOpenConnectorPreservesConnectContext(t *testing.T) {
	const network = "telemetry_context_test"
	started := make(chan struct{})
	mysql.RegisterDialContext(network, func(ctx context.Context, _ string) (net.Conn, error) {
		close(started)
		<-ctx.Done()
		return nil, ctx.Err()
	})

	connector, err := (&TracingDriver{Parent: mysql.MySQLDriver{}}).OpenConnector("user:password@" + network + "(unused)/db")
	if err != nil {
		t.Fatalf("OpenConnector() error = %v", err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	connectDone := make(chan error, 1)
	go func() {
		_, connectErr := connector.Connect(ctx)
		connectDone <- connectErr
	}()
	select {
	case <-started:
	case <-time.After(time.Second):
		t.Fatal("underlying context dial did not start")
	}
	cancel()
	select {
	case connectErr := <-connectDone:
		if !errors.Is(connectErr, context.Canceled) {
			t.Fatalf("Connect() error = %v, want context canceled", connectErr)
		}
	case <-time.After(time.Second):
		t.Fatal("Connect() did not return after context cancellation")
	}
}

func (c *observingTestConn) QueryContext(context.Context, string, []driver.NamedValue) (driver.Rows, error) {
	c.queries++
	return nil, errors.New("query probe")
}

func (c *observingTestConn) ExecContext(context.Context, string, []driver.NamedValue) (driver.Result, error) {
	c.execs++
	return nil, errors.New("exec probe")
}

func (c *optionalTestConn) ResetSession(context.Context) error { return c.resetErr }
func (c *optionalTestConn) IsValid() bool                      { return c.valid }
func (c *optionalTestConn) CheckNamedValue(value *driver.NamedValue) error {
	c.checked = true
	value.Value = "converted"
	return nil
}

func TestTracingConnForwardsOptionalConnectionContracts(t *testing.T) {
	resetErr := errors.New("closed connection")
	underlying := &optionalTestConn{resetErr: resetErr, valid: false}
	conn := &TracingConn{Conn: underlying}

	if err := conn.ResetSession(context.Background()); !errors.Is(err, resetErr) {
		t.Fatalf("ResetSession() error = %v, want %v", err, resetErr)
	}
	if conn.IsValid() {
		t.Fatal("IsValid() = true, want wrapped false")
	}
	value := &driver.NamedValue{Ordinal: 1, Value: "raw"}
	if err := conn.CheckNamedValue(value); err != nil {
		t.Fatalf("CheckNamedValue() error = %v", err)
	}
	if !underlying.checked || value.Value != "converted" {
		t.Fatalf("CheckNamedValue() did not delegate: checked=%t value=%v", underlying.checked, value.Value)
	}
}

func TestTracingConnOptionalConnectionFallbacks(t *testing.T) {
	conn := &TracingConn{Conn: &basicTestConn{}}

	if err := conn.ResetSession(context.Background()); err != nil {
		t.Fatalf("ResetSession() fallback error = %v", err)
	}
	if !conn.IsValid() {
		t.Fatal("IsValid() fallback = false, want true")
	}
	if err := conn.CheckNamedValue(&driver.NamedValue{}); !errors.Is(err, driver.ErrSkip) {
		t.Fatalf("CheckNamedValue() fallback error = %v, want driver.ErrSkip", err)
	}
}

func TestSQLStatementCounterIsRequestScopedAndCountsDriverAttempts(t *testing.T) {
	underlying := &observingTestConn{}
	conn := &TracingConn{Conn: underlying}
	observed, counter := WithSQLStatementCounter(context.Background())

	_, _ = conn.QueryContext(observed, "SELECT 1", nil)
	_, _ = conn.ExecContext(observed, "UPDATE t SET v = 1", nil)
	_, _ = conn.QueryContext(context.Background(), "SELECT 2", nil)

	if got := counter.Count(); got != 2 {
		t.Fatalf("request SQL count = %d, want 2", got)
	}
	if counter.TotalDuration() <= 0 || counter.MaxDuration() <= 0 || counter.MaxDuration() > counter.TotalDuration() {
		t.Fatalf("request SQL duration total=%s max=%s", counter.TotalDuration(), counter.MaxDuration())
	}
	if underlying.queries != 2 || underlying.execs != 1 {
		t.Fatalf("underlying attempts = queries %d execs %d", underlying.queries, underlying.execs)
	}
}

func TestNestedSQLStatementCountersPreserveWholeRequestAndStageCounts(t *testing.T) {
	conn := &TracingConn{Conn: &observingTestConn{}}
	requestCtx, requestCounter := WithSQLStatementCounter(context.Background())
	_, _ = conn.QueryContext(requestCtx, "SELECT outer", nil)
	stageCtx, stageCounter := WithSQLStatementCounter(requestCtx)
	_, _ = conn.QueryContext(stageCtx, "SELECT nested", nil)

	if got := requestCounter.Count(); got != 2 {
		t.Fatalf("whole-request SQL count = %d, want 2", got)
	}
	if got := stageCounter.Count(); got != 1 {
		t.Fatalf("nested-stage SQL count = %d, want 1", got)
	}
	if requestCounter.TotalDuration() < stageCounter.TotalDuration() || stageCounter.MaxDuration() <= 0 {
		t.Fatalf("nested durations request=%s stage=%s/%s", requestCounter.TotalDuration(), stageCounter.TotalDuration(), stageCounter.MaxDuration())
	}
}

func TestSQLStatementDurationIncludesResultSetConsumption(t *testing.T) {
	const rowDelay = 20 * time.Millisecond
	conn := &TracingConn{Conn: &successfulObservingTestConn{delay: rowDelay}}
	observed, counter := WithSQLStatementCounter(context.Background())

	rows, err := conn.QueryContext(observed, "SELECT slow_rows", nil)
	if err != nil {
		t.Fatalf("QueryContext() error = %v", err)
	}
	if got := counter.Count(); got != 1 {
		t.Fatalf("request SQL count before consumption = %d, want 1", got)
	}
	if got := counter.TotalDuration(); got != 0 {
		t.Fatalf("request SQL duration before consumption = %s, want 0", got)
	}
	values := make([]driver.Value, 1)
	if err := rows.Next(values); err != nil {
		t.Fatalf("Next() error = %v", err)
	}
	if err := rows.Next(values); !errors.Is(err, io.EOF) {
		t.Fatalf("terminal Next() error = %v, want EOF", err)
	}
	if counter.TotalDuration() < rowDelay || counter.MaxDuration() < rowDelay {
		t.Fatalf("result-set duration total=%s max=%s, want >= %s", counter.TotalDuration(), counter.MaxDuration(), rowDelay)
	}
	beforeClose := counter.TotalDuration()
	if err := rows.Close(); err != nil {
		t.Fatalf("Close() error = %v", err)
	}
	if got := counter.TotalDuration(); got != beforeClose {
		t.Fatalf("Close() double-counted duration: before=%s after=%s", beforeClose, got)
	}
}

func TestSQLStatementCounterFollowsPreparedStatementFallback(t *testing.T) {
	const rowDelay = 20 * time.Millisecond
	conn := &TracingConn{Conn: &fallbackTestConn{delay: rowDelay}}
	observed, counter := WithSQLStatementCounter(context.Background())

	rows, err := conn.QueryContext(observed, "SELECT value WHERE id = ?", []driver.NamedValue{{Ordinal: 1, Value: "id"}})
	if !errors.Is(err, driver.ErrSkip) || rows != nil {
		t.Fatalf("connection QueryContext() = %#v, %v; want driver.ErrSkip", rows, err)
	}
	if got := counter.Count(); got != 0 {
		t.Fatalf("fallback probe SQL count = %d, want 0", got)
	}
	stmt, err := conn.PrepareContext(observed, "SELECT value WHERE id = ?")
	if err != nil {
		t.Fatalf("PrepareContext() error = %v", err)
	}
	rows, err = stmt.(driver.StmtQueryContext).QueryContext(observed, []driver.NamedValue{{Ordinal: 1, Value: "id"}})
	if err != nil {
		t.Fatalf("prepared QueryContext() error = %v", err)
	}
	values := make([]driver.Value, 1)
	if err := rows.Next(values); err != nil {
		t.Fatalf("Next() error = %v", err)
	}
	if err := rows.Close(); err != nil {
		t.Fatalf("Close() error = %v", err)
	}
	if got := counter.Count(); got != 1 {
		t.Fatalf("prepared SQL count = %d, want 1", got)
	}
	if counter.TotalDuration() < rowDelay {
		t.Fatalf("prepared result-set duration = %s, want >= %s", counter.TotalDuration(), rowDelay)
	}
}