package otelcol

import (
	"context"
	"encoding/json"
	"strings"

	"github.com/google/uuid"
	"github.com/rs/zerolog"
	collectortracev1 "go.opentelemetry.io/proto/otlp/collector/trace/v1"
	commonv1 "go.opentelemetry.io/proto/otlp/common/v1"
	tracev1 "go.opentelemetry.io/proto/otlp/trace/v1"

	"github.com/hatchet-dev/hatchet/pkg/analytics"
	"github.com/hatchet-dev/hatchet/pkg/repository"
	"github.com/hatchet-dev/hatchet/pkg/repository/sqlcv1"
)

const (
	// keep these in sync with attributes sent from the sdks
	AttrHatchetTaskRunID     = "hatchet.step_run_id"     // Task run external ID from SDK
	AttrHatchetWorkflowRunID = "hatchet.workflow_run_id" // Workflow run ID from SDK
	AttrHatchetRetryCount    = "hatchet.retry_count"     // Retry count from SDK
)

type otelCollectorImpl struct {
	collectortracev1.UnimplementedTraceServiceServer

	repo         repository.Repository
	l            *zerolog.Logger
	maxBatchSize int
	a            analytics.Analytics
}

func (oc *otelCollectorImpl) Export(ctx context.Context, req *collectortracev1.ExportTraceServiceRequest) (*collectortracev1.ExportTraceServiceResponse, error) {
	tenant, ok := ctx.Value("tenant").(*sqlcv1.Tenant)
	if !ok {
		oc.l.Warn().Msg("no tenant in context for trace export")
		return &collectortracev1.ExportTraceServiceResponse{}, nil
	}

	tenantId := tenant.ID

	olapRepo := oc.repo.OLAP()
	if olapRepo == nil {
		oc.l.Debug().Msg("OLAP repository not configured, discarding spans")
		return &collectortracev1.ExportTraceServiceResponse{}, nil
	}

	spans := oc.convertOTLPToSpanData(req.GetResourceSpans(), tenant.ID)

	if len(spans) == 0 {
		return &collectortracev1.ExportTraceServiceResponse{}, nil
	}

	oc.a.Count(ctx, analytics.OtelSpan, analytics.Create)

	var rejected int64
	if oc.maxBatchSize > 0 && len(spans) > oc.maxBatchSize {
		rejected = int64(len(spans) - oc.maxBatchSize)
		oc.l.Warn().Int("total", len(spans)).Int("max", oc.maxBatchSize).Int64("rejected", rejected).Msg("span batch exceeds max size, truncating")
		spans = spans[:oc.maxBatchSize]
	}

	err := olapRepo.CreateSpans(ctx, tenantId, &repository.CreateSpansOpts{
		TenantID: tenantId,
		Spans:    spans,
	})

	if err != nil {
		oc.l.Error().Err(err).Msg("failed to store spans")
		return &collectortracev1.ExportTraceServiceResponse{
			PartialSuccess: &collectortracev1.ExportTracePartialSuccess{
				RejectedSpans: int64(len(spans)) + rejected,
				ErrorMessage:  err.Error(),
			},
		}, nil
	}

	err = olapRepo.CreateSpanLookupTableEntries(ctx, tenantId, &repository.CreateSpansOpts{
		TenantID: tenantId,
		Spans:    spans,
	})

	if err != nil {
		oc.l.Error().Err(err).Msg("failed to create span lookup table entries")
		return &collectortracev1.ExportTraceServiceResponse{
			PartialSuccess: &collectortracev1.ExportTracePartialSuccess{
				RejectedSpans: int64(len(spans)) + rejected,
				ErrorMessage:  err.Error(),
			},
		}, nil
	}

	oc.l.Debug().Int("span_count", len(spans)).Str("tenant_id", tenantId.String()).Msg("stored spans")

	if rejected > 0 {
		return &collectortracev1.ExportTraceServiceResponse{
			PartialSuccess: &collectortracev1.ExportTracePartialSuccess{
				RejectedSpans: rejected,
				ErrorMessage:  "batch size exceeded maximum limit",
			},
		}, nil
	}

	return &collectortracev1.ExportTraceServiceResponse{}, nil
}

func (oc *otelCollectorImpl) convertOTLPToSpanData(resourceSpans []*tracev1.ResourceSpans, tenantID uuid.UUID) []*repository.SpanData {
	var spans []*repository.SpanData

	for _, rs := range resourceSpans {
		resourceAttrs := oc.serializeAttributes(rs.GetResource().GetAttributes())

		for _, ss := range rs.GetScopeSpans() {
			scopeName := ss.GetScope().GetName()

			for _, span := range ss.GetSpans() {
				spanData := &repository.SpanData{
					TraceID:              span.GetTraceId(),
					SpanID:               span.GetSpanId(),
					ParentSpanID:         span.GetParentSpanId(),
					Name:                 span.GetName(),
					Kind:                 span.GetKind(),
					StartTimeUnixNano:    span.GetStartTimeUnixNano(),
					EndTimeUnixNano:      span.GetEndTimeUnixNano(),
					StatusCode:           span.GetStatus().GetCode(),
					StatusMessage:        span.GetStatus().GetMessage(),
					Attributes:           oc.serializeAttributes(span.GetAttributes()),
					Events:               oc.serializeEvents(span.GetEvents()),
					Links:                oc.serializeLinks(span.GetLinks()),
					ResourceAttributes:   resourceAttrs,
					TenantID:             tenantID,
					InstrumentationScope: scopeName,
				}

				oc.extractHatchetCorrelation(span.GetAttributes(), spanData)

				spans = append(spans, spanData)
			}
		}
	}

	return spans
}

func (oc *otelCollectorImpl) extractHatchetCorrelation(attrs []*commonv1.KeyValue, spanData *repository.SpanData) {
	for _, attr := range attrs {
		switch attr.GetKey() {
		case AttrHatchetTaskRunID:
			if strVal := attr.GetValue().GetStringValue(); strVal != "" {
				uuid := uuid.MustParse(strVal)
				spanData.TaskRunExternalID = &uuid
			}
		case AttrHatchetWorkflowRunID:
			if strVal := attr.GetValue().GetStringValue(); strVal != "" {
				uuid := uuid.MustParse(strVal)
				spanData.WorkflowRunID = &uuid
			}
		case AttrHatchetRetryCount:
			if intVal := attr.GetValue().GetIntValue(); intVal > 0 {
				spanData.RetryCount = int32(intVal) //nolint:gosec
			}
		}
	}

	// Only re-parent the top-level SDK step_run consumer span under the
	// engine's deterministic workflow_run root. Other SDK spans (user spans,
	// child-run producers) keep their natural parent so the hierarchy is
	// preserved in both Hatchet's OLAP view and external OTel backends.
	if spanData.WorkflowRunID != nil && len(spanData.ParentSpanID) > 0 &&
		strings.HasPrefix(spanData.Name, "hatchet.start_step_run") {
		spanData.ParentSpanID = repository.DeriveWorkflowRunSpanID(*spanData.WorkflowRunID)
	}
}

func (oc *otelCollectorImpl) serializeAttributes(attrs []*commonv1.KeyValue) []byte {
	if len(attrs) == 0 {
		return nil
	}

	attrMap := make(map[string]any, len(attrs))
	for _, kv := range attrs {
		attrMap[kv.GetKey()] = oc.anyValueToInterface(kv.GetValue())
	}

	data, err := json.Marshal(attrMap)
	if err != nil {
		oc.l.Warn().Err(err).Msg("failed to serialize attributes")
		return nil
	}

	return data
}

func (oc *otelCollectorImpl) serializeEvents(events []*tracev1.Span_Event) []byte {
	if len(events) == 0 {
		return nil
	}

	eventList := make([]map[string]any, 0, len(events))
	for _, event := range events {
		eventMap := map[string]any{
			"name":                     event.GetName(),
			"time_unix_nano":           event.GetTimeUnixNano(),
			"dropped_attributes_count": event.GetDroppedAttributesCount(),
		}

		if len(event.GetAttributes()) > 0 {
			attrMap := make(map[string]any, len(event.GetAttributes()))
			for _, kv := range event.GetAttributes() {
				attrMap[kv.GetKey()] = oc.anyValueToInterface(kv.GetValue())
			}
			eventMap["attributes"] = attrMap
		}

		eventList = append(eventList, eventMap)
	}

	data, err := json.Marshal(eventList)
	if err != nil {
		oc.l.Warn().Err(err).Msg("failed to serialize events")
		return nil
	}

	return data
}

func (oc *otelCollectorImpl) serializeLinks(links []*tracev1.Span_Link) []byte {
	if len(links) == 0 {
		return nil
	}

	linkList := make([]map[string]any, 0, len(links))
	for _, link := range links {
		linkMap := map[string]any{
			"trace_id":                 link.GetTraceId(),
			"span_id":                  link.GetSpanId(),
			"trace_state":              link.GetTraceState(),
			"dropped_attributes_count": link.GetDroppedAttributesCount(),
		}

		if len(link.GetAttributes()) > 0 {
			attrMap := make(map[string]any, len(link.GetAttributes()))
			for _, kv := range link.GetAttributes() {
				attrMap[kv.GetKey()] = oc.anyValueToInterface(kv.GetValue())
			}
			linkMap["attributes"] = attrMap
		}

		linkList = append(linkList, linkMap)
	}

	data, err := json.Marshal(linkList)
	if err != nil {
		oc.l.Warn().Err(err).Msg("failed to serialize links")
		return nil
	}

	return data
}

func (oc *otelCollectorImpl) anyValueToInterface(v *commonv1.AnyValue) interface{} {
	if v == nil {
		return nil
	}

	switch val := v.GetValue().(type) {
	case *commonv1.AnyValue_StringValue:
		return val.StringValue
	case *commonv1.AnyValue_BoolValue:
		return val.BoolValue
	case *commonv1.AnyValue_IntValue:
		return val.IntValue
	case *commonv1.AnyValue_DoubleValue:
		return val.DoubleValue
	case *commonv1.AnyValue_ArrayValue:
		arr := make([]any, 0, len(val.ArrayValue.GetValues()))
		for _, item := range val.ArrayValue.GetValues() {
			arr = append(arr, oc.anyValueToInterface(item))
		}
		return arr
	case *commonv1.AnyValue_KvlistValue:
		m := make(map[string]any, len(val.KvlistValue.GetValues()))
		for _, kv := range val.KvlistValue.GetValues() {
			m[kv.GetKey()] = oc.anyValueToInterface(kv.GetValue())
		}
		return m
	case *commonv1.AnyValue_BytesValue:
		return val.BytesValue
	default:
		return nil
	}
}
