package storage

import (
	"encoding/json"
	"testing"
)

const mixedPayload = `{
 "resourceSpans":[
  {"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"api"}}]},"schemaUrl":"keep-me",
   "scopeSpans":[{"scope":{"name":"lib"},"spans":[
     {"traceId":"aaa","spanId":"s1","parentSpanId":"","name":"GET /","startTimeUnixNano":"100","endTimeUnixNano":"300","status":{"code":0}},
     {"traceId":"bbb","spanId":"s2","parentSpanId":"","name":"other","startTimeUnixNano":"150","endTimeUnixNano":"160","status":{"code":2}},
     {"traceId":"aaa","spanId":"s3","parentSpanId":"s1","name":"db","startTimeUnixNano":"120","endTimeUnixNano":"200","status":{"code":0},
      "attributes":[{"key":"http.status_code","value":{"intValue":"503"}}]}
   ]}]},
  {"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"only-bbb"}}]},
   "scopeSpans":[{"spans":[{"traceId":"bbb","spanId":"s4","name":"x","startTimeUnixNano":"1","endTimeUnixNano":"2"}]}]}
 ],
 "resourceLogs":[{"scopeLogs":[{"logRecords":[{"traceId":"aaa","body":{"stringValue":"hi"}},{"traceId":"ccc"}]}]}]
}`

func TestTrimToTraceKeepsOnlyRequestedTrace(t *testing.T) {
	out, ok := TrimToTrace([]byte(mixedPayload), "aaa")
	if !ok {
		t.Fatal("expected trace aaa to be found")
	}
	var p struct {
		ResourceSpans []struct {
			SchemaURL  string `json:"schemaUrl"`
			ScopeSpans []struct {
				Scope struct{ Name string } `json:"scope"`
				Spans []struct {
					SpanID string `json:"spanId"`
				} `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(out, &p); err != nil {
		t.Fatal(err)
	}
	if len(p.ResourceSpans) != 1 {
		t.Fatalf("resource with only other traces should be dropped, got %d resources", len(p.ResourceSpans))
	}
	rs := p.ResourceSpans[0]
	if rs.SchemaURL != "keep-me" || rs.ScopeSpans[0].Scope.Name != "lib" {
		t.Errorf("unknown/sibling fields must be preserved: %+v", rs)
	}
	got := []string{}
	for _, s := range rs.ScopeSpans[0].Spans {
		got = append(got, s.SpanID)
	}
	if len(got) != 2 || got[0] != "s1" || got[1] != "s3" {
		t.Errorf("spans = %v, want [s1 s3]", got)
	}
	if len(p.ResourceLogs) != 1 || len(p.ResourceLogs[0].ScopeLogs[0].LogRecords) != 1 {
		t.Errorf("logs not trimmed to aaa: %+v", p.ResourceLogs)
	}
}

func TestTrimToTraceMissAndNonTrace(t *testing.T) {
	if _, ok := TrimToTrace([]byte(mixedPayload), "zzz"); ok {
		t.Error("unknown trace should report false")
	}
	if _, ok := TrimToTrace([]byte(`{"resourceMetrics":[]}`), "aaa"); ok {
		t.Error("metrics payload should report false")
	}
	if _, ok := TrimToTrace([]byte(`not json`), "aaa"); ok {
		t.Error("invalid json should report false")
	}
}

func TestExtractSpanFacts(t *testing.T) {
	facts := extractSpanFacts([]byte(mixedPayload))
	if len(facts) != 4 {
		t.Fatalf("got %d facts, want 4", len(facts))
	}
	by := map[string]spanFact{}
	for _, f := range facts {
		by[f.SpanID] = f
	}
	if f := by["s1"]; f.Service != "api" || f.StartNs != 100 || f.EndNs != 300 || f.IsError {
		t.Errorf("s1 = %+v", f)
	}
	if !by["s2"].IsError {
		t.Error("status code 2 must be an error")
	}
	if !by["s3"].IsError {
		t.Error("http.status_code 503 (intValue) must be an error")
	}
	if by["s4"].Service != "only-bbb" {
		t.Errorf("s4 service = %q", by["s4"].Service)
	}
}

func TestAllTraceIDsDistinctAcrossSpansAndLogs(t *testing.T) {
	ids := allTraceIDs([]byte(mixedPayload))
	want := map[string]bool{"aaa": true, "bbb": true, "ccc": true}
	if len(ids) != len(want) {
		t.Fatalf("ids = %v", ids)
	}
	for _, id := range ids {
		if !want[id] {
			t.Errorf("unexpected id %q", id)
		}
	}
}