package tracing

import (
	"context"
	"database/sql"
	"regexp"
	"strings"
	"testing"
	"time"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/stretchr/testify/require"
	"go.opentelemetry.io/otel/attribute"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/sdk/trace/tracetest"
	"go.opentelemetry.io/otel/trace"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
)

func TestSQLOpenerTracesQueryAndExecWithoutSQLText(t *testing.T) {
	t.Parallel()

	dsn := "observability_tracing_db"
	mockDB, mock, err := sqlmock.NewWithDSN(dsn)
	require.NoError(t, err)
	mock.MatchExpectationsInOrder(false)
	mock.ExpectQuery(regexp.QuoteMeta("SELECT 1")).WillReturnRows(sqlmock.NewRows([]string{"value"}).AddRow(1))
	mock.ExpectExec(regexp.QuoteMeta("UPDATE things SET name = ?")).WithArgs("secret-name").WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectClose()
	mock.ExpectClose()

	recorder := tracetest.NewSpanRecorder()
	tracerProvider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder))
	t.Cleanup(func() {
		require.NoError(t, tracerProvider.Shutdown(context.Background()))
	})

	tracedDB, err := NewSQLOpener(tracerProvider, DBOptions{})("sqlmock", dsn)
	require.NoError(t, err)

	var value int
	require.NoError(t, tracedDB.QueryRowContext(context.Background(), "SELECT 1").Scan(&value))
	require.Equal(t, 1, value)
	_, err = tracedDB.ExecContext(context.Background(), "UPDATE things SET name = ?", "secret-name")
	require.NoError(t, err)

	require.NoError(t, tracedDB.Close())
	require.NoError(t, mockDB.Close())
	require.NoError(t, mock.ExpectationsWereMet())

	spans := recorder.Ended()
	require.Len(t, spans, 2)
	require.ElementsMatch(t, []string{"select", "update"}, []string{
		spanStringAttribute(t, spans[0], dbOperationKey),
		spanStringAttribute(t, spans[1], dbOperationKey),
	})
	for _, span := range spans {
		require.Equal(t, "db.query", span.Name())
		require.Equal(t, trace.SpanKindClient, span.SpanKind())
		require.Equal(t, dbSystemMySQL, spanStringAttribute(t, span, "db.system"))
		require.False(t, spanHasAttribute(span, "db.statement"))
		require.False(t, spanHasAttribute(span, "db.query.text"))
		require.False(t, spanHasAttributeValue(span, "secret-name"))
	}
}

func TestSQLOpenerTracesCoreDBAdminStartupPreflight(t *testing.T) {
	t.Parallel()

	dsn := "observability_admin_startup_db"
	mockDB, mock, err := sqlmock.NewWithDSN(dsn, sqlmock.MonitorPingsOption(true))
	require.NoError(t, err)
	mock.ExpectPing()
	mock.ExpectQuery(regexp.QuoteMeta("SELECT VERSION()")).
		WillReturnRows(sqlmock.NewRows([]string{"version"}).AddRow("8.0.11-TiDB-v8.5.1"))
	mock.ExpectQuery(regexp.QuoteMeta("SELECT @@tidb_txn_mode")).
		WillReturnRows(sqlmock.NewRows([]string{"tidb_txn_mode"}).AddRow("optimistic"))
	mock.ExpectClose()
	mock.ExpectClose()

	recorder := tracetest.NewSpanRecorder()
	tracerProvider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder))
	t.Cleanup(func() {
		require.NoError(t, tracerProvider.Shutdown(context.Background()))
	})

	tracedOpener := NewSQLOpener(tracerProvider, DBOptions{})
	admin := coredb.NewAdminWithOptions(coredb.WithSQLOpener(func(driverName string, dataSourceName string) (*sql.DB, error) {
		require.Equal(t, "mysql", driverName)
		require.Contains(t, dataSourceName, "parseTime=true")
		require.Contains(t, dataSourceName, "time_zone=%27%2B00%3A00%27")
		return tracedOpener("sqlmock", dsn)
	}))

	db, err := admin.Open(context.Background(), "user:secret@tcp(127.0.0.1:4000)/agentteam", time.Second)
	require.NoError(t, err)
	require.NoError(t, admin.RunStartupPreflight(context.Background(), db))

	db.Close()
	require.NoError(t, mockDB.Close())
	require.NoError(t, mock.ExpectationsWereMet())

	spans := recorder.Ended()
	require.Len(t, spans, 2)
	for _, span := range spans {
		require.Equal(t, "db.query", span.Name())
		require.Equal(t, trace.SpanKindClient, span.SpanKind())
		require.Equal(t, dbSystemMySQL, spanStringAttribute(t, span, "db.system"))
		require.Equal(t, "select", spanStringAttribute(t, span, dbOperationKey))
		require.False(t, spanHasAttribute(span, "db.statement"))
		require.False(t, spanHasAttribute(span, "db.query.text"))
		require.False(t, spanHasAttributeValue(span, "secret"))
	}
}

func TestSQLOpenerCanTraceSQLStatementTemplate(t *testing.T) {
	t.Setenv("OTEL_SEMCONV_STABILITY_OPT_IN", "")

	dsn := "observability_tracing_db_statement_template"
	mockDB, mock, err := sqlmock.NewWithDSN(dsn)
	require.NoError(t, err)
	mock.MatchExpectationsInOrder(false)
	mock.ExpectQuery(regexp.QuoteMeta("SELECT 1")).WillReturnRows(sqlmock.NewRows([]string{"value"}).AddRow(1))
	mock.ExpectExec(regexp.QuoteMeta("UPDATE things SET name = ?")).WithArgs("secret-name").WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectClose()
	mock.ExpectClose()

	recorder := tracetest.NewSpanRecorder()
	tracerProvider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder))
	t.Cleanup(func() {
		require.NoError(t, tracerProvider.Shutdown(context.Background()))
	})

	tracedDB, err := NewSQLOpener(tracerProvider, DBOptions{StatementEnabled: true})("sqlmock", dsn)
	require.NoError(t, err)

	var value int
	require.NoError(t, tracedDB.QueryRowContext(context.Background(), "SELECT 1").Scan(&value))
	require.Equal(t, 1, value)
	_, err = tracedDB.ExecContext(context.Background(), "UPDATE things SET name = ?", "secret-name")
	require.NoError(t, err)

	require.NoError(t, tracedDB.Close())
	require.NoError(t, mockDB.Close())
	require.NoError(t, mock.ExpectationsWereMet())

	spans := recorder.Ended()
	require.Len(t, spans, 2)
	require.True(t, spansHaveAttributeContaining(spans, "db.statement", "SELECT 1"))
	require.True(t, spansHaveAttributeContaining(spans, "db.statement", "UPDATE things SET name = ?"))
	for _, span := range spans {
		require.False(t, spanContainsString(span, "secret-name"))
	}
}

func TestSQLOpenerCanTraceStableSQLQueryText(t *testing.T) {
	t.Setenv("OTEL_SEMCONV_STABILITY_OPT_IN", "database")

	dsn := "observability_tracing_db_query_text"
	mockDB, mock, err := sqlmock.NewWithDSN(dsn)
	require.NoError(t, err)
	mock.ExpectExec(regexp.QuoteMeta("UPDATE things SET name = ?")).WithArgs("secret-name").WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectClose()
	mock.ExpectClose()

	recorder := tracetest.NewSpanRecorder()
	tracerProvider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder))
	t.Cleanup(func() {
		require.NoError(t, tracerProvider.Shutdown(context.Background()))
	})

	tracedDB, err := NewSQLOpener(tracerProvider, DBOptions{StatementEnabled: true})("sqlmock", dsn)
	require.NoError(t, err)

	_, err = tracedDB.ExecContext(context.Background(), "UPDATE things SET name = ?", "secret-name")
	require.NoError(t, err)

	require.NoError(t, tracedDB.Close())
	require.NoError(t, mockDB.Close())
	require.NoError(t, mock.ExpectationsWereMet())

	spans := recorder.Ended()
	require.Len(t, spans, 1)
	require.True(t, spanHasAttributeContaining(spans[0], "db.query.text", "UPDATE things SET name = ?"))
	require.False(t, spanHasAttribute(spans[0], "db.statement"))
	require.False(t, spanContainsString(spans[0], "secret-name"))
}

func TestOperationClass(t *testing.T) {
	t.Parallel()

	require.Equal(t, "select", operationClass(" SELECT * FROM things"))
	require.Equal(t, "update", operationClass("\n/* leading comment */\nUPDATE things SET name = ?"))
	require.Equal(t, "insert", operationClass("-- comment\nINSERT INTO things(id) VALUES (?)"))
	require.Equal(t, dbOperationUnknown, operationClass(""))
	require.Equal(t, dbOperationUnknown, operationClass("/* unterminated"))
}

func spanStringAttribute(t *testing.T, span sdktrace.ReadOnlySpan, key string) string {
	t.Helper()
	for _, attr := range span.Attributes() {
		if string(attr.Key) == key {
			return attr.Value.AsString()
		}
	}
	require.Failf(t, "missing span attribute", "key %q", key)
	return ""
}

func spanHasAttribute(span sdktrace.ReadOnlySpan, key string) bool {
	for _, attr := range span.Attributes() {
		if string(attr.Key) == key {
			return true
		}
	}
	return false
}

func spanHasAttributeValue(span sdktrace.ReadOnlySpan, value string) bool {
	for _, attr := range span.Attributes() {
		if attr.Value.Type() == attribute.STRING && attr.Value.AsString() == value {
			return true
		}
	}
	return false
}

func spansHaveAttributeContaining(spans []sdktrace.ReadOnlySpan, key string, substring string) bool {
	for _, span := range spans {
		if spanHasAttributeContaining(span, key, substring) {
			return true
		}
	}
	return false
}

func spanHasAttributeContaining(span sdktrace.ReadOnlySpan, key string, substring string) bool {
	for _, attr := range span.Attributes() {
		if string(attr.Key) == key && attr.Value.Type() == attribute.STRING && strings.Contains(attr.Value.AsString(), substring) {
			return true
		}
	}
	return false
}

func spanContainsString(span sdktrace.ReadOnlySpan, substring string) bool {
	for _, attr := range span.Attributes() {
		if attr.Value.Type() == attribute.STRING && strings.Contains(attr.Value.AsString(), substring) {
			return true
		}
	}
	for _, event := range span.Events() {
		if strings.Contains(event.Name, substring) {
			return true
		}
		for _, attr := range event.Attributes {
			if attr.Value.Type() == attribute.STRING && strings.Contains(attr.Value.AsString(), substring) {
				return true
			}
		}
	}
	return strings.Contains(span.Status().Description, substring)
}