package assistant_test

import (
	"errors"
	"io"
	"strings"
	"testing"

	"congo.gg/pkg/assistant"
)

// sseReader wraps a string reader with a no-op Close method for use as io.ReadCloser.
type sseReader struct {
	io.Reader
}

func (sseReader) Close() error { return nil }

// newSSEReader creates an io.ReadCloser from raw SSE text.
func newSSEReader(data string) io.ReadCloser {
	return sseReader{strings.NewReader(data)}
}

// simpleContentParser is a parseEvent function that treats data as literal content text.
func simpleContentParser(data string) (*assistant.StreamEvent, error) {
	if strings.HasPrefix(data, `{"type":"content"`) {
		// Extract text value: {"type":"content","text":"..."}
		start := strings.Index(data, `"text":"`) + 8
		end := strings.LastIndex(data, `"`)
		if start > 8 && end > start {
			return &assistant.StreamEvent{
				Type:    assistant.EventContentDelta,
				Content: data[start:end],
			}, nil
		}
	}

	if strings.HasPrefix(data, `{"type":"tool_start"`) {
		id := extractJSONField(data, "id")
		name := extractJSONField(data, "name")
		return &assistant.StreamEvent{
			Type: assistant.EventToolCallStart,
			ToolCall: &assistant.ToolCall{
				ID:   id,
				Name: name,
			},
		}, nil
	}

	if strings.HasPrefix(data, `{"type":"tool_delta"`) {
		args := extractJSONField(data, "arguments")
		return &assistant.StreamEvent{
			Type: assistant.EventToolCallDelta,
			ToolCall: &assistant.ToolCall{
				Arguments: args,
			},
		}, nil
	}

	if strings.HasPrefix(data, `{"type":"error"`) {
		msg := extractJSONField(data, "message")
		return &assistant.StreamEvent{
			Type:  assistant.EventError,
			Error: errors.New(msg),
		}, nil
	}

	if strings.HasPrefix(data, `{"type":"done"`) {
		return &assistant.StreamEvent{
			Type: assistant.EventDone,
		}, nil
	}

	return nil, errors.New("unknown event type")
}

// extractJSONField extracts a simple string field from a JSON-like string.
func extractJSONField(data, field string) string {
	key := `"` + field + `":"`
	start := strings.Index(data, key)
	if start < 0 {
		return ""
	}
	start += len(key)
	end := strings.Index(data[start:], `"`)
	if end < 0 {
		return ""
	}
	return data[start : start+end]
}

func TestStreamReaderContentDelta(t *testing.T) {
	sse := "data: {\"type\":\"content\",\"text\":\"Hello\"}\n\n" +
		"data: {\"type\":\"content\",\"text\":\" World\"}\n\n" +
		"data: [DONE]\n\n"

	reader := assistant.NewStreamReader(newSSEReader(sse), simpleContentParser)
	defer reader.Close()

	// First event: "Hello"
	event := reader.Next()
	if event == nil {
		t.Fatal("expected event, got nil")
	}
	if event.Type != assistant.EventContentDelta {
		t.Errorf("Type = %d, want EventContentDelta (%d)", event.Type, assistant.EventContentDelta)
	}
	if event.Content != "Hello" {
		t.Errorf("Content = %q, want %q", event.Content, "Hello")
	}

	// Second event: " World"
	event = reader.Next()
	if event == nil {
		t.Fatal("expected event, got nil")
	}
	if event.Type != assistant.EventContentDelta {
		t.Errorf("Type = %d, want EventContentDelta", event.Type)
	}
	if event.Content != " World" {
		t.Errorf("Content = %q, want %q", event.Content, " World")
	}

	// Done event
	event = reader.Next()
	if event == nil {
		t.Fatal("expected done event, got nil")
	}
	if event.Type != assistant.EventDone {
		t.Errorf("Type = %d, want EventDone (%d)", event.Type, assistant.EventDone)
	}

	// Accumulated content
	if reader.Content() != "Hello World" {
		t.Errorf("Content() = %q, want %q", reader.Content(), "Hello World")
	}
}

func TestStreamReaderToolCallStartAndDelta(t *testing.T) {
	sse := "data: {\"type\":\"tool_start\",\"id\":\"call_1\",\"name\":\"search\"}\n\n" +
		"data: {\"type\":\"tool_delta\",\"arguments\":\"{\\\"q\\\"\"}\n\n" +
		"data: {\"type\":\"tool_delta\",\"arguments\":\":\\\"test\\\"}\"}\n\n" +
		"data: [DONE]\n\n"

	// Custom parser that handles tool_delta with escaped JSON
	parser := func(data string) (*assistant.StreamEvent, error) {
		if strings.HasPrefix(data, `{"type":"tool_start"`) {
			id := extractJSONField(data, "id")
			name := extractJSONField(data, "name")
			return &assistant.StreamEvent{
				Type: assistant.EventToolCallStart,
				ToolCall: &assistant.ToolCall{
					ID:   id,
					Name: name,
				},
			}, nil
		}
		if strings.HasPrefix(data, `{"type":"tool_delta"`) {
			// Extract arguments - handle escaped quotes
			key := `"arguments":"`
			start := strings.Index(data, key)
			if start < 0 {
				return nil, errors.New("no arguments")
			}
			start += len(key)
			// Find closing quote (not escaped)
			end := start
			for end < len(data) {
				if data[end] == '\\' {
					end += 2
					continue
				}
				if data[end] == '"' {
					break
				}
				end++
			}
			args := data[start:end]
			args = strings.ReplaceAll(args, `\"`, `"`)
			args = strings.ReplaceAll(args, `\\`, `\`)
			return &assistant.StreamEvent{
				Type: assistant.EventToolCallDelta,
				ToolCall: &assistant.ToolCall{
					Arguments: args,
				},
			}, nil
		}
		return nil, errors.New("unknown")
	}

	reader := assistant.NewStreamReader(newSSEReader(sse), parser)
	defer reader.Close()

	// Tool call start
	event := reader.Next()
	if event == nil {
		t.Fatal("expected tool_start event, got nil")
	}
	if event.Type != assistant.EventToolCallStart {
		t.Errorf("Type = %d, want EventToolCallStart (%d)", event.Type, assistant.EventToolCallStart)
	}
	if event.ToolCall.ID != "call_1" {
		t.Errorf("ToolCall.ID = %q, want %q", event.ToolCall.ID, "call_1")
	}
	if event.ToolCall.Name != "search" {
		t.Errorf("ToolCall.Name = %q, want %q", event.ToolCall.Name, "search")
	}

	// Tool call delta 1
	event = reader.Next()
	if event == nil {
		t.Fatal("expected tool_delta event, got nil")
	}
	if event.Type != assistant.EventToolCallDelta {
		t.Errorf("Type = %d, want EventToolCallDelta (%d)", event.Type, assistant.EventToolCallDelta)
	}

	// Tool call delta 2
	event = reader.Next()
	if event == nil {
		t.Fatal("expected second tool_delta event, got nil")
	}
	if event.Type != assistant.EventToolCallDelta {
		t.Errorf("Type = %d, want EventToolCallDelta", event.Type)
	}

	// Done
	event = reader.Next()
	if event == nil {
		t.Fatal("expected done event, got nil")
	}
	if event.Type != assistant.EventDone {
		t.Errorf("Type = %d, want EventDone", event.Type)
	}

	// Check accumulated tool calls
	toolCalls := reader.ToolCalls()
	if len(toolCalls) != 1 {
		t.Fatalf("ToolCalls length = %d, want 1", len(toolCalls))
	}
	if toolCalls[0].ID != "call_1" {
		t.Errorf("ToolCalls[0].ID = %q, want %q", toolCalls[0].ID, "call_1")
	}
	if toolCalls[0].Name != "search" {
		t.Errorf("ToolCalls[0].Name = %q, want %q", toolCalls[0].Name, "search")
	}
	expectedArgs := `{"q":"test"}`
	if toolCalls[0].Arguments != expectedArgs {
		t.Errorf("ToolCalls[0].Arguments = %q, want %q", toolCalls[0].Arguments, expectedArgs)
	}
}

func TestStreamReaderDoneViaDONEMarker(t *testing.T) {
	sse := "data: {\"type\":\"content\",\"text\":\"hi\"}\n\n" +
		"data: [DONE]\n\n"

	reader := assistant.NewStreamReader(newSSEReader(sse), simpleContentParser)
	defer reader.Close()

	// Content event
	event := reader.Next()
	if event == nil || event.Type != assistant.EventContentDelta {
		t.Fatal("expected content event")
	}

	// DONE marker
	event = reader.Next()
	if event == nil || event.Type != assistant.EventDone {
		t.Fatal("expected done event from [DONE] marker")
	}

	if !reader.Done() {
		t.Error("Done() = false after [DONE] marker, want true")
	}

	// Subsequent calls return nil
	event = reader.Next()
	if event != nil {
		t.Errorf("expected nil after done, got event type %d", event.Type)
	}
}

func TestStreamReaderDoneViaEOF(t *testing.T) {
	// No [DONE] marker, stream just ends
	sse := "data: {\"type\":\"content\",\"text\":\"hi\"}\n\n"

	reader := assistant.NewStreamReader(newSSEReader(sse), simpleContentParser)
	defer reader.Close()

	// Content event
	event := reader.Next()
	if event == nil || event.Type != assistant.EventContentDelta {
		t.Fatal("expected content event")
	}

	// EOF causes done
	event = reader.Next()
	if event == nil {
		t.Fatal("expected done event on EOF")
	}
	if event.Type != assistant.EventDone {
		t.Errorf("Type = %d, want EventDone (%d)", event.Type, assistant.EventDone)
	}

	if !reader.Done() {
		t.Error("Done() = false after EOF, want true")
	}
}

func TestStreamReaderErrorEvent(t *testing.T) {
	sse := "data: {\"type\":\"error\",\"message\":\"rate limited\"}\n\n"

	reader := assistant.NewStreamReader(newSSEReader(sse), simpleContentParser)
	defer reader.Close()

	event := reader.Next()
	if event == nil {
		t.Fatal("expected error event, got nil")
	}
	if event.Type != assistant.EventError {
		t.Errorf("Type = %d, want EventError (%d)", event.Type, assistant.EventError)
	}
	if event.Error == nil {
		t.Fatal("Error is nil")
	}
	if event.Error.Error() != "rate limited" {
		t.Errorf("Error = %q, want %q", event.Error.Error(), "rate limited")
	}

	// After error event, Err() should return the error
	if reader.Err() == nil {
		t.Error("Err() = nil after error event, want non-nil")
	}
}

func TestStreamReaderReadError(t *testing.T) {
	// Use a pipe and close the write end with an error
	pr, pw := io.Pipe()
	pw.CloseWithError(errors.New("connection reset"))

	reader := assistant.NewStreamReader(pr, simpleContentParser)
	defer reader.Close()

	event := reader.Next()
	if event == nil {
		t.Fatal("expected error event, got nil")
	}
	if event.Type != assistant.EventError {
		t.Errorf("Type = %d, want EventError (%d)", event.Type, assistant.EventError)
	}
	if event.Error == nil {
		t.Fatal("Error is nil")
	}

	if reader.Err() == nil {
		t.Error("Err() = nil after read error, want non-nil")
	}
}

func TestStreamReaderCollect(t *testing.T) {
	sse := "data: {\"type\":\"content\",\"text\":\"Hello\"}\n\n" +
		"data: {\"type\":\"content\",\"text\":\" there\"}\n\n" +
		"data: [DONE]\n\n"

	reader := assistant.NewStreamReader(newSSEReader(sse), simpleContentParser)
	defer reader.Close()

	resp, err := reader.Collect()
	if err != nil {
		t.Fatalf("Collect returned error: %v", err)
	}
	if resp.Content != "Hello there" {
		t.Errorf("Content = %q, want %q", resp.Content, "Hello there")
	}
	if resp.FinishReason != "stop" {
		t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "stop")
	}
	if len(resp.ToolCalls) != 0 {
		t.Errorf("ToolCalls length = %d, want 0", len(resp.ToolCalls))
	}
}

func TestStreamReaderCollectWithToolCalls(t *testing.T) {
	sse := "data: {\"type\":\"content\",\"text\":\"Let me search.\"}\n\n" +
		"data: {\"type\":\"tool_start\",\"id\":\"call_1\",\"name\":\"search\"}\n\n" +
		"data: [DONE]\n\n"

	reader := assistant.NewStreamReader(newSSEReader(sse), simpleContentParser)
	defer reader.Close()

	resp, err := reader.Collect()
	if err != nil {
		t.Fatalf("Collect returned error: %v", err)
	}
	if resp.Content != "Let me search." {
		t.Errorf("Content = %q, want %q", resp.Content, "Let me search.")
	}
	if resp.FinishReason != "tool_calls" {
		t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "tool_calls")
	}
	if len(resp.ToolCalls) != 1 {
		t.Fatalf("ToolCalls length = %d, want 1", len(resp.ToolCalls))
	}
	if resp.ToolCalls[0].Name != "search" {
		t.Errorf("ToolCalls[0].Name = %q, want %q", resp.ToolCalls[0].Name, "search")
	}
}

func TestStreamReaderCollectError(t *testing.T) {
	sse := "data: {\"type\":\"content\",\"text\":\"partial\"}\n\n" +
		"data: {\"type\":\"error\",\"message\":\"server error\"}\n\n"

	reader := assistant.NewStreamReader(newSSEReader(sse), simpleContentParser)
	defer reader.Close()

	_, err := reader.Collect()
	if err == nil {
		t.Fatal("Collect should return error")
	}
	if err.Error() != "server error" {
		t.Errorf("error = %q, want %q", err.Error(), "server error")
	}
}

func TestStreamReaderClose(t *testing.T) {
	sse := "data: {\"type\":\"content\",\"text\":\"hello\"}\n\n" +
		"data: {\"type\":\"content\",\"text\":\" world\"}\n\n" +
		"data: [DONE]\n\n"

	reader := assistant.NewStreamReader(newSSEReader(sse), simpleContentParser)

	// Read one event
	event := reader.Next()
	if event == nil || event.Type != assistant.EventContentDelta {
		t.Fatal("expected content event")
	}

	// Close the reader
	err := reader.Close()
	if err != nil {
		t.Fatalf("Close returned error: %v", err)
	}

	// After close, Done should be true
	if !reader.Done() {
		t.Error("Done() = false after Close(), want true")
	}

	// Next should return nil after close
	event = reader.Next()
	if event != nil {
		t.Errorf("expected nil after Close(), got event type %d", event.Type)
	}
}

func TestStreamReaderSkipsComments(t *testing.T) {
	sse := ": this is a comment\n\n" +
		"data: {\"type\":\"content\",\"text\":\"hello\"}\n\n" +
		": another comment\n\n" +
		"data: [DONE]\n\n"

	reader := assistant.NewStreamReader(newSSEReader(sse), simpleContentParser)
	defer reader.Close()

	event := reader.Next()
	if event == nil {
		t.Fatal("expected content event, got nil")
	}
	if event.Type != assistant.EventContentDelta {
		t.Errorf("Type = %d, want EventContentDelta", event.Type)
	}
	if event.Content != "hello" {
		t.Errorf("Content = %q, want %q", event.Content, "hello")
	}

	event = reader.Next()
	if event == nil || event.Type != assistant.EventDone {
		t.Fatal("expected done event")
	}
}

func TestStreamReaderSkipsEmptyLines(t *testing.T) {
	sse := "\n\n\ndata: {\"type\":\"content\",\"text\":\"hi\"}\n\n\n\ndata: [DONE]\n\n"

	reader := assistant.NewStreamReader(newSSEReader(sse), simpleContentParser)
	defer reader.Close()

	event := reader.Next()
	if event == nil {
		t.Fatal("expected content event, got nil")
	}
	if event.Content != "hi" {
		t.Errorf("Content = %q, want %q", event.Content, "hi")
	}
}

func TestStreamReaderContentAccumulation(t *testing.T) {
	sse := "data: {\"type\":\"content\",\"text\":\"A\"}\n\n" +
		"data: {\"type\":\"content\",\"text\":\"B\"}\n\n" +
		"data: {\"type\":\"content\",\"text\":\"C\"}\n\n" +
		"data: [DONE]\n\n"

	reader := assistant.NewStreamReader(newSSEReader(sse), simpleContentParser)
	defer reader.Close()

	// Read first event - Content() should have "A"
	reader.Next()
	if reader.Content() != "A" {
		t.Errorf("Content() after 1st event = %q, want %q", reader.Content(), "A")
	}

	// Read second event - Content() should have "AB"
	reader.Next()
	if reader.Content() != "AB" {
		t.Errorf("Content() after 2nd event = %q, want %q", reader.Content(), "AB")
	}

	// Read third event - Content() should have "ABC"
	reader.Next()
	if reader.Content() != "ABC" {
		t.Errorf("Content() after 3rd event = %q, want %q", reader.Content(), "ABC")
	}
}

func TestStreamReaderToolCallsAccumulation(t *testing.T) {
	sse := "data: {\"type\":\"tool_start\",\"id\":\"call_1\",\"name\":\"fn_a\"}\n\n" +
		"data: {\"type\":\"tool_start\",\"id\":\"call_2\",\"name\":\"fn_b\"}\n\n" +
		"data: [DONE]\n\n"

	reader := assistant.NewStreamReader(newSSEReader(sse), simpleContentParser)
	defer reader.Close()

	reader.Next() // tool_start 1
	calls := reader.ToolCalls()
	if len(calls) != 1 {
		t.Fatalf("ToolCalls after 1st start = %d, want 1", len(calls))
	}

	reader.Next() // tool_start 2
	calls = reader.ToolCalls()
	if len(calls) != 2 {
		t.Fatalf("ToolCalls after 2nd start = %d, want 2", len(calls))
	}
	if calls[0].Name != "fn_a" {
		t.Errorf("ToolCalls[0].Name = %q, want %q", calls[0].Name, "fn_a")
	}
	if calls[1].Name != "fn_b" {
		t.Errorf("ToolCalls[1].Name = %q, want %q", calls[1].Name, "fn_b")
	}
}

func TestStreamReaderDoneAndErrInitialState(t *testing.T) {
	sse := "data: {\"type\":\"content\",\"text\":\"hello\"}\n\ndata: [DONE]\n\n"

	reader := assistant.NewStreamReader(newSSEReader(sse), simpleContentParser)
	defer reader.Close()

	// Before reading anything
	if reader.Done() {
		t.Error("Done() = true before reading, want false")
	}
	if reader.Err() != nil {
		t.Errorf("Err() = %v before reading, want nil", reader.Err())
	}
}

func TestStreamReaderToolCallsReturnsCopy(t *testing.T) {
	sse := "data: {\"type\":\"tool_start\",\"id\":\"call_1\",\"name\":\"fn\"}\n\n" +
		"data: [DONE]\n\n"

	reader := assistant.NewStreamReader(newSSEReader(sse), simpleContentParser)
	defer reader.Close()

	reader.Next()

	calls1 := reader.ToolCalls()
	calls2 := reader.ToolCalls()

	// Modify the first slice and verify the second is unaffected
	if len(calls1) != 1 || len(calls2) != 1 {
		t.Fatal("expected 1 tool call in each slice")
	}
	calls1[0].Name = "modified"
	if calls2[0].Name == "modified" {
		t.Error("ToolCalls() did not return a copy; modifying one slice affected the other")
	}
}

func TestStreamReaderEmptyStream(t *testing.T) {
	sse := "data: [DONE]\n\n"

	reader := assistant.NewStreamReader(newSSEReader(sse), simpleContentParser)
	defer reader.Close()

	event := reader.Next()
	if event == nil || event.Type != assistant.EventDone {
		t.Fatal("expected done event from empty stream")
	}

	if reader.Content() != "" {
		t.Errorf("Content() = %q, want empty string", reader.Content())
	}
	if len(reader.ToolCalls()) != 0 {
		t.Errorf("ToolCalls() length = %d, want 0", len(reader.ToolCalls()))
	}
}

func TestStreamReaderSkipsUnparseableEvents(t *testing.T) {
	sse := "data: {\"type\":\"unknown_garbage\"}\n\n" +
		"data: {\"type\":\"content\",\"text\":\"good\"}\n\n" +
		"data: [DONE]\n\n"

	reader := assistant.NewStreamReader(newSSEReader(sse), simpleContentParser)
	defer reader.Close()

	// The unknown event should be skipped, first readable event is "good"
	event := reader.Next()
	if event == nil {
		t.Fatal("expected content event, got nil")
	}
	if event.Type != assistant.EventContentDelta {
		t.Errorf("Type = %d, want EventContentDelta", event.Type)
	}
	if event.Content != "good" {
		t.Errorf("Content = %q, want %q", event.Content, "good")
	}
}

func TestEventTypeConstants(t *testing.T) {
	// Verify event type constants are distinct
	types := map[assistant.EventType]string{
		assistant.EventContentDelta:  "EventContentDelta",
		assistant.EventToolCallStart: "EventToolCallStart",
		assistant.EventToolCallDelta: "EventToolCallDelta",
		assistant.EventDone:          "EventDone",
		assistant.EventError:         "EventError",
	}
	if len(types) != 5 {
		t.Errorf("expected 5 distinct event types, got %d", len(types))
	}
}