package mirko

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"time"

	"github.com/gchaincl/sqlhooks"
	"golang.org/x/net/trace"
)

type sqlHooks struct{}

func (h *sqlHooks) Before(ctx context.Context, query string, args ...interface{}) (context.Context, error) {
	tr, ok := trace.FromContext(ctx)
	if ok {
		tr.LazyPrintf("SQL query: %s", query)
		tr.LazyPrintf("SQL  args: %+v", args)
	}
	return context.WithValue(ctx, "begin", time.Now()), nil
}

func (h *sqlHooks) After(ctx context.Context, query string, args ...interface{}) (context.Context, error) {
	begin := ctx.Value("begin").(time.Time)
	tr, ok := trace.FromContext(ctx)
	if ok {
		tr.LazyPrintf("SQL  took: %s", time.Since(begin).String())
	}
	return ctx, nil
}

func TraceSQL(driver driver.Driver, wrapped string) {
	sql.Register(wrapped, sqlhooks.Wrap(driver, &sqlHooks{}))
}