package postgres

import (
	"encoding/json"
	"testing"
	"time"
)

func TestMergeEnvelopeJSON_MergesNestedPayloads(t *testing.T) {
	existing := json.RawMessage(`{
		"event":{"type":"span.start"},
		"attributes":{"name":"tool.call","span_kind":"tool"},
		"payload":{"input":{"query":"status"},"prompt_preview":"summarize"}
	}`)
	next := json.RawMessage(`{
		"event":{"type":"span.end"},
		"payload":{"result_preview":"ok","duration_ms":42}
	}`)

	merged := mergeEnvelopeJSON(existing, next)

	var got map[string]any
	if err := json.Unmarshal(merged, &got); err != nil {
		t.Fatalf("unmarshal merged payload: %v", err)
	}

	payload, ok := got["payload"].(map[string]any)
	if !ok {
		t.Fatal("expected merged payload object")
	}
	if _, ok := payload["input"].(map[string]any); !ok {
		t.Fatal("expected input from first event to be preserved")
	}
	if payload["result_preview"] != "ok" {
		t.Fatalf("expected result_preview to be merged, got %#v", payload["result_preview"])
	}
	if payload["duration_ms"] != float64(42) {
		t.Fatalf("expected duration_ms to be merged, got %#v", payload["duration_ms"])
	}

	attrs, ok := got["attributes"].(map[string]any)
	if !ok || attrs["name"] != "tool.call" {
		t.Fatalf("expected attributes to be preserved, got %#v", got["attributes"])
	}

	event, ok := got["event"].(map[string]any)
	if !ok || event["type"] != "span.end" {
		t.Fatalf("expected later event metadata to win, got %#v", got["event"])
	}
}

func TestMergeSpanEvent_PreservesStartPayloadDetails(t *testing.T) {
	existing := &SpanRow{
		Name:   "tool.call",
		Kind:   "tool",
		Status: "success",
		Payload: json.RawMessage(`{
			"attributes":{"name":"tool.call","span_kind":"tool"},
			"payload":{"input":{"command":"ls"}}
		}`),
	}

	mergeSpanEvent(existing, SpanRow{
		Status: "success",
		Payload: json.RawMessage(`{
			"payload":{"result_preview":"done","duration_ms":12}
		}`),
	})

	var got map[string]any
	if err := json.Unmarshal(existing.Payload, &got); err != nil {
		t.Fatalf("unmarshal merged span payload: %v", err)
	}

	payload := got["payload"].(map[string]any)
	if _, ok := payload["input"].(map[string]any); !ok {
		t.Fatal("expected input to remain after merge")
	}
	if payload["result_preview"] != "done" {
		t.Fatalf("expected result_preview after merge, got %#v", payload["result_preview"])
	}
}

func TestFindRunIndexForSpan_MatchesContainingRunWindow(t *testing.T) {
	start := time.Date(2026, 3, 23, 10, 0, 0, 0, time.UTC)
	run1End := start.Add(2 * time.Minute)
	run2Start := start.Add(3 * time.Minute)
	run2End := start.Add(7 * time.Minute)

	runs := []RunRow{
		{
			RunID:     "run-1",
			StartedAt: start,
			EndedAt:   &run1End,
		},
		{
			RunID:     "run-2",
			StartedAt: run2Start,
			EndedAt:   &run2End,
		},
	}

	idx := findRunIndexForSpan(runs, start.Add(4*time.Minute))
	if idx != 1 {
		t.Fatalf("expected span to attach to run-2, got index %d", idx)
	}
}

func TestFindRunIndexForSpan_MatchesOpenRun(t *testing.T) {
	start := time.Date(2026, 3, 23, 10, 0, 0, 0, time.UTC)
	run1End := start.Add(2 * time.Minute)
	run2Start := start.Add(3 * time.Minute)

	runs := []RunRow{
		{
			RunID:     "run-1",
			StartedAt: start,
			EndedAt:   &run1End,
		},
		{
			RunID:     "run-2",
			StartedAt: run2Start,
		},
	}

	idx := findRunIndexForSpan(runs, start.Add(5*time.Minute))
	if idx != 1 {
		t.Fatalf("expected span to attach to open run-2, got index %d", idx)
	}
}