package storage

import (
	"encoding/json"
	"strconv"
	"strings"
)

// spanFact is the slice of a span needed for trace summaries.
type spanFact struct {
	TraceID      string
	SpanID       string
	ParentSpanID string
	Service      string
	Name         string
	StartNs      uint64
	EndNs        uint64
	IsError      bool
}

type otlpValue struct {
	StringValue string      `json:"stringValue"`
	IntValue    json.Number `json:"intValue"`
}

type otlpAttr struct {
	Key   string    `json:"key"`
	Value otlpValue `json:"value"`
}

type otlpSpan struct {
	TraceID           string      `json:"traceId"`
	SpanID            string      `json:"spanId"`
	ParentSpanID      string      `json:"parentSpanId"`
	Name              string      `json:"name"`
	StartTimeUnixNano json.Number `json:"startTimeUnixNano"`
	EndTimeUnixNano   json.Number `json:"endTimeUnixNano"`
	Status            struct {
		Code int `json:"code"`
	} `json:"status"`
	Attributes []otlpAttr `json:"attributes"`
}

// httpStatusOf reads http.status_code (old) / http.response.status_code
// (new) as a number, tolerating string- or int-encoded OTLP values.
func httpStatusOf(attrs []otlpAttr) int {
	for _, a := range attrs {
		if a.Key != "http.status_code" && a.Key != "http.response.status_code" {
			continue
		}
		if n, err := a.Value.IntValue.Int64(); err == nil {
			return int(n)
		}
		if n, err := strconv.Atoi(a.Value.StringValue); err == nil {
			return n
		}
	}
	return 0
}

// isErrorSpan mirrors the UI's rule: OTLP error status or HTTP status >= 400.
func isErrorSpan(s otlpSpan) bool {
	return s.Status.Code == 2 || httpStatusOf(s.Attributes) >= 400
}

func parseUint(n json.Number) uint64 {
	v, err := strconv.ParseUint(n.String(), 10, 64)
	if err != nil {
		return 0
	}
	return v
}

// extractSpanFacts flattens an OTLP traces payload into per-span facts.
func extractSpanFacts(data []byte) []spanFact {
	var p struct {
		ResourceSpans []struct {
			Resource struct {
				Attributes []otlpAttr `json:"attributes"`
			} `json:"resource"`
			ScopeSpans []struct {
				Spans []otlpSpan `json:"spans"`
			} `json:"scopeSpans"`
		} `json:"resourceSpans"`
	}
	if err := json.Unmarshal(data, &p); err != nil {
		return nil
	}
	var out []spanFact
	for _, rs := range p.ResourceSpans {
		service := "unknown-service"
		for _, a := range rs.Resource.Attributes {
			if a.Key == "service.name" && a.Value.StringValue != "" {
				service = a.Value.StringValue
			}
		}
		for _, ss := range rs.ScopeSpans {
			for _, sp := range ss.Spans {
				if sp.TraceID == "" {
					continue
				}
				start := parseUint(sp.StartTimeUnixNano)
				end := parseUint(sp.EndTimeUnixNano)
				if end < start {
					end = start
				}
				out = append(out, spanFact{
					TraceID: sp.TraceID, SpanID: sp.SpanID, ParentSpanID: sp.ParentSpanID,
					Service: service, Name: sp.Name, StartNs: start, EndNs: end, IsError: isErrorSpan(sp),
				})
			}
		}
	}
	return out
}

// allTraceIDs lists the distinct trace ids in a payload (a single OTLP batch
// routinely carries spans/logs of many traces, so the scalar trace_id column
// alone cannot answer "which entries contain trace X").
func allTraceIDs(data []byte) []string {
	var p struct {
		ResourceSpans []struct {
			ScopeSpans []struct {
				Spans []struct {
					TraceID string `json:"traceId"`
				} `json:"spans"`
			} `json:"scopeSpans"`
		} `json:"resourceSpans"`
		ResourceLogs []struct {
			ScopeLogs []struct {
				LogRecords []struct {
					TraceID string `json:"traceId"`
				} `json:"logRecords"`
			} `json:"scopeLogs"`
		} `json:"resourceLogs"`
	}
	if err := json.Unmarshal(data, &p); err != nil {
		return nil
	}
	seen := map[string]struct{}{}
	var out []string
	add := func(id string) {
		if id == "" {
			return
		}
		if _, ok := seen[id]; !ok {
			seen[id] = struct{}{}
			out = append(out, id)
		}
	}
	for _, rs := range p.ResourceSpans {
		for _, ss := range rs.ScopeSpans {
			for _, sp := range ss.Spans {
				add(sp.TraceID)
			}
		}
	}
	for _, rl := range p.ResourceLogs {
		for _, sl := range rl.ScopeLogs {
			for _, rec := range sl.LogRecords {
				add(rec.TraceID)
			}
		}
	}
	return out
}

// filterNested walks raw (a JSON array) down path (object keys holding arrays
// at each level), keeping only leaf elements for which keep returns true and
// dropping any parent that ends up empty. Unknown fields are preserved
// verbatim. It returns the pruned array and the number of leaves kept.
func filterNested(raw json.RawMessage, path []string, keep func(json.RawMessage) bool) (json.RawMessage, int, error) {
	var items []json.RawMessage
	if err := json.Unmarshal(raw, &items); err != nil {
		return nil, 0, err
	}
	kept := 0
	out := make([]json.RawMessage, 0, len(items))
	for _, item := range items {
		if len(path) == 0 {
			if keep(item) {
				out = append(out, item)
				kept++
			}
			continue
		}
		var obj map[string]json.RawMessage
		if err := json.Unmarshal(item, &obj); err != nil {
			return nil, 0, err
		}
		child, ok := obj[path[0]]
		if !ok {
			continue
		}
		pruned, n, err := filterNested(child, path[1:], keep)
		if err != nil {
			return nil, 0, err
		}
		if n == 0 {
			continue
		}
		prunedJSON := pruned
		obj[path[0]] = prunedJSON
		rebuilt, err := json.Marshal(obj)
		if err != nil {
			return nil, 0, err
		}
		out = append(out, rebuilt)
		kept += n
	}
	b, err := json.Marshal(out)
	return b, kept, err
}

// TrimToTrace returns data reduced to the spans and log records of traceID,
// and false when nothing in it belongs to the trace. Payloads that are not
// OTLP traces/logs (metrics) report false.
func TrimToTrace(data []byte, traceID string) ([]byte, bool) {
	var top map[string]json.RawMessage
	if err := json.Unmarshal(data, &top); err != nil {
		return nil, false
	}
	match := func(item json.RawMessage) bool {
		var t struct {
			TraceID string `json:"traceId"`
		}
		return json.Unmarshal(item, &t) == nil && strings.EqualFold(t.TraceID, traceID)
	}
	total := 0
	for key, path := range map[string][]string{
		"resourceSpans": {"scopeSpans", "spans"},
		"resourceLogs":  {"scopeLogs", "logRecords"},
	} {
		raw, ok := top[key]
		if !ok {
			continue
		}
		pruned, n, err := filterNested(raw, path, match)
		if err != nil || n == 0 {
			delete(top, key)
			continue
		}
		top[key] = pruned
		total += n
	}
	if total == 0 {
		return nil, false
	}
	out, err := json.Marshal(top)
	return out, err == nil
}