package protocol

import (
	"bytes"
	"context"
	"errors"
	"io"
	"strings"
	"testing"
	"time"
)

func TestRequestMatrix(t *testing.T) {
	chat := []byte(`{"model":"old","max_tokens":100,"stream":true,"messages":[{"role":"system","content":"help"},{"role":"user","content":[{"type":"text","text":"look"},{"type":"image_url","image_url":{"url":"data:image/png;base64,YQ=="}}]},{"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"read","arguments":"{\"path\":\"a\"}"}}]},{"role":"tool","tool_call_id":"c1","content":"ok"},{"role":"user","content":"continue"}],"tools":[{"type":"function","function":{"name":"read","parameters":{"type":"object","properties":{"path":{"type":"string"}}}}}],"tool_choice":"auto"}`)
	for _, from := range []Kind{Chat, Responses, Anthropic} {
		src := chat
		var e error
		if from != Chat {
			src, e = ConvertRequest(Chat, from, chat, "new")
			if e != nil {
				t.Fatal(e)
			}
		}
		for _, to := range []Kind{Chat, Responses, Anthropic} {
			t.Run(string(from)+"/"+string(to), func(t *testing.T) {
				b, e := ConvertRequest(from, to, src, "target")
				if e != nil {
					t.Fatal(e)
				}
				m, e := decode(b)
				if e != nil || m["model"] != "target" {
					t.Fatalf("%s %v", b, e)
				}
				r, e := readRequest(to, m)
				if e != nil {
					t.Fatal(e)
				}
				var calls, results, images int
				for _, m := range r.messages {
					for _, p := range m.parts {
						switch p.typ {
						case "call":
							calls++
							if p.id != "c1" {
								t.Fatal(p)
							}
						case "result":
							results++
						case "image":
							images++
						}
					}
				}
				if calls != 1 || results != 1 || images != 1 {
					t.Fatalf("lost content: %s", b)
				}
			})
		}
	}
}
func TestStrictRejection(t *testing.T) {
	cases := []struct {
		k  Kind
		b  string
		to Kind
	}{
		{Chat, `{"messages":[],"response_format":{"type":"json_object"}}`, Responses},
		{Chat, `{"messages":[]}`, Anthropic},
		{Responses, `{"input":[],"previous_response_id":"r"}`, Chat},
		{Anthropic, `{"max_tokens":10,"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"x"}]}]}`, Chat},
		{Chat, `{"messages":[],"n":2}`, Responses},
		{Chat, `{"messages":[],"stop":["x"]}`, Responses},
		{Chat, `{"max_tokens":10,"messages":[],"tools":[{"type":"function","function":{"name":"x","parameters":{},"strict":true}}]}`, Anthropic},
	}
	for _, c := range cases {
		if b, e := ConvertRequest(c.k, c.to, []byte(c.b), "m"); e == nil {
			t.Fatalf("accepted %s => %s", c.b, b)
		}
	}
}
func TestResponseMatrix(t *testing.T) {
	chat := []byte(`{"id":"r","model":"m","choices":[{"index":0,"message":{"role":"assistant","content":"working","tool_calls":[{"id":"c","type":"function","function":{"name":"read","arguments":"{}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":4,"prompt_tokens_details":{"cached_tokens":3}}}`)
	for _, from := range []Kind{Chat, Responses, Anthropic} {
		src, e := ConvertResponse(Chat, from, chat)
		if e != nil {
			t.Fatal(e)
		}
		for _, to := range []Kind{Chat, Responses, Anthropic} {
			b, e := ConvertResponse(from, to, src)
			if e != nil {
				t.Fatal(from, to, e)
			}
			m, _ := decode(b)
			r, e := readResponse(to, m)
			if e != nil || len(r.parts) != 2 || r.finish != "tool" {
				t.Fatalf("%s %v", b, e)
			}
			u := ExtractUsage(to, b)
			if u.Input != 12 || u.Output != 4 || u.CacheRead != 3 || !u.Known {
				t.Fatal(u)
			}
		}
	}
}

type flushBuffer struct {
	bytes.Buffer
	flushes int
}

func (b *flushBuffer) Flush() { b.flushes++ }
func chatStream() string {
	return strings.Join([]string{
		`data: {"id":"r","model":"m","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"}}]}`,
		`data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"c1","type":"function","function":{"name":"read","arguments":"{\"x\":"}},{"index":1,"id":"c2","type":"function","function":{"name":"write","arguments":"{"}}]}}]}`,
		`data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"function":{"arguments":"}"}},{"index":0,"function":{"arguments":"1}"}}]}}]}`,
		`data: {"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}`,
		`data: {"choices":[],"usage":{"prompt_tokens":10,"completion_tokens":7}}`,
		`data: [DONE]`,
	}, "\n\n") + "\n\n"
}
func TestStreamMatrix(t *testing.T) {
	for _, from := range []Kind{Chat, Responses, Anthropic} {
		src := chatStream()
		if from != Chat {
			var b bytes.Buffer
			if _, e := Stream(context.Background(), &b, strings.NewReader(src), Chat, from, "m"); e != nil {
				t.Fatal(e)
			}
			src = b.String()
		}
		for _, to := range []Kind{Chat, Responses, Anthropic} {
			t.Run(string(from)+"/"+string(to), func(t *testing.T) {
				var b flushBuffer
				u, e := Stream(context.Background(), &b, strings.NewReader(src), from, to, "m")
				if e != nil {
					t.Fatal(e)
				}
				if u.Input != 10 || u.Output != 7 || !u.Known {
					t.Fatal(u)
				}
				if b.flushes < 3 {
					t.Fatal("not flushed")
				}
				if from == to && b.String() != src {
					t.Fatal("passthrough changed bytes")
				}
				if to == Responses {
					for _, event := range []string{"response.output_item.added", "response.function_call_arguments.delta", "response.function_call_arguments.done", "response.output_item.done", "response.completed"} {
						if !strings.Contains(b.String(), event) {
							t.Fatal("missing", event)
						}
					}
				}
				if to == Anthropic && !strings.Contains(b.String(), "message_stop") {
					t.Fatal("missing stop")
				}
			})
		}
	}
}
func TestUsageCumulative(t *testing.T) {
	src := "data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":5,\"output_tokens\":1,\"cache_read_input_tokens\":8,\"cache_creation_input_tokens\":2}}}\n\n" + "data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":9}}\n\n"
	var b bytes.Buffer
	u, e := Stream(context.Background(), &b, strings.NewReader(src), Anthropic, Anthropic, "")
	if e != nil {
		t.Fatal(e)
	}
	if u.Input != 15 || u.Output != 9 || u.CacheRead != 8 || u.CacheWrite != 2 {
		t.Fatal(u)
	}
	raw, e := decode(u.Raw)
	if e != nil || num(raw["input_tokens"]) != 5 || num(raw["output_tokens"]) != 9 || num(raw["cache_read_input_tokens"]) != 8 {
		t.Fatal(string(u.Raw), e)
	}
}
func TestRealtimeBeforeEOF(t *testing.T) {
	r, w := io.Pipe()
	defer r.Close()
	defer w.Close()
	done := make(chan error, 1)
	out := &signalWriter{signal: make(chan struct{}, 1)}
	go func() { _, e := Stream(context.Background(), out, r, Chat, Responses, "m"); done <- e }()
	first, rest, _ := strings.Cut(chatStream(), "\n\n")
	if _, e := io.WriteString(w, first+"\n\n"); e != nil {
		t.Fatal(e)
	}
	select {
	case <-out.signal:
	case <-time.After(time.Second):
		t.Fatal("first delta buffered until EOF")
	}
	if _, e := io.WriteString(w, rest); e != nil {
		t.Fatal(e)
	}
	select {
	case e := <-done:
		if e != nil {
			t.Fatal(e)
		}
	case <-time.After(time.Second):
		t.Fatal("terminal marker waits for EOF")
	}
}

type signalWriter struct{ signal chan struct{} }

func (w *signalWriter) Write(p []byte) (int, error) {
	select {
	case w.signal <- struct{}{}:
	default:
	}
	return len(p), nil
}
func TestUnknownUsageAndUniqueID(t *testing.T) {
	src := "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"a\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"
	var previous string
	for i := 0; i < 2; i++ {
		var b bytes.Buffer
		u, e := Stream(context.Background(), &b, strings.NewReader(src), Chat, Anthropic, "m")
		if e != nil || u.Known {
			t.Fatal(u, e)
		}
		line := strings.Split(b.String(), "\n")[1]
		m, e := decode([]byte(strings.TrimPrefix(line, "data: ")))
		if e != nil {
			t.Fatal(e)
		}
		msg := obj(m["message"])
		if obj(msg["usage"]) == nil {
			t.Fatal("missing initial usage")
		}
		id := str(msg["id"])
		if id == "" || id == previous {
			t.Fatal("nonunique id")
		}
		previous = id
	}
}
func TestStreamErrors(t *testing.T) {
	for _, src := range []string{"data: nope\n\n", "data: {\"error\":{\"message\":\"bad\"}}\n\n", strings.Replace(chatStream(), "data: [DONE]\n\n", "", 1), strings.Replace(chatStream(), `"arguments":"1}"`, `"arguments":"oops}"`, 1)} {
		var b bytes.Buffer
		if _, e := Stream(context.Background(), &b, strings.NewReader(src), Chat, Responses, "m"); e == nil {
			t.Fatal("accepted broken stream", src)
		}
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, e := Stream(ctx, io.Discard, strings.NewReader(chatStream()), Chat, Chat, ""); !errors.Is(e, context.Canceled) {
		t.Fatal(e)
	}
	if _, e := Stream(context.Background(), badWriter{}, strings.NewReader(chatStream()), Chat, Anthropic, ""); !errors.Is(e, io.ErrClosedPipe) {
		t.Fatal(e)
	}
}

type badWriter struct{}

func (badWriter) Write([]byte) (int, error) { return 0, io.ErrClosedPipe }
func TestSSEMultilineAndCRLF(t *testing.T) {
	src := ": comment\r\n\r\ndata: {\r\ndata: \"usage\":{\"prompt_tokens\":2,\"completion_tokens\":3}}\r\n\r\ndata: [DONE]\r\n\r\n"
	var b flushBuffer
	u, e := Stream(context.Background(), &b, strings.NewReader(src), Chat, Chat, "")
	if e != nil || u.Input != 2 || u.Output != 3 || b.String() != src {
		t.Fatal(u, e, b.String())
	}
}