package telemetry

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

	"github.com/go-sql-driver/mysql"
	pkgerrors "github.com/nexuscrm/backend/pkg/errors"
	"github.com/nexuscrm/shared/pkg/constants"
)

func init() {
	sql.Register("mysql-tracing", &TracingDriver{Parent: mysql.MySQLDriver{}})
}

// TracingDriver wraps the MySQL driver to inject trace context into SQL comments
type TracingDriver struct { //nolint:revive
	Parent mysql.MySQLDriver
}

// Open returns a wrapped connection
func (d *TracingDriver) Open(name string) (driver.Conn, error) {
	conn, err := d.Parent.Open(name)
	if err != nil {
		return nil, err
	}
	return &TracingConn{Conn: conn}, nil
}

// TracingConn wraps a driver.Conn
type TracingConn struct {
	driver.Conn
}

// Ping implements driver.Pinger
func (c *TracingConn) Ping(ctx context.Context) error {
	if pinger, ok := c.Conn.(driver.Pinger); ok {
		return pinger.Ping(ctx)
	}
	return nil
}

// BeginTx implements driver.ConnBeginTx
func (c *TracingConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
	if connBeginTx, ok := c.Conn.(driver.ConnBeginTx); ok {
		return connBeginTx.BeginTx(ctx, opts)
	}
	return c.Conn.Begin()
}

// PrepareContext implements driver.ConnPrepareContext
func (c *TracingConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	query = injectTraceComment(ctx, query)
	if connPrepareCtx, ok := c.Conn.(driver.ConnPrepareContext); ok {
		return connPrepareCtx.PrepareContext(ctx, query)
	}
	return c.Conn.Prepare(query)
}

// QueryContext implements driver.QueryerContext
func (c *TracingConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	query = injectTraceComment(ctx, query)
	if queryer, ok := c.Conn.(driver.QueryerContext); ok {
		return queryer.QueryContext(ctx, query, args)
	}
	return nil, pkgerrors.NewInternalError("underlying connection does not support QueryContext", nil)
}

// ExecContext implements driver.ExecerContext
func (c *TracingConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
	query = injectTraceComment(ctx, query)
	if execer, ok := c.Conn.(driver.ExecerContext); ok {
		return execer.ExecContext(ctx, query, args)
	}
	return nil, pkgerrors.NewInternalError("underlying connection does not support ExecContext", nil)
}

// injectTraceComment appends the W3C traceparent to the SQL query
func injectTraceComment(ctx context.Context, query string) string {
	// Use W3C Traceparent only
	if val := ctx.Value(constants.ContextKeyTraceparent); val != nil {
		if str, ok := val.(string); ok && str != "" {
			// Prevent SQL injection via header (basic sanitization)
			clean := strings.ReplaceAll(str, "*/", "")
			return fmt.Sprintf("%s /* traceparent='%s' */", query, clean)
		}
	}
	return query
}