package openapi

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestSessionToolTurns(t *testing.T) {
	callCount := 0
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		callCount++
		var resp ChatCompletionResponse

		if callCount == 1 {
			// First call: Model decides to call a tool
			resp = ChatCompletionResponse{
				ID: "resp1",
				Choices: []Choice{
					{
						Message: Message{
							Role: RoleAssistant,
							ToolCalls: []ToolCall{
								{
									ID:   "call_1",
									Type: "function",
									Function: ToolCallFunction{
										Name:      "get_weather",
										Arguments: `{"location":"Stockholm"}`,
									},
								},
							},
						},
						FinishReason: "tool_calls",
					},
				},
			}
		} else {
			// Second call: Model gives final response after tool result
			resp = ChatCompletionResponse{
				ID: "resp2",
				Choices: []Choice{
					{
						Message: Message{
							Role:    RoleAssistant,
							Content: "The weather in Stockholm is sunny.",
						},
						FinishReason: "stop",
					},
				},
			}
		}
		json.NewEncoder(w).Encode(resp)
	}))
	defer server.Close()

	client := NewClient(server.URL, "test-key")
	session := NewSession(client, "test-model")

	type weatherArgs struct {
		Location string `json:"location"`
	}

	toolExecuted := false
	session.RegisterTool("get_weather", "Get weather", 
		map[string]any{"type": "object", "properties": map[string]any{"location": map[string]any{"type": "string"}}},
		func(ctx context.Context, s *Session, args json.RawMessage) (any, error) {
			var w weatherArgs
			json.Unmarshal(args, &w)
			if w.Location != "Stockholm" {
				t.Errorf("expected Stockholm, got %s", w.Location)
			}
			toolExecuted = true
			return map[string]any{"temp": 20, "condition": "sunny"}, nil
		})

	hooksCalled := make(map[string]bool)
	session.SetHooks(SessionHooks{
		BeforeMessage: func(ctx context.Context, s *Session, msg *Message) error {
			hooksCalled["BeforeMessage"] = true
			return nil
		},
		AfterMessage: func(ctx context.Context, s *Session, msg *Message) error {
			hooksCalled["AfterMessage"] = true
			return nil
		},
		BeforeToolCall: func(ctx context.Context, s *Session, tc *ToolCall) error {
			hooksCalled["BeforeToolCall"] = true
			return nil
		},
		AfterToolCall: func(ctx context.Context, s *Session, tc *ToolCall, result any, err error) error {
			hooksCalled["AfterToolCall"] = true
			return nil
		},
		OnDone: func(ctx context.Context, s *Session) error {
			hooksCalled["OnDone"] = true
			return nil
		},
	})

	ctx := context.Background()
	chunks, errs := session.Chat(ctx, "How is the weather in Stockholm?")

	var finalContent string
	for chunk := range chunks {
		if len(chunk.Choices) > 0 {
			finalContent += chunk.Choices[0].Delta.Content
		}
	}

	if err := <-errs; err != nil {
		t.Fatal(err)
	}

	if !toolExecuted {
		t.Error("expected tool to be executed")
	}

	hookNames := []string{"BeforeMessage", "AfterMessage", "BeforeToolCall", "AfterToolCall", "OnDone"}
	for _, name := range hookNames {
		if !hooksCalled[name] {
			t.Errorf("expected hook %s to be called", name)
		}
	}

	if finalContent != "The weather in Stockholm is sunny." {
		t.Errorf("expected final response, got %s", finalContent)
	}

	if len(session.Messages()) != 4 {
		t.Errorf("expected 4 messages in history, got %d", len(session.Messages()))
	}
}