package llm

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

// sse serves the chunks as a streaming completion.
func sse(t *testing.T, check func(req map[string]any), chunks ...string) *httptest.Server {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		var req map[string]any
		if err := json.Unmarshal(body, &req); err != nil {
			t.Errorf("bad request body: %v", err)
		}
		if check != nil {
			check(req)
		}
		w.Header().Set("Content-Type", "text/event-stream")
		for _, c := range chunks {
			fmt.Fprintf(w, "data: %s\n\n", c)
		}
		fmt.Fprint(w, "data: [DONE]\n\n")
	}))
	t.Cleanup(srv.Close)
	return srv
}

func delta(d string) string { return `{"choices":[{"delta":` + d + `}]}` }

func TestStreamToolCalls(t *testing.T) {
	srv := sse(t, func(req map[string]any) {
		if req["stream"] != true || req["tool_choice"] != "auto" || fmt.Sprint(req["stream_options"]) != "map[include_usage:true]" {
			t.Errorf("request = %v", req)
		}
		tools, _ := req["tools"].([]any)
		if len(tools) != 1 {
			t.Errorf("tools = %v", req["tools"])
		}
	},
		delta(`{"reasoning_content":"let me "}`),
		delta(`{"reasoning_content":"look"}`),
		delta(`{"content":"Reading."}`),
		delta(`{"tool_calls":[{"index":0,"id":"c1","function":{"name":"Read","arguments":"{\"pa"}}]}`),
		delta(`{"tool_calls":[{"index":0,"function":{"arguments":"th\":\"a.go\"}"}}]}`),
		delta(`{"tool_calls":[{"index":1,"id":"c2","function":{"name":"Safe_Bash","arguments":"{}"}}]}`),
		`{"choices":[],"usage":{"prompt_tokens":1200,"completion_tokens":34}}`,
	)
	c := New(srv.URL, "", "m")
	var reply, reasoning strings.Builder
	r, err := c.Complete(context.Background(), []Message{{Role: "user", Content: "hi"}},
		[]Tool{{Name: "Read", Parameters: map[string]any{"type": "object"}}},
		func(text string, isReasoning bool) {
			if isReasoning {
				reasoning.WriteString(text)
			} else {
				reply.WriteString(text)
			}
		})
	if err != nil {
		t.Fatal(err)
	}
	if r.Reply != "Reading." || r.Reasoning != "let me look" || reply.String() != "Reading." || reasoning.String() != "let me look" {
		t.Errorf("reply %q reasoning %q streamed %q / %q", r.Reply, r.Reasoning, reply.String(), reasoning.String())
	}
	if len(r.Calls) != 2 || r.Calls[0].Function.Arguments != `{"path":"a.go"}` || r.Calls[1].ID != "c2" {
		t.Errorf("calls = %+v", r.Calls)
	}
	if r.Usage != (Usage{PromptTokens: 1200, CompletionTokens: 34}) {
		t.Errorf("usage = %+v", r.Usage)
	}
}

func TestInlineThinkingSplitAcrossChunks(t *testing.T) {
	srv := sse(t, nil,
		delta(`{"content":"<thi"}`),
		delta(`{"content":"nk>hmm, </th"}`),
		delta(`{"content":"ink>\n\nAnswer"}`),
	)
	var reply, reasoning strings.Builder
	r, err := New(srv.URL, "", "m").Complete(context.Background(), nil, nil, func(text string, isReasoning bool) {
		if isReasoning {
			reasoning.WriteString(text)
		} else {
			reply.WriteString(text)
		}
	})
	if err != nil {
		t.Fatal(err)
	}
	if r.Reply != "Answer" || reply.String() != "Answer" || strings.TrimSpace(reasoning.String()) != "hmm," {
		t.Errorf("reply %q streamed %q reasoning %q", r.Reply, reply.String(), reasoning.String())
	}
}

func TestNonStreamingFallbackAndErrors(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer k" {
			w.WriteHeader(http.StatusUnauthorized)
			fmt.Fprint(w, `{"error":{"message":"bad key"}}`)
			return
		}
		fmt.Fprint(w, `{"choices":[{"message":{"role":"assistant","content":"plain"}}]}`)
	}))
	defer srv.Close()
	r, err := New(srv.URL, "k", "m").Complete(context.Background(), nil, nil, nil)
	if err != nil || r.Reply != "plain" {
		t.Errorf("got %q, %v", r.Reply, err)
	}
	_, err = New(srv.URL, "wrong", "m").Complete(context.Background(), nil, nil, nil)
	if err == nil || !strings.Contains(err.Error(), "bad key") {
		t.Errorf("err = %v", err)
	}
}

func TestRedact(t *testing.T) {
	if got := string(redact([]byte("Authorization: Bearer sk-123"), "sk-123")); strings.Contains(got, "sk-123") {
		t.Error(got)
	}
}

func TestContextSize(t *testing.T) {
	routes := map[string]string{
		"/vllm/v1/models":  `{"data":[{"id":"other","max_model_len":8192},{"id":"m","max_model_len":32768}]}`,
		"/llama/v1/models": `{"data":[{"id":"m"}]}`,
		"/llama/props":     `{"default_generation_settings":{"n_ctx":16384}}`,
		"/other/v1/models": `{"data":[{"id":"m","max_model_len":-1}]}`,
	}
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, ok := routes[r.URL.Path]
		if !ok {
			w.WriteHeader(http.StatusNotFound)
		}
		fmt.Fprint(w, body)
	}))
	defer srv.Close()
	for server, want := range map[string]int{"vllm": 32768, "llama": 16384, "other": 0} {
		got, err := New(srv.URL+"/"+server+"/v1", "", "m").ContextSize(context.Background())
		if err != nil || got != want {
			t.Errorf("%s: ContextSize = %d, %v; want %d", server, got, err, want)
		}
	}
}

// A redirect, even to the same host name on another port, never gets the
// key: Go would forward the Authorization header there.
func TestRedirectKeepsKey(t *testing.T) {
	got := make(chan string, 4)
	other := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got <- r.Header.Get("Authorization")
	}))
	t.Cleanup(other.Close)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, other.URL+r.URL.Path, http.StatusTemporaryRedirect)
	}))
	t.Cleanup(srv.Close)
	c := New(srv.URL, "sk-SECRET", "m")
	if _, err := c.Complete(context.Background(), []Message{{Role: "user", Content: "hi"}}, nil, nil); !errors.Is(err, ErrRedirect) {
		t.Errorf("Complete: %v", err)
	}
	if _, err := c.Models(context.Background()); err == nil {
		t.Error("Models followed the redirect")
	}
	close(got)
	for auth := range got {
		t.Errorf("the redirect target got %q", auth)
	}
}

// A stream is bounded whatever MaxTokens says, tool calls' names
// included, and one that stops without saying it is done is an error.
func TestStreamLimits(t *testing.T) {
	name := strings.Repeat("x", 1<<20)
	var chunks []string
	for range 100 {
		chunks = append(chunks, delta(`{"tool_calls":[{"index":0,"function":{"name":"`+name+`"}}]}`))
	}
	c := New(sse(t, nil, chunks...).URL, "", "m")
	c.MaxTokens = 100
	if _, err := c.Complete(context.Background(), []Message{{Role: "user", Content: "hi"}}, nil, nil); !errors.Is(err, ErrTruncated) {
		t.Errorf("a 100 MB tool name: %v", err)
	}

	cut := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		fmt.Fprintf(w, "data: %s\n\n", delta(`{"content":"Running."}`))
		fmt.Fprintf(w, "data: %s\n\n", delta(`{"tool_calls":[{"index":0,"id":"c1","function":{"name":"Safe_Bash","arguments":"{\"script\":\"ls\"}"}}]}`))
	}))
	t.Cleanup(cut.Close)
	r, err := New(cut.URL, "", "m").Complete(context.Background(), []Message{{Role: "user", Content: "hi"}}, nil, nil)
	if err == nil || len(r.Calls) != 0 {
		t.Errorf("a stream cut short: %+v, %v", r, err)
	}
}