package telemetry

import (
	"context"
	"database/sql"
	"github.com/redis/go-redis/v9"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/trace"
	"io"
	"net"
	"sync"
)

type spanBody struct {
	io.ReadCloser
	span trace.Span
	once sync.Once
}

func (b *spanBody) finish(err error) {
	b.once.Do(func() {
		if err != nil && err != io.EOF {
			b.span.SetStatus(codes.Error, "body read failed")
		}
		b.span.End()
	})
}
func (b *spanBody) Read(p []byte) (int, error) {
	n, err := b.ReadCloser.Read(p)
	if err != nil {
		b.finish(err)
	}
	return n, err
}
func (b *spanBody) Close() error { err := b.ReadCloser.Close(); b.finish(err); return err }

type RedisHook struct{ Runtime *Runtime }

func (h RedisHook) DialHook(next redis.DialHook) redis.DialHook {
	return func(ctx context.Context, network, addr string) (conn net.Conn, err error) {
		err = h.Runtime.Do(ctx, "redis.connect", func(ctx context.Context) error { conn, err = next(ctx, network, addr); return err })
		return
	}
}
func (h RedisHook) ProcessHook(next redis.ProcessHook) redis.ProcessHook {
	return func(ctx context.Context, cmd redis.Cmder) error {
		var result error
		_ = h.Runtime.Do(ctx, "redis.command", func(ctx context.Context) error {
			result = next(ctx, cmd)
			if result == redis.Nil {
				return nil
			}
			return result
		})
		return result
	}
}
func (h RedisHook) ProcessPipelineHook(next redis.ProcessPipelineHook) redis.ProcessPipelineHook {
	return func(ctx context.Context, cmds []redis.Cmder) error {
		return h.Runtime.Do(ctx, "redis.pipeline", func(ctx context.Context) error { return next(ctx, cmds) })
	}
}

// SQL records execution spans without SQL text or arguments. Query spans cover
// obtaining rows; iteration/Scan remain caller-owned. Use Do around transactions.
type SQL struct {
	DB      *sql.DB
	Runtime *Runtime
}

func (d SQL) ExecContext(ctx context.Context, query string, args ...any) (result sql.Result, err error) {
	err = d.Runtime.Do(ctx, "db.exec", func(ctx context.Context) error { result, err = d.DB.ExecContext(ctx, query, args...); return err })
	return
}
func (d SQL) QueryContext(ctx context.Context, query string, args ...any) (rows *sql.Rows, err error) {
	err = d.Runtime.Do(ctx, "db.query", func(ctx context.Context) error { rows, err = d.DB.QueryContext(ctx, query, args...); return err })
	return
}