package protocol

import (
	"context"
	"fmt"
	"io"
	"reflect"
	"strings"
	"testing"
)

func TestMergeUsageSizeBoundary(t *testing.T) {
	for _, kind := range []Kind{Chat, Responses, Anthropic} {
		t.Run(string(kind), func(t *testing.T) {
			m := object{"padding": strings.Repeat("x", maxStreamUsageBytes-len(`{"padding":""}`))}
			u, err := mergeUsage(kind, Usage{}, object{"usage": m})
			if err != nil || len(u.Raw) != maxStreamUsageBytes {
				t.Fatalf("exact limit: size=%d err=%v", len(u.Raw), err)
			}
			for _, update := range []object{
				{"usage": object{"padding": str(m["padding"]) + "x"}},
				{"usage": object{"new_key": 1}},
			} {
				got, err := mergeUsage(kind, u, update)
				if err == nil || !strings.Contains(err.Error(), "usage exceeds 16384 bytes") || !reflect.DeepEqual(got, u) {
					t.Fatalf("overflow: retained=%d err=%v", len(got.Raw), err)
				}
			}
			for i := 0; i < 3; i++ {
				got, err := mergeUsage(kind, u, object{"usage": m})
				if err != nil || !reflect.DeepEqual(got, u) {
					t.Fatalf("replacement at limit: %v", err)
				}
			}
		})
	}
}

func TestStreamUsageSizeLimits(t *testing.T) {
	for _, from := range []Kind{Chat, Responses, Anthropic} {
		for _, to := range []Kind{Chat, Responses, Anthropic} {
			for _, cumulative := range []bool{false, true} {
				t.Run(fmt.Sprintf("%s/%s/cumulative=%t", from, to, cumulative), func(t *testing.T) {
					var input strings.Builder
					writeUsage := func(u object) {
						m := object{"usage": u}
						switch from {
						case Anthropic:
							m["type"] = "ping"
						case Responses:
							m = object{"type": "response.in_progress", "response": m}
						}
						input.WriteString(limitEvent(mustLimitJSON(t, m)))
					}
					writeUsage(object{"prompt_tokens": 7, "input_tokens": 7})
					if cumulative {
						for i := 0; i < 32; i++ {
							writeUsage(object{fmt.Sprintf("key_%d", i): strings.Repeat("x", 1024)})
						}
					} else {
						writeUsage(object{"padding": strings.Repeat("x", maxStreamUsageBytes)})
					}
					u, err := Stream(context.Background(), io.Discard, strings.NewReader(input.String()), from, to, "")
					if err == nil || !strings.Contains(err.Error(), "usage exceeds 16384 bytes") {
						t.Fatalf("expected usage limit error, got %v", err)
					}
					if !u.Known || u.Input != 7 || len(u.Raw) > maxStreamUsageBytes {
						t.Fatalf("lost valid usage: known=%t input=%d size=%d", u.Known, u.Input, len(u.Raw))
					}
				})
			}
		}
	}
}

func TestMergeUsagePartialUpdates(t *testing.T) {
	u, err := mergeUsage(Anthropic, Usage{}, object{"message": object{"usage": object{"input_tokens": 5, "cache_read_input_tokens": 2}}})
	if err != nil {
		t.Fatal(err)
	}
	u, err = mergeUsage(Anthropic, u, object{"usage": object{"output_tokens": 3}})
	if err != nil || !u.Known || u.Input != 7 || u.Output != 3 || u.CacheRead != 2 {
		t.Fatalf("partial update: %+v err=%v", u, err)
	}
	got, err := mergeUsage(Anthropic, u, object{})
	if err != nil || !reflect.DeepEqual(got, u) {
		t.Fatalf("missing usage changed totals: %+v err=%v", got, err)
	}
}