package tracing

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"strings"
	"unicode"

	"github.com/XSAM/otelsql"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/metric/noop"
	"go.opentelemetry.io/otel/trace"
	otelnoop "go.opentelemetry.io/otel/trace/noop"
)

const (
	dbSystemMySQL      = "mysql"
	dbOperationKey     = "db.operation"
	dbOperationUnknown = "<unknown>"
)

// SQLOpener opens a database/sql DB through AgentTeam's tracing boundary.
type SQLOpener func(driverName string, dataSourceName string) (*sql.DB, error)

// DBOptions controls database tracing behavior.
type DBOptions struct {
	StatementEnabled bool
}

// NewSQLOpener returns a database/sql opener that traces query and exec calls.
func NewSQLOpener(tracerProvider trace.TracerProvider, options DBOptions) SQLOpener {
	if tracerProvider == nil {
		tracerProvider = otelnoop.NewTracerProvider()
	}
	return func(driverName string, dataSourceName string) (*sql.DB, error) {
		return otelsql.Open(driverName, dataSourceName,
			otelsql.WithTracerProvider(tracerProvider),
			otelsql.WithMeterProvider(noop.NewMeterProvider()),
			otelsql.WithSQLCommenter(false),
			otelsql.WithAttributes(attribute.String("db.system", dbSystemMySQL)),
			otelsql.WithSpanNameFormatter(func(_ context.Context, _ otelsql.Method, _ string) string {
				return "db.query"
			}),
			otelsql.WithAttributesGetter(func(_ context.Context, _ otelsql.Method, query string, _ []driver.NamedValue) []attribute.KeyValue {
				return []attribute.KeyValue{
					attribute.String(dbOperationKey, operationClass(query)),
				}
			}),
			otelsql.WithSpanOptions(otelsql.SpanOptions{
				DisableQuery:         !options.StatementEnabled,
				DisableErrSkip:       true,
				OmitConnPrepare:      true,
				OmitConnResetSession: true,
				OmitConnectorConnect: true,
				OmitRows:             true,
				SpanFilter:           dbSpanFilter,
			}),
		)
	}
}

func dbSpanFilter(_ context.Context, method otelsql.Method, _ string, _ []driver.NamedValue) bool {
	switch method {
	case otelsql.MethodConnExec, otelsql.MethodConnQuery, otelsql.MethodStmtExec, otelsql.MethodStmtQuery:
		return true
	default:
		return false
	}
}

func operationClass(query string) string {
	trimmed := trimLeadingSQLComments(strings.TrimSpace(query))
	if trimmed == "" {
		return dbOperationUnknown
	}

	end := 0
	for index, r := range trimmed {
		if !unicode.IsLetter(r) {
			break
		}
		end = index + len(string(r))
	}
	if end == 0 {
		return dbOperationUnknown
	}
	return strings.ToLower(trimmed[:end])
}

func trimLeadingSQLComments(query string) string {
	for {
		query = strings.TrimSpace(query)
		switch {
		case strings.HasPrefix(query, "--"):
			newline := strings.IndexByte(query, '\n')
			if newline == -1 {
				return ""
			}
			query = query[newline+1:]
		case strings.HasPrefix(query, "/*"):
			end := strings.Index(query, "*/")
			if end == -1 {
				return ""
			}
			query = query[end+2:]
		default:
			return query
		}
	}
}