package llm

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func newCompletionTestClient(t *testing.T, handler http.HandlerFunc) *Client {
	t.Helper()
	srv := httptest.NewServer(handler)
	t.Cleanup(srv.Close)
	client, err := NewClient(Config{Provider: "test", URL: srv.URL}, slog.Default())
	if err != nil {
		t.Fatalf("NewClient: %v", err)
	}
	return client
}

func openCompletionTestStream(t *testing.T, client *Client) *CompletionStream {
	t.Helper()
	stream, err := client.CompletionsStream(
		t.Context(), CompletionsParams{Model: "test-model"},
	)
	if err != nil {
		t.Fatalf("CompletionsStream: %v", err)
	}
	t.Cleanup(func() { _ = stream.Close() })
	return stream
}

func TestCompletionsStreamReconstructsReasoning(t *testing.T) {
	client := newCompletionTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		for _, data := range []string{
			`{"choices":[{"index":0,"delta":{"role":"assistant"}}]}`,
			`{"choices":[{"index":0,"delta":{"reasoning_content":"thinking"}}]}`,
			`{"choices":[{"index":0,"delta":{"reasoning_content":" more"}}]}`,
			`{"choices":[{"index":0,"delta":{"content":"answer"},"finish_reason":"stop"}]}`,
		} {
			fmt.Fprintf(w, "data: %s\n\n", data)
		}
		fmt.Fprint(w, "data: [DONE]\n\n")
	})

	stream := openCompletionTestStream(t, client)
	if got := stream.Result(); got.Message.Text() != "" {
		t.Fatalf("initial result text = %q, want empty", got.Message.Text())
	}
	var chunks []StreamChunk
	for stream.Next() {
		chunks = append(chunks, stream.Current())
	}
	if err := stream.Err(); err != nil {
		t.Fatalf("stream.Err: %v", err)
	}
	if len(chunks) != 4 {
		t.Fatalf("chunks = %d, want 4", len(chunks))
	}
	result := stream.Result()
	if result.Message.ReasoningContent != "thinking more" {
		t.Errorf("reasoning = %q", result.Message.ReasoningContent)
	}
	if result.Message.Text() != "answer" {
		t.Errorf("text = %q", result.Message.Text())
	}
	if result.FinishReason != "stop" {
		t.Errorf("finish reason = %q", result.FinishReason)
	}
	if stream.Next() {
		t.Error("Next returned true after termination")
	}
}

func TestCompletionsStreamReportsCancellationWithPartialResult(t *testing.T) {
	firstSent := make(chan struct{})
	client := newCompletionTestClient(t, func(w http.ResponseWriter, r *http.Request) {
		flusher := w.(http.Flusher)
		fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"first\"}}]}\n\n")
		flusher.Flush()
		close(firstSent)
		<-r.Context().Done()
	})

	ctx, cancel := context.WithCancel(t.Context())
	stream, err := client.CompletionsStream(ctx, CompletionsParams{Model: "test-model"})
	if err != nil {
		t.Fatalf("CompletionsStream: %v", err)
	}
	defer stream.Close()
	<-firstSent
	if !stream.Next() {
		t.Fatalf("first Next = false: %v", stream.Err())
	}
	cancel()
	if stream.Next() {
		t.Fatal("Next returned true after cancellation")
	}
	if !errors.Is(stream.Err(), context.Canceled) {
		t.Fatalf("Err = %v, want context.Canceled", stream.Err())
	}
	result := stream.Result()
	if got := result.Message.Text(); got != "first" {
		t.Errorf("partial text = %q, want first", got)
	}
}

func TestCompletionsStreamReconstructsFragmentedParallelToolCalls(t *testing.T) {
	client := newCompletionTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
		for _, data := range []string{
			`{"choices":[{"index":0,"delta":{"role":"assistant"}}]}`,
			`{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_0","function":{"name":"ec"}},{"index":1,"id":"call_1","function":{"name":"pr"}}]}}]}`,
			`{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"name":"ho","arguments":"{\"value\":"}},{"index":1,"function":{"name":"intf","arguments":"{\"format\":"}}]}}]}`,
			`{"choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"function":{"arguments":"\"%s\"}"}},{"index":0,"function":{"arguments":"\"one\"}"}}]},"finish_reason":"tool_calls"}]}`,
		} {
			fmt.Fprintf(w, "data: %s\n\n", data)
		}
		fmt.Fprint(w, "data: [DONE]\n\n")
	})

	stream := openCompletionTestStream(t, client)
	var chunks []StreamChunk
	for stream.Next() {
		chunks = append(chunks, stream.Current())
	}
	if err := stream.Err(); err != nil {
		t.Fatalf("stream.Err: %v", err)
	}
	if len(chunks) != 4 || len(chunks[1].Choices[0].Delta.ToolCalls) != 2 {
		t.Fatalf("tool-call chunks not exposed unchanged: %#v", chunks)
	}
	result := stream.Result()
	toolCalls := result.Message.ToolCalls
	if len(toolCalls) != 2 {
		t.Fatalf("tool calls = %d, want 2", len(toolCalls))
	}
	if toolCalls[0].ID != "call_0" || toolCalls[0].Function.Name != "echo" ||
		toolCalls[0].Function.Arguments != `{"value":"one"}` {
		t.Errorf("tool call 0 = %#v", toolCalls[0])
	}
	if toolCalls[1].ID != "call_1" || toolCalls[1].Function.Name != "printf" ||
		toolCalls[1].Function.Arguments != `{"format":"%s"}` {
		t.Errorf("tool call 1 = %#v", toolCalls[1])
	}
	if result.FinishReason != "tool_calls" {
		t.Errorf("finish reason = %q", result.FinishReason)
	}
}

func TestCompletionsReturnsUsageAndTimings(t *testing.T) {
	client := newCompletionTestClient(t, func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/chat/completions" {
			t.Fatalf("unexpected path: %s", r.URL.Path)
		}
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{
			"id": "test-123",
			"model": "test-model",
			"choices": [{
				"index": 0,
				"message": {"role": "assistant", "content": "hello"},
				"finish_reason": "stop"
			}],
			"usage": {
				"prompt_tokens": 10,
				"completion_tokens": 5,
				"total_tokens": 15,
				"prompt_tokens_details": {"cached_tokens": 3},
				"completion_tokens_details": {"reasoning_tokens": 2}
			},
			"timings": {
				"prompt_n": 10,
				"prompt_ms": 50.5,
				"prompt_per_token_ms": 5.05,
				"prompt_per_second": 198.0,
				"predicted_n": 5,
				"predicted_ms": 100.0,
				"predicted_per_token_ms": 20.0,
				"predicted_per_second": 50.0
			}
		}`)
	})

	res, err := client.Completions(context.Background(), CompletionsParams{Model: "test-model"})
	if err != nil {
		t.Fatalf("Completions: %v", err)
	}
	if res.Message.Text() != "hello" || res.Usage == nil || res.Timings == nil {
		t.Fatalf("unexpected result: %#v", res)
	}
	if res.Usage.PromptTokens != 10 || res.Usage.CompletionTokens != 5 || res.Usage.TotalTokens != 15 {
		t.Errorf("usage = %#v", res.Usage)
	}
	if res.Usage.PromptTokensDetails == nil || res.Usage.PromptTokensDetails.CachedTokens != 3 {
		t.Errorf("prompt token details = %#v", res.Usage.PromptTokensDetails)
	}
	if res.Usage.CompletionTokensDetails == nil || res.Usage.CompletionTokensDetails.ReasoningTokens != 2 {
		t.Errorf("completion token details = %#v", res.Usage.CompletionTokensDetails)
	}
	if res.Timings.PromptMS != 50.5 || res.Timings.PredictedMS != 100 {
		t.Errorf("timings = %#v", res.Timings)
	}
}

func TestCompletionsStreamSSEAndMetadata(t *testing.T) {
	large := strings.Repeat("x", 70*1024)
	client := newCompletionTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
		fmt.Fprint(w, ": comment\n")
		fmt.Fprintf(w, "data:{\"choices\":[{\"index\":0,\"delta\":{\"content\":%q}}],\n", large)
		fmt.Fprint(w, "data: \"usage\":{\"prompt_tokens\":10,\"total_tokens\":15}}\n\n")
		fmt.Fprint(w, "data: {\"choices\":[],\"timings\":{\"prompt_ms\":50.5}}\n\n")
		fmt.Fprint(w, "data:[DONE]\n\n")
	})

	stream := openCompletionTestStream(t, client)
	var count int
	for stream.Next() {
		count++
	}
	if err := stream.Err(); err != nil {
		t.Fatalf("stream.Err: %v", err)
	}
	result := stream.Result()
	if count != 2 || result.Message.Text() != large {
		t.Fatalf("count = %d, text length = %d", count, len(result.Message.Text()))
	}
	if result.Usage == nil || result.Usage.PromptTokens != 10 || result.Timings == nil || result.Timings.PromptMS != 50.5 {
		t.Errorf("metadata = usage %#v, timings %#v", result.Usage, result.Timings)
	}
}

func TestCompletionsStreamFailuresPreservePartialResult(t *testing.T) {
	tests := []struct {
		name     string
		response string
		wantErr  error
	}{
		{
			name:     "EOF before done",
			response: "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"}}]}\n\n",
			wantErr:  io.ErrUnexpectedEOF,
		},
		{
			name:     "malformed JSON",
			response: "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"}}]}\n\ndata: {bad}\n\n",
		},
		{
			name:     "unexpected choice",
			response: "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"}}]}\n\ndata: {\"choices\":[{\"index\":1,\"delta\":{}}]}\n\n",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			client := newCompletionTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
				fmt.Fprint(w, tt.response)
			})
			stream := openCompletionTestStream(t, client)
			for stream.Next() {
			}
			if stream.Err() == nil {
				t.Fatal("Err = nil")
			}
			if tt.wantErr != nil && !errors.Is(stream.Err(), tt.wantErr) {
				t.Errorf("Err = %v, want %v", stream.Err(), tt.wantErr)
			}
			result := stream.Result()
			if got := result.Message.Text(); got != "partial" {
				t.Errorf("partial text = %q", got)
			}
		})
	}
}

func TestCompletionStreamCloseIsIdempotent(t *testing.T) {
	client := newCompletionTestClient(t, func(w http.ResponseWriter, r *http.Request) {
		w.(http.Flusher).Flush()
		<-r.Context().Done()
	})
	stream := openCompletionTestStream(t, client)
	if err := stream.Close(); err != nil {
		t.Fatalf("first Close: %v", err)
	}
	if err := stream.Close(); err != nil {
		t.Fatalf("second Close: %v", err)
	}
	if stream.Err() != nil {
		t.Errorf("Err after Close = %v", stream.Err())
	}
	if stream.Next() {
		t.Error("Next after Close returned true")
	}
}

func TestCompletionsStreamRequiresCompleteDoneEvent(t *testing.T) {
	client := newCompletionTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
		fmt.Fprint(w, "data: [DONE]")
	})
	stream := openCompletionTestStream(t, client)
	if stream.Next() {
		t.Fatal("Next returned true")
	}
	if !errors.Is(stream.Err(), io.ErrUnexpectedEOF) {
		t.Fatalf("Err = %v, want io.ErrUnexpectedEOF", stream.Err())
	}
}