package protocol

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"strings"
	"testing"
)

func limitEvent(s string) string { return "data: " + s + "\n\n" }
func TestSSEReadLimits(t *testing.T) {
	for _, to := range []Kind{Chat, Anthropic} {
		var out bytes.Buffer
		_, err := streamWithLimits(context.Background(), &out, strings.NewReader(strings.Repeat("x", 10000)), Chat, to, "", 64, 8, 64)
		if err == nil || !strings.Contains(err.Error(), "size limit") || out.Len() != 0 {
			t.Fatalf("%s: %v", to, err)
		}
	}
	for _, input := range []string{"12345678", "1234567\n"} {
		line, err := readSSELine(context.Background(), bufio.NewReader(strings.NewReader(input)), 8)
		if line != input || (err != nil && err != io.EOF) {
			t.Fatalf("%q %v", line, err)
		}
	}
	var out bytes.Buffer
	_, err := streamWithLimits(context.Background(), &out, strings.NewReader(":123456\n:123456\n"), Chat, Chat, "", 12, 8, 16)
	if err == nil || out.String() != ":123456\n" {
		t.Fatalf("%q %v", out.String(), err)
	}
}
func TestStreamPassthroughTotalLimit(t *testing.T) {
	frame := "data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":null}]}\n\n"
	var out bytes.Buffer
	_, err := streamWithLimits(context.Background(), &out, strings.NewReader(strings.Repeat(frame, 4)), Chat, Chat, "", 512, 8, 2*len(frame))
	if err == nil || !strings.Contains(err.Error(), "passthrough") || out.Len() != 2*len(frame) {
		t.Fatalf("out=%d %q %v", out.Len(), out.String(), err)
	}
}
func TestStreamCumulativeLimits(t *testing.T) {
	for _, field := range []string{"text", "arguments", "identity"} {
		delta := object{"content": strings.Repeat("x", 100)}
		if field != "text" {
			call := object{"index": 0}
			if field == "arguments" {
				call["function"] = object{"arguments": strings.Repeat("x", 100)}
			} else {
				call["id"] = strings.Repeat("x", 100)
			}
			delta = object{"tool_calls": []any{call}}
		}
		event := limitEvent(fmt.Sprintf("{\"choices\":[{\"index\":0,\"delta\":%s}]}", mustLimitJSON(t, delta)))
		_, err := streamWithLimits(context.Background(), io.Discard, strings.NewReader(strings.Repeat(event, 6)), Chat, Responses, "", 512, 8, 512)
		if err == nil || !strings.Contains(err.Error(), "retained content") {
			t.Fatalf("%s: %v", field, err)
		}
	}
}
func TestStreamBlockLimits(t *testing.T) {
	for _, kind := range []string{"text", "thinking", "reasoning", "summary", "tool"} {
		from := Anthropic
		in := limitEvent("{\"type\":\"message_start\",\"message\":{\"content\":[]}}")
		if kind == "reasoning" || kind == "summary" || kind == "tool" {
			from = Responses
			in = ""
		}
		if kind == "summary" {
			in = limitEvent("{\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"reasoning\",\"id\":\"r\"}}")
		}
		for i := 0; i < 4; i++ {
			m := object{"type": "content_block_start", "index": i, "content_block": object{"type": kind, "text": ""}}
			switch kind {
			case "thinking":
				m["content_block"] = object{"type": "thinking", "thinking": "", "signature": ""}
			case "reasoning", "tool":
				item := object{"type": "reasoning", "id": "r"}
				if kind == "tool" {
					item = object{"type": "function_call", "call_id": "c", "name": "f", "arguments": ""}
				}
				m = object{"type": "response.output_item.added", "output_index": i, "item": item}
			case "summary":
				m = object{"type": "response.reasoning_summary_part.added", "output_index": 0, "item_id": "r", "summary_index": i, "part": object{"type": "summary_text", "text": ""}}
			}
			in += limitEvent(mustLimitJSON(t, m))
		}
		_, err := streamWithLimits(context.Background(), io.Discard, strings.NewReader(in), from, Chat, "", 1024, 3, 1024)
		if err == nil || !strings.Contains(err.Error(), "3 blocks") {
			t.Fatalf("%s: %v", kind, err)
		}
	}
}

func TestReasoningIdentityCumulativeLimit(t *testing.T) {
	var in strings.Builder
	for i := 0; i < 6; i++ {
		in.WriteString(limitEvent(mustLimitJSON(t, object{"type": "response.output_item.added", "output_index": i, "item": object{"type": "reasoning", "id": strings.Repeat("r", 100)}})))
	}
	_, err := streamWithLimits(context.Background(), io.Discard, strings.NewReader(in.String()), Responses, Chat, "", 512, 8, 512)
	if err == nil || !strings.Contains(err.Error(), "retained content") {
		t.Fatal(err)
	}
}
func TestStreamBufferSemantics(t *testing.T) {
	var out bytes.Buffer
	s := &streamState{ctx: context.Background(), w: &out, from: Chat, to: Responses, source: map[int]*streamBlock{}, byteLimit: 64, blockLimit: 4}
	b, err := s.add(0, part{typ: "call"})
	if err != nil {
		t.Fatal(err)
	}
	if err = s.delta(b, "{\"a\":"); err != nil {
		t.Fatal(err)
	}
	if out.Len() != 0 {
		t.Fatal("incomplete identity emitted")
	}
	b.id, b.name = "c", "f"
	if err = s.delta(b, "1}"); err != nil {
		t.Fatal(err)
	}
	if b.args != "{\"a\":1}" || b.sent != 7 {
		t.Fatal(b.part)
	}
	if err = s.close(b); err != nil {
		t.Fatal(err)
	}
	text, err := s.add(1, part{typ: "text"})
	if err != nil {
		t.Fatal(err)
	}
	for _, v := range []string{"ab", "cd"} {
		if err = s.delta(text, v); err != nil {
			t.Fatal(err)
		}
	}
	if text.text != "abcd" {
		t.Fatal(text.text)
	}
	if err = s.finish("stop"); err != nil {
		t.Fatal(err)
	}
	if err = s.end(); err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(out.String(), "abcd") || !strings.Contains(out.String(), "response.completed") {
		t.Fatal(out.String())
	}
	s.finished = false
	text.closed = false
	s.byteLimit = s.retained
	if err = s.delta(text, "x"); err == nil || text.text != "abcd" {
		t.Fatal("overflow changed content")
	}
}
func mustLimitJSON(t *testing.T, m object) string {
	t.Helper()
	b, err := json.Marshal(m)
	if err != nil {
		t.Fatal(err)
	}
	return string(b)
}