package source

import (
	"bytes"
	"encoding/json"
	"math"
	"reflect"
	"strings"
	"testing"
)

type typedMetadataMap map[string]any
type typedMetadataSlice []typedMetadataMap
type typedMetadataArray [2]any

func TestCanonicalizeChunkAnnotations(t *testing.T) {
	original := ChunkAnnotations{
		"domain/items": json.RawMessage(` { "z": [3, 2, 1], "a": 1.0 } `),
	}
	canonical, err := CanonicalizeChunkAnnotations(original)
	if err != nil {
		t.Fatalf("CanonicalizeChunkAnnotations() error = %v, want nil", err)
	}
	if got, want := string(canonical["domain/items"]), `{"a":1.0,"z":[3,2,1]}`; got != want {
		t.Fatalf("canonical annotation = %q, want %q", got, want)
	}
	original["domain/items"][0] = '['
	if got := string(canonical["domain/items"]); got != `{"a":1.0,"z":[3,2,1]}` {
		t.Fatalf("canonical annotation changed after input mutation: %q", got)
	}
	canonical["domain/items"][0] = '['
	if original["domain/items"][0] == '[' && bytes.Equal(original["domain/items"], canonical["domain/items"]) {
		t.Fatal("input and canonical annotation share value storage")
	}
}

func TestCanonicalizeChunkAnnotationsRejectsInvalidValues(t *testing.T) {
	tests := []struct {
		name        string
		annotations ChunkAnnotations
		want        string
	}{
		{name: "blank namespace", annotations: ChunkAnnotations{" \t": json.RawMessage(`true`)}, want: "namespace must not be empty"},
		{name: "untrimmed namespace", annotations: ChunkAnnotations{" items ": json.RawMessage(`true`)}, want: "leading or trailing whitespace"},
		{name: "invalid JSON", annotations: ChunkAnnotations{"items": json.RawMessage(`{"x":`)}, want: "valid JSON"},
		{name: "trailing JSON", annotations: ChunkAnnotations{"items": json.RawMessage(`true false`)}, want: "exactly one JSON value"},
		{name: "non-finite number", annotations: ChunkAnnotations{"items": json.RawMessage(`NaN`)}, want: "valid JSON"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			_, err := CanonicalizeChunkAnnotations(tt.annotations)
			if err == nil || !strings.Contains(err.Error(), tt.want) {
				t.Fatalf("CanonicalizeChunkAnnotations() error = %v, want containing %q", err, tt.want)
			}
		})
	}
}

func TestValidateChunkAnnotationsRequiresCanonicalJSON(t *testing.T) {
	if err := ValidateChunkAnnotations(ChunkAnnotations{"items": json.RawMessage(` {"b":2,"a":1}`)}); err == nil || !strings.Contains(err.Error(), "canonical JSON") {
		t.Fatalf("ValidateChunkAnnotations() error = %v, want canonical JSON error", err)
	}
	if err := ValidateChunkAnnotations(ChunkAnnotations{"items": json.RawMessage(`{"a":1,"b":2}`)}); err != nil {
		t.Fatalf("ValidateChunkAnnotations(canonical) error = %v, want nil", err)
	}
}

func TestCloneChunkPlanDoesNotShareAnnotationBytes(t *testing.T) {
	plan := validChunkPlan(planDocument())
	cloned := CloneChunkPlan(plan)
	cloned.Annotations["plan"][0] = '['
	cloned.Ranges[0].Annotations["range"][0] = '['
	if string(plan.Annotations["plan"]) != `{"value":1}` || string(plan.Ranges[0].Annotations["range"]) != `{"value":2}` {
		t.Fatal("CloneChunkPlan() shares annotation value storage")
	}
}

func TestValidateChunkPlanRanges(t *testing.T) {
	doc := planDocument()
	tests := []struct {
		name   string
		mutate func(*ChunkPlan)
		want   string
	}{
		{name: "source mismatch", mutate: func(plan *ChunkPlan) { plan.SourceDigest = "sha256:other" }, want: "does not match"},
		{name: "missing ranges", mutate: func(plan *ChunkPlan) { plan.Ranges = nil }, want: "ranges must not be empty"},
		{name: "missing start", mutate: func(plan *ChunkPlan) { plan.Ranges[0].StartUnitID = 99 }, want: "start_unit_id 99 was not found"},
		{name: "missing end", mutate: func(plan *ChunkPlan) { plan.Ranges[0].EndUnitID = 99 }, want: "end_unit_id 99 was not found"},
		{name: "backward range", mutate: func(plan *ChunkPlan) { plan.Ranges[0] = ChunkRange{StartUnitID: 30, EndUnitID: 10} }, want: "appears after end_unit_id"},
		{name: "duplicate start", mutate: func(plan *ChunkPlan) { plan.Ranges[1].StartUnitID = plan.Ranges[0].StartUnitID }, want: "does not appear after"},
		{name: "backward starts", mutate: func(plan *ChunkPlan) {
			plan.Ranges = []ChunkRange{{StartUnitID: 30, EndUnitID: 50}, {StartUnitID: 20, EndUnitID: 40}}
		}, want: "does not appear after"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			plan := validChunkPlan(doc)
			tt.mutate(&plan)
			err := ValidateChunkPlan(doc, plan)
			if err == nil || !strings.Contains(err.Error(), tt.want) {
				t.Fatalf("ValidateChunkPlan() error = %v, want containing %q", err, tt.want)
			}
		})
	}

	for name, ranges := range map[string][]ChunkRange{
		"gap":     {{StartUnitID: 10, EndUnitID: 20}, {StartUnitID: 40, EndUnitID: 50}},
		"overlap": {{StartUnitID: 10, EndUnitID: 30}, {StartUnitID: 20, EndUnitID: 50}},
	} {
		t.Run(name, func(t *testing.T) {
			plan := validChunkPlan(doc)
			plan.Ranges = ranges
			if err := ValidateChunkPlan(doc, plan); err != nil {
				t.Fatalf("ValidateChunkPlan() error = %v, want nil", err)
			}
		})
	}
}

func TestDigestChunkPlanIsStableAndCoversLogicalPlan(t *testing.T) {
	doc := planDocument()
	plan := validChunkPlan(doc)
	first, err := DigestChunkPlan(plan)
	if err != nil {
		t.Fatalf("DigestChunkPlan() error = %v, want nil", err)
	}
	reformatted := CloneChunkPlan(plan)
	reformatted.Annotations["plan"] = json.RawMessage(` { "value" : 1 } `)
	second, err := DigestChunkPlan(reformatted)
	if err != nil {
		t.Fatalf("DigestChunkPlan(reformatted) error = %v, want nil", err)
	}
	if first != second {
		t.Fatalf("digests = %q and %q, want stable canonical annotation digest", first, second)
	}

	changes := []func(*ChunkPlan){
		func(value *ChunkPlan) { value.Ranges[0].EndUnitID = 30 },
		func(value *ChunkPlan) { value.Annotations["plan"] = json.RawMessage(`{"value":2}`) },
		func(value *ChunkPlan) { value.Ranges[0].Annotations["range"] = json.RawMessage(`{"value":3}`) },
	}
	for i, change := range changes {
		changed := CloneChunkPlan(plan)
		change(&changed)
		digest, err := DigestChunkPlan(changed)
		if err != nil {
			t.Fatalf("DigestChunkPlan(change %d) error = %v", i, err)
		}
		if digest == first {
			t.Fatalf("DigestChunkPlan(change %d) = %q, want changed digest", i, digest)
		}
	}
}

func TestMaterializeChunkPlanExactOutputAndMutationSafety(t *testing.T) {
	doc := planDocument()
	plan := validChunkPlan(doc)
	plan.Ranges = []ChunkRange{
		{StartUnitID: 10, EndUnitID: 30, Annotations: ChunkAnnotations{"range": json.RawMessage(`{"value":2}`)}},
		{StartUnitID: 20, EndUnitID: 50, Annotations: ChunkAnnotations{"range": json.RawMessage(`{"value":3}`)}},
	}
	chunks, err := MaterializeChunkPlan(doc, plan)
	if err != nil {
		t.Fatalf("MaterializeChunkPlan() error = %v, want nil", err)
	}
	if len(chunks) != 2 {
		t.Fatalf("chunks = %d, want 2", len(chunks))
	}
	first := chunks[0]
	if first.ID != "chunk-000001" || first.SourceID != doc.ID || first.Index != 0 || first.MediaType != "application/json" {
		t.Fatalf("first chunk identity = %#v", first)
	}
	if want := (SourceRef{SourceID: doc.ID, StartUnitID: 10, EndUnitID: 30}); first.Ref != want {
		t.Fatalf("first ref = %#v, want %#v", first.Ref, want)
	}
	if got, want := unitIDs(first.Units), []int{10, 20, 30}; !reflect.DeepEqual(got, want) {
		t.Fatalf("first unit ids = %#v, want %#v", got, want)
	}
	wantContent, _ := json.Marshal(struct {
		Units []SourceUnit `json:"units"`
	}{Units: doc.Units[:3]})
	if !bytes.Equal(first.Content, wantContent) {
		t.Fatalf("first content = %s, want %s", first.Content, wantContent)
	}
	if !reflect.DeepEqual(first.Metadata, map[string]any{"start_unit_id": 10, "end_unit_id": 30, "unit_count": 3}) {
		t.Fatalf("first metadata = %#v", first.Metadata)
	}
	if string(first.Annotations["range"]) != `{"value":2}` || string(first.PlanAnnotations["plan"]) != `{"value":1}` {
		t.Fatalf("first annotations = %#v / %#v", first.Annotations, first.PlanAnnotations)
	}
	if got, want := unitIDs(chunks[1].Units), []int{20, 30, 40, 50}; !reflect.DeepEqual(got, want) {
		t.Fatalf("overlapping unit ids = %#v, want %#v", got, want)
	}

	again, err := MaterializeChunkPlan(doc, plan)
	if err != nil {
		t.Fatalf("MaterializeChunkPlan(repeated) error = %v", err)
	}
	if !reflect.DeepEqual(chunks, again) {
		t.Fatalf("repeated materialization differs:\nfirst: %#v\nagain: %#v", chunks, again)
	}
	chunks[0].Units[0].Text = "mutated"
	chunks[0].Annotations["range"][0] = '['
	chunks[0].PlanAnnotations["plan"][0] = '['
	if doc.Units[0].Text == "mutated" || string(plan.Ranges[0].Annotations["range"]) != `{"value":2}` || string(plan.Annotations["plan"]) != `{"value":1}` {
		t.Fatal("materialized chunk shares owned plan or source storage")
	}
}

func TestMaterializeChunkPlanDeepClonesUnitMetadata(t *testing.T) {
	doc := planDocument()
	doc.Units[0].Metadata = map[string]any{"nested": map[string]any{"values": []any{json.RawMessage(`{"ok":true}`)}}}
	chunks, err := MaterializeChunkPlan(doc, validChunkPlan(doc))
	if err != nil {
		t.Fatal(err)
	}
	nested := chunks[0].Units[0].Metadata["nested"].(map[string]any)
	nested["values"].([]any)[0].(json.RawMessage)[0] = '['
	nested["changed"] = true
	original := doc.Units[0].Metadata["nested"].(map[string]any)
	if _, exists := original["changed"]; exists || string(original["values"].([]any)[0].(json.RawMessage)) != `{"ok":true}` {
		t.Fatalf("source metadata changed through materialized chunk: %#v", doc.Units[0].Metadata)
	}
}

func TestMaterializeChunkPlanClonesConcreteJSONMetadata(t *testing.T) {
	doc := planDocument()
	doc.Units[0].Metadata = map[string]any{
		"typed_map":   typedMetadataMap{"bytes": []byte("map")},
		"typed_slice": typedMetadataSlice{{"raw": json.RawMessage(`{"slice":true}`)}},
		"typed_array": typedMetadataArray{map[string]any{"bytes": []byte("array")}, []any{json.RawMessage(`{"array":true}`)}},
		"interface":   any(typedMetadataMap{"bytes": []byte("interface")}),
		"raw":         json.RawMessage(`{"raw":true}`),
		"bytes":       []byte("bytes"),
	}

	chunks, err := MaterializeChunkPlan(doc, validChunkPlan(doc))
	if err != nil {
		t.Fatal(err)
	}
	again, err := MaterializeChunkPlan(doc, validChunkPlan(doc))
	if err != nil {
		t.Fatal(err)
	}
	metadata := chunks[0].Units[0].Metadata
	metadata["typed_map"].(typedMetadataMap)["bytes"].([]byte)[0] = 'M'
	metadata["typed_slice"].(typedMetadataSlice)[0]["raw"].(json.RawMessage)[0] = '['
	metadata["typed_array"].(typedMetadataArray)[0].(map[string]any)["bytes"].([]byte)[0] = 'A'
	metadata["typed_array"].(typedMetadataArray)[1].([]any)[0].(json.RawMessage)[0] = '['
	metadata["interface"].(typedMetadataMap)["bytes"].([]byte)[0] = 'I'
	metadata["raw"].(json.RawMessage)[0] = '['
	metadata["bytes"].([]byte)[0] = 'B'

	for name, candidate := range map[string]map[string]any{
		"source": doc.Units[0].Metadata,
		"again":  again[0].Units[0].Metadata,
	} {
		if got := string(candidate["typed_map"].(typedMetadataMap)["bytes"].([]byte)); got != "map" {
			t.Fatalf("%s typed map bytes = %q, want map", name, got)
		}
		if got := string(candidate["typed_slice"].(typedMetadataSlice)[0]["raw"].(json.RawMessage)); got != `{"slice":true}` {
			t.Fatalf("%s typed slice raw = %q", name, got)
		}
		array := candidate["typed_array"].(typedMetadataArray)
		if got := string(array[0].(map[string]any)["bytes"].([]byte)); got != "array" || string(array[1].([]any)[0].(json.RawMessage)) != `{"array":true}` {
			t.Fatalf("%s typed array = %#v", name, array)
		}
		if got := string(candidate["interface"].(typedMetadataMap)["bytes"].([]byte)); got != "interface" {
			t.Fatalf("%s interface bytes = %q, want interface", name, got)
		}
		if got := string(candidate["raw"].(json.RawMessage)); got != `{"raw":true}` {
			t.Fatalf("%s raw = %q", name, got)
		}
		if got := string(candidate["bytes"].([]byte)); got != "bytes" {
			t.Fatalf("%s bytes = %q, want bytes", name, got)
		}
	}
}

func TestMaterializeChunkPlanRejectsInvalidMetadata(t *testing.T) {
	cyclic := make(map[string]any)
	cyclic["self"] = cyclic
	for _, tc := range []struct {
		name  string
		value any
		want  string
	}{
		{name: "cycle", value: cyclic, want: "metadata.cycle.self contains a cycle"},
		{name: "unsupported", value: func() {}, want: "metadata.unsupported has unsupported type func()"},
		{name: "nonfinite", value: math.NaN(), want: "metadata.nonfinite has a non-finite number"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			doc := planDocument()
			doc.Units[0].Metadata = map[string]any{tc.name: tc.value}
			_, first := MaterializeChunkPlan(doc, validChunkPlan(doc))
			_, second := MaterializeChunkPlan(doc, validChunkPlan(doc))
			if first == nil || !strings.Contains(first.Error(), "clone chunk plan range[0] units: source unit[0] metadata: "+tc.want) {
				t.Fatalf("first MaterializeChunkPlan() error = %v, want %q", first, tc.want)
			}
			if second == nil || second.Error() != first.Error() {
				t.Fatalf("MaterializeChunkPlan() errors = %v and %v, want deterministic error", first, second)
			}
		})
	}
}

func planDocument() *SourceDocument {
	doc := &SourceDocument{ID: "source-plan", Kind: "test", Format: "application/test", Digest: "sha256:source-plan"}
	for _, id := range []int{10, 20, 30, 40, 50} {
		doc.Units = append(doc.Units, SourceUnit{ID: id, Kind: "line", Text: "unit", Ref: SourceRef{SourceID: doc.ID, StartUnitID: id, EndUnitID: id}})
	}
	return doc
}

func validChunkPlan(doc *SourceDocument) ChunkPlan {
	return ChunkPlan{
		SourceDigest: doc.Digest,
		Ranges: []ChunkRange{
			{StartUnitID: 10, EndUnitID: 20, Annotations: ChunkAnnotations{"range": json.RawMessage(`{"value":2}`)}},
			{StartUnitID: 30, EndUnitID: 50},
		},
		Annotations: ChunkAnnotations{"plan": json.RawMessage(`{"value":1}`)},
	}
}

func unitIDs(units []SourceUnit) []int {
	ids := make([]int, len(units))
	for i, unit := range units {
		ids[i] = unit.ID
	}
	return ids
}