package spans

import (
	"bytes"
	"context"
	"fmt"
	"hash/crc32"
	"net"
	"strconv"
	"strings"
	"sync/atomic"

	"github.com/ydb-platform/ydb-go-sdk/v3/internal/kv"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/repeater"
	"github.com/ydb-platform/ydb-go-sdk/v3/meta"
	"github.com/ydb-platform/ydb-go-sdk/v3/pkg/xstring"
	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
)

func traceparent(traceID string, spanID string) string {
	b := xstring.Buffer()
	defer b.Free()

	b.WriteString("00-")
	b.WriteString(traceID)
	b.WriteByte('-')
	b.WriteString(spanID)
	b.WriteString("-01")

	return b.String()
}

// driver makes driver with publishing traces
func driver(adapter Adapter) trace.Driver { //nolint:gocyclo,funlen
	return trace.Driver{
		OnRepeaterWakeUp: func(info trace.DriverRepeaterWakeUpStartInfo) func(trace.DriverRepeaterWakeUpDoneInfo) {
			if adapter.Details()&trace.DriverRepeaterEvents == 0 {
				return nil
			}
			start := childSpanWithReplaceCtx(
				adapter,
				info.Context,
				safeCall(info.Call),
			)

			return func(info trace.DriverRepeaterWakeUpDoneInfo) {
				finish(
					start,
					info.Error,
				)
			}
		},
		OnConnDial: func(info trace.DriverConnDialStartInfo) func(trace.DriverConnDialDoneInfo) {
			// We don't create a child span for the gRPC dial — it would
			// pollute the user-facing span tree. We only inject the W3C
			// `traceparent` header so server-side YDB spans link to the
			// caller's trace.
			if adapter.Details()&trace.DriverConnEvents == 0 {
				return nil
			}
			if traceID, valid := adapter.SpanFromContext(safeContextPtr(info.Context)).TraceID(); valid {
				withContextPtr(info.Context, func(c context.Context) context.Context {
					return meta.WithTraceParent(c, traceID)
				})
			}

			return nil
		},
		OnConnInvoke: func(info trace.DriverConnInvokeStartInfo) func(trace.DriverConnInvokeDoneInfo) {
			// The QueryService unary RPC. Instead of a child span, attach
			// network.peer.address / network.peer.port to the surrounding
			// client span (e.g. ydb.CreateSession, ydb.Commit) and inject
			// traceparent so the server sees this RPC as a child of that
			// span.
			if adapter.Details()&trace.DriverConnEvents == 0 {
				return nil
			}
			parent := adapter.SpanFromContext(safeContextPtr(info.Context))
			annotateNetworkPeer(parent, info.Endpoint)
			// Also propagate the same peer/node attributes up to the
			// surrounding top-level CLIENT span (e.g. the outer
			// ydb.ExecuteQuery span emitted by OnExec/OnQuery/...) so that
			// every ydb.* CLIENT span carries ydb.node.id / ydb.node.dc
			// alongside network.peer.* — see the documentation block at
			// the top of spans/query.go.
			if clientSpan := clientSpanFromContext(safeContextPtr(info.Context)); clientSpan != nil {
				annotateNetworkPeer(clientSpan, info.Endpoint)
			}
			if id, valid := parent.ID(); valid {
				if traceID, valid := parent.TraceID(); valid {
					withContextPtr(info.Context, func(c context.Context) context.Context {
						return meta.WithTraceParent(c, traceparent(traceID, id))
					})
				}
			}

			return func(info trace.DriverConnInvokeDoneInfo) {
				if info.Error != nil {
					logError(parent, info.Error)
				} else if len(info.Issues) > 0 {
					issues := make([]string, len(info.Issues))
					for i, issue := range info.Issues {
						issues[i] = fmt.Sprintf("%+v", issue)
					}
					parent.Log("conn.Invoke issues", kv.Strings("issues", issues))
				}
			}
		},
		OnConnNewStream: func(info trace.DriverConnNewStreamStartInfo) func(trace.DriverConnNewStreamDoneInfo) {
			// The QueryService streaming RPC. Same treatment as Invoke:
			// no child span, just augment the parent client span.
			if adapter.Details()&trace.DriverConnStreamEvents == 0 {
				return nil
			}
			withContextPtr(info.Context, withGrpcStreamMsgCounters)
			parent := adapter.SpanFromContext(safeContextPtr(info.Context))
			annotateNetworkPeer(parent, info.Endpoint)
			if clientSpan := clientSpanFromContext(safeContextPtr(info.Context)); clientSpan != nil {
				annotateNetworkPeer(clientSpan, info.Endpoint)
			}
			if id, valid := parent.ID(); valid {
				if traceID, valid := parent.TraceID(); valid {
					withContextPtr(info.Context, func(c context.Context) context.Context {
						return meta.WithTraceParent(c, traceparent(traceID, id))
					})
				}
			}

			return func(info trace.DriverConnNewStreamDoneInfo) {
				if info.Error != nil {
					logError(parent, info.Error)
				}
			}
		},
		OnConnStreamRecvMsg: func(info trace.DriverConnStreamRecvMsgStartInfo) func(trace.DriverConnStreamRecvMsgDoneInfo) {
			if adapter.Details()&trace.DriverConnStreamEvents == 0 {
				return nil
			}
			start := adapter.SpanFromContext(safeContextPtr(info.Context))
			counters := grpcStreamMsgCountersFromContext(safeContextPtr(info.Context))
			if counters != nil {
				counters.updateReceivedMessages()
			}

			return func(info trace.DriverConnStreamRecvMsgDoneInfo) {
				if skipEOF(info.Error) != nil {
					logError(start, info.Error)
				}
			}
		},
		OnConnStreamSendMsg: func(info trace.DriverConnStreamSendMsgStartInfo) func(trace.DriverConnStreamSendMsgDoneInfo) {
			if adapter.Details()&trace.DriverConnStreamEvents == 0 {
				return nil
			}
			start := adapter.SpanFromContext(safeContextPtr(info.Context))
			counters := grpcStreamMsgCountersFromContext(safeContextPtr(info.Context))
			if counters != nil {
				counters.updateSentMessages()
			}

			return func(info trace.DriverConnStreamSendMsgDoneInfo) {
				if info.Error != nil {
					logError(start, info.Error)
				}
			}
		},
		OnConnStreamCloseSend: func(info trace.DriverConnStreamCloseSendStartInfo) func(
			trace.DriverConnStreamCloseSendDoneInfo,
		) {
			if adapter.Details()&trace.DriverConnStreamEvents == 0 {
				return nil
			}

			var (
				call  = info.Call
				start = adapter.SpanFromContext(safeContextPtr(info.Context))
			)

			return func(info trace.DriverConnStreamCloseSendDoneInfo) {
				if info.Error != nil {
					start.Log(safeCall(call), kv.Error(info.Error))
				}
			}
		},
		OnConnStreamFinish: func(info trace.DriverConnStreamFinishInfo) {
			if adapter.Details()&trace.DriverConnStreamEvents == 0 {
				return
			}

			var (
				call     = info.Call
				ctx      = safeContext(info.Context)
				start    = adapter.SpanFromContext(ctx)
				counters = grpcStreamMsgCountersFromContext(ctx)
			)

			attributes := make([]kv.KeyValue, 0, 2)
			if counters != nil {
				attributes = append(attributes,
					kv.Int64("received_messages", counters.receivedMessages()),
					kv.Int64("sent_messages", counters.sentMessages()),
				)
			}

			if info.Error != nil {
				attributes = append(attributes, kv.Error(info.Error))
			}

			start.Log(safeCall(call), attributes...)
		},
		OnConnPark: func(info trace.DriverConnParkStartInfo) func(trace.DriverConnParkDoneInfo) {
			if adapter.Details()&trace.DriverConnEvents == 0 {
				return nil
			}

			var (
				call  = info.Call
				start = adapter.SpanFromContext(safeContextPtr(info.Context))
			)

			return func(info trace.DriverConnParkDoneInfo) {
				if info.Error != nil {
					start.Log(safeCall(call), kv.Error(info.Error))
				} else {
					start.Log(safeCall(call))
				}
			}
		},
		OnConnClose: func(info trace.DriverConnCloseStartInfo) func(trace.DriverConnCloseDoneInfo) {
			if adapter.Details()&trace.DriverConnEvents == 0 {
				return nil
			}

			var (
				call  = info.Call
				start = adapter.SpanFromContext(safeContextPtr(info.Context))
			)

			return func(info trace.DriverConnCloseDoneInfo) {
				if info.Error != nil {
					start.Log(safeCall(call), kv.Error(info.Error))
				} else {
					start.Log(safeCall(call))
				}
			}
		},
		OnConnBan: func(info trace.DriverConnBanStartInfo) func(trace.DriverConnBanDoneInfo) {
			if adapter.Details()&trace.DriverConnEvents == 0 {
				return nil
			}
			s := adapter.SpanFromContext(safeContextPtr(info.Context))
			s.Log(safeCall(info.Call),
				kv.String("cause", safeError(info.Cause)),
			)

			return nil
		},
		OnConnStateChange: func(info trace.DriverConnStateChangeStartInfo) func(trace.DriverConnStateChangeDoneInfo) {
			if adapter.Details()&trace.DriverConnEvents == 0 {
				return nil
			}
			s := adapter.SpanFromContext(safeContextPtr(info.Context))
			oldState := safeStringer(info.State)
			functionID := safeCall(info.Call)

			return func(info trace.DriverConnStateChangeDoneInfo) {
				s.Log(functionID,
					kv.String("old state", oldState),
					kv.String("new state", safeStringer(info.State)),
				)
			}
		},
		OnBalancerInit: func(info trace.DriverBalancerInitStartInfo) func(trace.DriverBalancerInitDoneInfo) {
			if adapter.Details()&trace.DriverBalancerEvents == 0 {
				return nil
			}
			start := childSpanWithReplaceCtx(
				adapter,
				info.Context,
				safeCall(info.Call),
				kv.String("name", info.Name),
			)

			return func(info trace.DriverBalancerInitDoneInfo) {
				finish(start, info.Error)
			}
		},
		OnBalancerClusterDiscoveryAttempt: func(info trace.DriverBalancerClusterDiscoveryAttemptStartInfo) func(
			trace.DriverBalancerClusterDiscoveryAttemptDoneInfo,
		) {
			if adapter.Details()&trace.DriverBalancerEvents == 0 {
				return nil
			}
			// Driver bootstrap fires the trace event from two places:
			//   - clusterDiscoveryAttemptWithDial (outer, also responsible
			//     for opening the gRPC connection)
			//   - clusterDiscoveryAttempt (inner, called from the outer one)
			// We only want a single ydb.Driver.Initialize span, and only on
			// the very first init — periodic discovery refreshes triggered
			// by the repeater come with EventTick / EventForce and must not
			// produce a span.
			isOuter := strings.HasSuffix(safeCall(info.Call), "clusterDiscoveryAttemptWithDial")
			isInit := repeater.EventType(safeContextPtr(info.Context)) == repeater.EventInit
			if !isOuter || !isInit {
				return nil
			}
			start := childSpanWithReplaceCtx(
				adapter,
				info.Context,
				SpanNameDriverInitialize,
				kv.String("address", info.Address),
			)

			return func(info trace.DriverBalancerClusterDiscoveryAttemptDoneInfo) {
				finish(start, info.Error)
			}
		},
		OnBalancerUpdate: func(info trace.DriverBalancerUpdateStartInfo) func(trace.DriverBalancerUpdateDoneInfo) {
			if adapter.Details()&trace.DriverBalancerEvents == 0 {
				return nil
			}
			start := childSpanWithReplaceCtx(adapter, info.Context,
				safeCall(info.Call),
				kv.String("database", info.Database),
				kv.Bool("need_local_dc", info.NeedLocalDC),
			)

			return func(info trace.DriverBalancerUpdateDoneInfo) {
				endpoints := make([]string, 0, len(info.Endpoints))
				added := make([]string, 0, len(info.Added))
				dropped := make([]string, 0, len(info.Dropped))
				for _, e := range info.Endpoints {
					if isNil(e) {
						continue
					}
					endpoints = append(endpoints, e.String())
				}
				for _, e := range info.Added {
					if isNil(e) {
						continue
					}
					added = append(added, e.String())
				}
				for _, e := range info.Dropped {
					if isNil(e) {
						continue
					}
					dropped = append(dropped, e.String())
				}
				start.Log(fmt.Sprintf("endpoints=%v", endpoints))
				start.Log(fmt.Sprintf("added=%v", added))
				start.Log(fmt.Sprintf("dropped=%v", dropped))
				start.End(
					kv.String("local_dc", info.LocalDC),
				)
			}
		},
		OnBalancerChooseEndpoint: func(
			info trace.DriverBalancerChooseEndpointStartInfo,
		) func(
			trace.DriverBalancerChooseEndpointDoneInfo,
		) {
			if adapter.Details()&trace.DriverBalancerEvents == 0 {
				return nil
			}
			parent := adapter.SpanFromContext(safeContextPtr(info.Context))
			functionID := safeCall(info.Call)

			return func(info trace.DriverBalancerChooseEndpointDoneInfo) {
				if info.Error != nil {
					parent.Error(info.Error)
				} else {
					parent.Log(functionID,
						kv.String("address", safeAddress(info.Endpoint)),
						kv.String("nodeID", safeNodeID(info.Endpoint)),
					)
				}
			}
		},
		OnGetCredentials: func(info trace.DriverGetCredentialsStartInfo) func(trace.DriverGetCredentialsDoneInfo) {
			if adapter.Details()&trace.DriverCredentialsEvents == 0 {
				return nil
			}
			parent := adapter.SpanFromContext(safeContextPtr(info.Context))
			functionID := safeCall(info.Call)

			return func(info trace.DriverGetCredentialsDoneInfo) {
				if info.Error != nil {
					parent.Error(info.Error)
				} else {
					var mask bytes.Buffer
					if len(info.Token) > 16 {
						mask.WriteString(info.Token[:4])
						mask.WriteString("****")
						mask.WriteString(info.Token[len(info.Token)-4:])
					} else {
						mask.WriteString("****")
					}
					fmt.Fprintf(&mask, "(CRC-32c: %08X)", crc32.Checksum([]byte(info.Token), crc32.IEEETable))
					parent.Log(functionID,
						kv.String("token", mask.String()),
					)
				}
			}
		},
		OnInit: func(info trace.DriverInitStartInfo) func(trace.DriverInitDoneInfo) {
			if adapter.Details()&trace.DriverEvents == 0 {
				return nil
			}
			start := childSpanWithReplaceCtx(
				adapter,
				info.Context,
				safeCall(info.Call),
				kv.String("endpoint", info.Endpoint),
				kv.String("database", info.Database),
				kv.Bool("secure", info.Secure),
			)

			return func(info trace.DriverInitDoneInfo) {
				finish(start, info.Error)
			}
		},
		OnClose: func(info trace.DriverCloseStartInfo) func(trace.DriverCloseDoneInfo) {
			if adapter.Details()&trace.DriverEvents == 0 {
				return nil
			}
			start := childSpanWithReplaceCtx(
				adapter,
				info.Context,
				safeCall(info.Call),
			)

			return func(info trace.DriverCloseDoneInfo) {
				finish(start, info.Error)
			}
		},
		OnPoolNew: func(info trace.DriverConnPoolNewStartInfo) func(trace.DriverConnPoolNewDoneInfo) {
			if adapter.Details()&trace.DriverEvents == 0 {
				return nil
			}
			start := childSpanWithReplaceCtx(
				adapter,
				info.Context,
				safeCall(info.Call),
			)

			return func(info trace.DriverConnPoolNewDoneInfo) {
				start.End()
			}
		},
	}
}

// annotateNetworkPeer attaches network.peer.address / network.peer.port to
// the given span based on the actual endpoint chosen for the RPC. The
// address comes in `host:port` form (e.g. "8e5d46f58bc6:2136"); we split
// it so the span has both attributes per OTel semantic conventions.
//
// Additionally, endpoint metadata is attached as YDB-specific attributes.
func annotateNetworkPeer(s Span, endpoint trace.EndpointInfo) {
	if isNil(endpoint) {
		return
	}
	addr := endpoint.Address()
	if addr == "" {
		return
	}
	attrs := make([]KeyValue, 0, 4)
	host, portStr, err := net.SplitHostPort(addr)
	if err != nil {
		attrs = append(attrs, kv.String(AttrNetworkPeerAddress, addr))
	} else {
		attrs = append(attrs, kv.String(AttrNetworkPeerAddress, host))
		if port, err := strconv.Atoi(portStr); err == nil {
			attrs = append(attrs, kv.Int(AttrNetworkPeerPort, port))
		}
	}
	attrs = append(attrs,
		kv.String(AttrYDBNodeDC, endpoint.Location()),
		kv.Int64(AttrYDBNodeID, int64(endpoint.NodeID())),
	)
	setSpanAttributes(s, attrs...)
}

type spanAttributesSetter interface {
	SetAttributes(attributes ...KeyValue)
}

func setSpanAttributes(s Span, attrs ...KeyValue) {
	if setter, ok := s.(spanAttributesSetter); ok {
		setter.SetAttributes(attrs...)
	}
}

type grpcStreamMsgCountersKey struct{}

func withGrpcStreamMsgCounters(ctx context.Context) context.Context {
	return context.WithValue(ctx, grpcStreamMsgCountersKey{}, &grpcStreamMsgCounters{})
}

func grpcStreamMsgCountersFromContext(ctx context.Context) *grpcStreamMsgCounters {
	value, ok := ctx.Value(grpcStreamMsgCountersKey{}).(*grpcStreamMsgCounters)
	if !ok {
		return nil
	}

	return value
}

type grpcStreamMsgCounters struct {
	sent     atomic.Int64
	received atomic.Int64
}

func (c *grpcStreamMsgCounters) updateSentMessages() {
	c.sent.Add(1)
}

func (c *grpcStreamMsgCounters) updateReceivedMessages() {
	c.received.Add(1)
}

func (c *grpcStreamMsgCounters) sentMessages() int64 {
	return c.sent.Load()
}

func (c *grpcStreamMsgCounters) receivedMessages() int64 {
	return c.received.Load()
}
