package mcptests

import (
	"fmt"
	"testing"
	"time"

	"github.com/maximhq/bifrost/core/schemas"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// =============================================================================
// MAX DEPTH TESTS - NON-CODE MODE
// =============================================================================

func TestAgent_MaxDepthEnforcement(t *testing.T) {
	t.Parallel()

	clientConfig := GetSampleHTTPClientConfig(GetTestConfig(t).HTTPServerURL)
	clientConfig.ToolsToExecute = []string{"*"}
	clientConfig.ToolsToAutoExecute = []string{"*"}
	manager := setupMCPManager(t, clientConfig)

	// Update tool manager config to set MaxAgentDepth = 5
	manager.UpdateToolManagerConfig(&schemas.MCPToolManagerConfig{
		MaxAgentDepth:        5,
		ToolExecutionTimeout: 30 * time.Second,
	})

	ctx := createTestContext()

	// LLM returns tool calls for 6+ iterations (exceeds max depth of 5)
	mockLLM := &MockLLMCaller{
		chatResponses: []*schemas.BifrostChatResponse{
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-1", "iteration 1"),
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-2", "iteration 2"),
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-3", "iteration 3"),
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-4", "iteration 4"),
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-5", "iteration 5"),
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-6", "iteration 6 - should not reach"),
			}),
			CreateChatResponseWithText("Final - should not reach"),
		},
	}

	initialResponse := mockLLM.chatResponses[0]
	mockLLM.chatCallCount = 1

	originalReq := &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: schemas.Ptr("Long task"),
				},
			},
		},
	}

	result, err := manager.CheckAndExecuteAgentForChatRequest(
		ctx,
		originalReq,
		initialResponse,
		mockLLM.MakeChatRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		},
	)

	require.Nil(t, err)
	require.NotNil(t, result)

	// Agent should stop at depth 5 (made 4 additional calls after initial)
	assert.LessOrEqual(t, mockLLM.chatCallCount, 4, "should stop at max depth")
	t.Logf("Agent stopped at %d iterations (max depth: 5)", mockLLM.chatCallCount)
}

func TestAgent_MaxDepthCustomValue(t *testing.T) {
	t.Parallel()

	clientConfig := GetSampleHTTPClientConfig(GetTestConfig(t).HTTPServerURL)
	clientConfig.ToolsToExecute = []string{"*"}
	clientConfig.ToolsToAutoExecute = []string{"*"}
	manager := setupMCPManager(t, clientConfig)

	// Set MaxAgentDepth = 3
	manager.UpdateToolManagerConfig(&schemas.MCPToolManagerConfig{
		MaxAgentDepth:        3,
		ToolExecutionTimeout: 30 * time.Second,
	})

	ctx := createTestContext()

	mockLLM := &MockLLMCaller{
		chatResponses: []*schemas.BifrostChatResponse{
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-1", "iter 1"),
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-2", "iter 2"),
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-3", "iter 3"),
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-4", "should not reach"),
			}),
		},
	}

	initialResponse := mockLLM.chatResponses[0]
	mockLLM.chatCallCount = 1

	originalReq := &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: schemas.Ptr("Test"),
				},
			},
		},
	}

	result, err := manager.CheckAndExecuteAgentForChatRequest(
		ctx,
		originalReq,
		initialResponse,
		mockLLM.MakeChatRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		},
	)

	require.Nil(t, err)
	require.NotNil(t, result)

	// Should stop at depth 3 (made 2 additional calls after initial)
	assert.LessOrEqual(t, mockLLM.chatCallCount, 2, "should stop at custom max depth of 3")
	t.Logf("Agent stopped at depth 3 with %d follow-up calls", mockLLM.chatCallCount)
}

func TestAgent_MaxDepthReached_ChatFormat(t *testing.T) {
	t.Parallel()

	clientConfig := GetSampleHTTPClientConfig(GetTestConfig(t).HTTPServerURL)
	clientConfig.ToolsToExecute = []string{"*"}
	clientConfig.ToolsToAutoExecute = []string{"*"}
	manager := setupMCPManager(t, clientConfig)

	manager.UpdateToolManagerConfig(&schemas.MCPToolManagerConfig{
		MaxAgentDepth:        2,
		ToolExecutionTimeout: 30 * time.Second,
	}) // Max depth = 2

	ctx := createTestContext()

	mockLLM := &MockLLMCaller{
		chatResponses: []*schemas.BifrostChatResponse{
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-1", "first"),
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-2", "second"),
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-3", "should not reach"),
			}),
		},
	}

	initialResponse := mockLLM.chatResponses[0]
	mockLLM.chatCallCount = 1

	originalReq := &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: schemas.Ptr("Test"),
				},
			},
		},
	}

	result, err := manager.CheckAndExecuteAgentForChatRequest(
		ctx,
		originalReq,
		initialResponse,
		mockLLM.MakeChatRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		},
	)

	require.Nil(t, err)
	require.NotNil(t, result)
	assert.LessOrEqual(t, mockLLM.chatCallCount, 1, "max depth 2 in Chat format")
}

func TestAgent_MaxDepthReached_ResponsesFormat(t *testing.T) {
	t.Parallel()

	clientConfig := GetSampleHTTPClientConfig(GetTestConfig(t).HTTPServerURL)
	clientConfig.ToolsToExecute = []string{"*"}
	clientConfig.ToolsToAutoExecute = []string{"*"}
	manager := setupMCPManager(t, clientConfig)

	manager.UpdateToolManagerConfig(&schemas.MCPToolManagerConfig{
		MaxAgentDepth:        2,
		ToolExecutionTimeout: 30 * time.Second,
	}) // Max depth = 2

	ctx := createTestContext()

	mockLLM := &MockLLMCaller{
		responsesResponses: []*schemas.BifrostResponsesResponse{
			CreateResponsesResponseWithToolCalls([]schemas.ResponsesToolMessage{
				{
					CallID:    schemas.Ptr("call-1"),
					Name: schemas.Ptr("bifrostInternal-echo"),
					Arguments: schemas.Ptr(`{"message": "first"}`),
				},
			}),
			CreateResponsesResponseWithToolCalls([]schemas.ResponsesToolMessage{
				{
					CallID:    schemas.Ptr("call-2"),
					Name: schemas.Ptr("bifrostInternal-echo"),
					Arguments: schemas.Ptr(`{"message": "second"}`),
				},
			}),
			CreateResponsesResponseWithText("Should not reach"),
		},
	}

	initialResponse := mockLLM.responsesResponses[0]
	mockLLM.responsesCallCount = 1

	originalReq := &schemas.BifrostResponsesRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ResponsesMessage{
			{
				Type: schemas.Ptr(schemas.ResponsesMessageTypeMessage),
				Role: schemas.Ptr(schemas.ResponsesInputMessageRoleUser),
				Content: &schemas.ResponsesMessageContent{
					ContentStr: schemas.Ptr("Test"),
				},
			},
		},
	}

	result, err := manager.CheckAndExecuteAgentForResponsesRequest(
		ctx,
		originalReq,
		initialResponse,
		mockLLM.MakeResponsesRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		},
	)

	require.Nil(t, err)
	require.NotNil(t, result)
	assert.LessOrEqual(t, mockLLM.responsesCallCount, 1, "max depth 2 in Responses format")
}

// =============================================================================
// MAX DEPTH TESTS - CODE MODE
// =============================================================================

func TestAgent_MaxDepth_CodeMode(t *testing.T) {
	t.Parallel()

	// Code mode client
	codeModeClient := GetSampleCodeModeClientConfig(t, GetTestConfig(t).HTTPServerURL)
	// Regular HTTP client with tools
	httpClient := GetSampleHTTPClientConfig(GetTestConfig(t).HTTPServerURL)
	httpClient.ID = "mcpserver"
	httpClient.ToolsToExecute = []string{"*"}
	httpClient.ToolsToAutoExecute = []string{"echo"}

	manager := setupMCPManager(t, codeModeClient, httpClient)
	manager.UpdateToolManagerConfig(&schemas.MCPToolManagerConfig{
		MaxAgentDepth:        3,
		ToolExecutionTimeout: 30 * time.Second,
	}) // Max depth = 3

	ctx := createTestContext()

	// Mock LLM that returns executeToolCode calls multiple times
	mockLLM := &MockLLMCaller{
		chatResponses: []*schemas.BifrostChatResponse{
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				{
					ID:   schemas.Ptr("call-1"),
					Type: schemas.Ptr("function"),
					Function: schemas.ChatAssistantMessageToolCallFunction{
						Name:      schemas.Ptr("executeToolCode"),
						Arguments: `{"code": "await mcpserver.echo({message: 'iter 1'}); return 'done1';"}`,
					},
				},
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				{
					ID:   schemas.Ptr("call-2"),
					Type: schemas.Ptr("function"),
					Function: schemas.ChatAssistantMessageToolCallFunction{
						Name:      schemas.Ptr("executeToolCode"),
						Arguments: `{"code": "await mcpserver.echo({message: 'iter 2'}); return 'done2';"}`,
					},
				},
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				{
					ID:   schemas.Ptr("call-3"),
					Type: schemas.Ptr("function"),
					Function: schemas.ChatAssistantMessageToolCallFunction{
						Name:      schemas.Ptr("executeToolCode"),
						Arguments: `{"code": "await mcpserver.echo({message: 'iter 3'}); return 'done3';"}`,
					},
				},
			}),
			CreateChatResponseWithText("Should not reach - max depth hit"),
		},
	}

	initialResponse := mockLLM.chatResponses[0]
	mockLLM.chatCallCount = 1

	originalReq := &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: schemas.Ptr("Code mode test"),
				},
			},
		},
	}

	result, err := manager.CheckAndExecuteAgentForChatRequest(
		ctx,
		originalReq,
		initialResponse,
		mockLLM.MakeChatRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		},
	)

	require.Nil(t, err)
	require.NotNil(t, result)
	assert.LessOrEqual(t, mockLLM.chatCallCount, 2, "max depth should apply to code mode")
	t.Logf("Code mode agent stopped with %d calls", mockLLM.chatCallCount)
}

func TestAgent_MaxDepth_CodeMode_ChatFormat(t *testing.T) {
	t.Parallel()

	codeModeClient := GetSampleCodeModeClientConfig(t, GetTestConfig(t).HTTPServerURL)
	httpClient := GetSampleHTTPClientConfig(GetTestConfig(t).HTTPServerURL)
	httpClient.ID = "server"
	httpClient.ToolsToExecute = []string{"*"}
	httpClient.ToolsToAutoExecute = []string{"*"}

	manager := setupMCPManager(t, codeModeClient, httpClient)
	manager.UpdateToolManagerConfig(&schemas.MCPToolManagerConfig{
		MaxAgentDepth:        2,
		ToolExecutionTimeout: 30 * time.Second,
	})

	ctx := createTestContext()

	mockLLM := &MockLLMCaller{
		chatResponses: []*schemas.BifrostChatResponse{
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				{
					ID:   schemas.Ptr("call-1"),
					Type: schemas.Ptr("function"),
					Function: schemas.ChatAssistantMessageToolCallFunction{
						Name:      schemas.Ptr("executeToolCode"),
						Arguments: `{"code": "return 'test1';"}`,
					},
				},
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				{
					ID:   schemas.Ptr("call-2"),
					Type: schemas.Ptr("function"),
					Function: schemas.ChatAssistantMessageToolCallFunction{
						Name:      schemas.Ptr("executeToolCode"),
						Arguments: `{"code": "return 'test2';"}`,
					},
				},
			}),
			CreateChatResponseWithText("Done"),
		},
	}

	initialResponse := mockLLM.chatResponses[0]
	mockLLM.chatCallCount = 1

	originalReq := &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: schemas.Ptr("Test"),
				},
			},
		},
	}

	result, err := manager.CheckAndExecuteAgentForChatRequest(
		ctx,
		originalReq,
		initialResponse,
		mockLLM.MakeChatRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		},
	)

	require.Nil(t, err)
	require.NotNil(t, result)
	t.Logf("Code mode Chat format test completed")
}

func TestAgent_MaxDepth_CodeMode_ResponsesFormat(t *testing.T) {
	t.Parallel()

	codeModeClient := GetSampleCodeModeClientConfig(t, GetTestConfig(t).HTTPServerURL)
	httpClient := GetSampleHTTPClientConfig(GetTestConfig(t).HTTPServerURL)
	httpClient.ID = "server"
	httpClient.ToolsToExecute = []string{"*"}
	httpClient.ToolsToAutoExecute = []string{"*"}

	manager := setupMCPManager(t, codeModeClient, httpClient)
	manager.UpdateToolManagerConfig(&schemas.MCPToolManagerConfig{
		MaxAgentDepth:        2,
		ToolExecutionTimeout: 30 * time.Second,
	})

	ctx := createTestContext()

	mockLLM := &MockLLMCaller{
		responsesResponses: []*schemas.BifrostResponsesResponse{
			CreateResponsesResponseWithToolCalls([]schemas.ResponsesToolMessage{
				{
					CallID:    schemas.Ptr("call-1"),
					Name:      schemas.Ptr("executeToolCode"),
					Arguments: schemas.Ptr(`{"code": "return 'test1';"}`),
				},
			}),
			CreateResponsesResponseWithToolCalls([]schemas.ResponsesToolMessage{
				{
					CallID:    schemas.Ptr("call-2"),
					Name:      schemas.Ptr("executeToolCode"),
					Arguments: schemas.Ptr(`{"code": "return 'test2';"}`),
				},
			}),
			CreateResponsesResponseWithText("Done"),
		},
	}

	initialResponse := mockLLM.responsesResponses[0]
	mockLLM.responsesCallCount = 1

	originalReq := &schemas.BifrostResponsesRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ResponsesMessage{
			{
				Type: schemas.Ptr(schemas.ResponsesMessageTypeMessage),
				Role: schemas.Ptr(schemas.ResponsesInputMessageRoleUser),
				Content: &schemas.ResponsesMessageContent{
					ContentStr: schemas.Ptr("Test"),
				},
			},
		},
	}

	result, err := manager.CheckAndExecuteAgentForResponsesRequest(
		ctx,
		originalReq,
		initialResponse,
		mockLLM.MakeResponsesRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		},
	)

	require.Nil(t, err)
	require.NotNil(t, result)
	t.Logf("Code mode Responses format test completed")
}

// =============================================================================
// AGENT TIMEOUT TESTS - NON-CODE MODE
// =============================================================================

func TestAgent_Timeout(t *testing.T) {
	t.Parallel()

	// Test that agent loop MUST timeout by creating a tool that takes longer than the timeout
	manager := setupMCPManager(t)

	// Register a slow tool that takes 500ms (longer than our 200ms timeout)
	slowToolSchema := schemas.ChatTool{
		Type: schemas.ChatToolTypeFunction,
		Function: &schemas.ChatToolFunction{
			Name:        "slow_tool",
			Description: schemas.Ptr("A tool that takes a long time"),
			Parameters: &schemas.ToolFunctionParameters{
				Type:       "object",
				Properties: schemas.NewOrderedMap(),
			},
		},
	}

	err := manager.RegisterTool(
		"slow_tool",
		"A tool that takes a long time",
		func(args any) (string, error) {
			// This will definitely exceed the 200ms timeout
			time.Sleep(500 * time.Millisecond)
			return `{"result": "should not reach here"}`, nil
		},
		slowToolSchema,
	)
	require.NoError(t, err)

	err = SetInternalClientAutoExecute(manager, []string{"slow_tool"})
	require.NoError(t, err)

	// Timeout set to 200ms - tool takes 500ms, so it MUST timeout
	ctx, cancel := createTestContextWithTimeout(200 * time.Millisecond)
	defer cancel()

	mockLLM := &MockLLMCaller{
		chatResponses: []*schemas.BifrostChatResponse{
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				CreateInProcessToolCall("call-1", "slow_tool", map[string]interface{}{}),
			}),
		},
	}

	initialResponse := mockLLM.chatResponses[0]

	originalReq := &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4o",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: schemas.Ptr("Test timeout"),
				},
			},
		},
	}

	// Agent MUST timeout since tool takes 500ms but timeout is 200ms
	_, bifrostErr := manager.CheckAndExecuteAgentForChatRequest(
		ctx,
		originalReq,
		initialResponse,
		mockLLM.MakeChatRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		},
	)

	// MUST have timeout error
	require.NotNil(t, bifrostErr, "Expected timeout error but got success - timeout is not being enforced!")
	t.Logf("✅ Timeout correctly enforced: %v", bifrostErr.Error)
}

func TestAgent_TimeoutDuringExecution(t *testing.T) {
	t.Parallel()

	// Test that timeout is enforced DURING tool execution (not just between iterations)
	// Tool takes 1 second, timeout is 150ms - MUST timeout mid-execution
	manager := setupMCPManager(t)

	slowToolSchema := schemas.ChatTool{
		Type: schemas.ChatToolTypeFunction,
		Function: &schemas.ChatToolFunction{
			Name:        "very_slow_tool",
			Description: schemas.Ptr("A tool that takes 1 full second"),
			Parameters: &schemas.ToolFunctionParameters{
				Type:       "object",
				Properties: schemas.NewOrderedMap(),
			},
		},
	}

	err := manager.RegisterTool(
		"very_slow_tool",
		"A tool that takes 1 full second",
		func(args any) (string, error) {
			// This takes 1 second - much longer than 150ms timeout
			time.Sleep(1 * time.Second)
			return `{"result": "should never complete"}`, nil
		},
		slowToolSchema,
	)
	require.NoError(t, err)

	err = SetInternalClientAutoExecute(manager, []string{"very_slow_tool"})
	require.NoError(t, err)

	// Timeout is 150ms, tool takes 1000ms - MUST timeout during execution
	ctx, cancel := createTestContextWithTimeout(150 * time.Millisecond)
	defer cancel()

	mockLLM := &MockLLMCaller{
		chatResponses: []*schemas.BifrostChatResponse{
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				CreateInProcessToolCall("call-1", "very_slow_tool", map[string]interface{}{}),
			}),
		},
	}

	_, bifrostErr := manager.CheckAndExecuteAgentForChatRequest(
		ctx,
		&schemas.BifrostChatRequest{
			Provider: schemas.OpenAI,
			Model:    "gpt-4o",
			Input:    []schemas.ChatMessage{{Role: schemas.ChatMessageRoleUser, Content: &schemas.ChatMessageContent{ContentStr: schemas.Ptr("Test")}}},
		},
		mockLLM.chatResponses[0],
		mockLLM.MakeChatRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		},
	)

	// MUST timeout during execution
	require.NotNil(t, bifrostErr, "Expected timeout during tool execution but got success - timeout not enforced mid-execution!")
	t.Logf("✅ Timeout during execution correctly enforced: %v", bifrostErr.Error)
}

func TestAgent_Timeout_ChatFormat(t *testing.T) {
	t.Parallel()

	// Chat format MUST timeout - tool takes 400ms, timeout is 150ms
	manager := setupMCPManager(t)

	slowToolSchema := schemas.ChatTool{
		Type: schemas.ChatToolTypeFunction,
		Function: &schemas.ChatToolFunction{
			Name:        "slow_chat_tool",
			Description: schemas.Ptr("Tool for chat format timeout test"),
			Parameters:  &schemas.ToolFunctionParameters{Type: "object", Properties: schemas.NewOrderedMap()},
		},
	}

	err := manager.RegisterTool("slow_chat_tool", "Tool for timeout test",
		func(args any) (string, error) {
			time.Sleep(400 * time.Millisecond) // Longer than 150ms timeout
			return `{"status": "should not complete"}`, nil
		}, slowToolSchema)
	require.NoError(t, err)

	err = SetInternalClientAutoExecute(manager, []string{"slow_chat_tool"})
	require.NoError(t, err)

	ctx, cancel := createTestContextWithTimeout(150 * time.Millisecond)
	defer cancel()

	mockLLM := &MockLLMCaller{
		chatResponses: []*schemas.BifrostChatResponse{
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				CreateInProcessToolCall("call-1", "slow_chat_tool", map[string]interface{}{}),
			}),
		},
	}

	_, bifrostErr := manager.CheckAndExecuteAgentForChatRequest(ctx,
		&schemas.BifrostChatRequest{Provider: schemas.OpenAI, Model: "gpt-4o",
			Input: []schemas.ChatMessage{{Role: schemas.ChatMessageRoleUser, Content: &schemas.ChatMessageContent{ContentStr: schemas.Ptr("Test")}}}},
		mockLLM.chatResponses[0], mockLLM.MakeChatRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		})

	require.NotNil(t, bifrostErr, "Chat format timeout not enforced!")
	t.Logf("✅ Chat format timeout enforced: %v", bifrostErr.Error)
}

func TestAgent_Timeout_ResponsesFormat(t *testing.T) {
	t.Parallel()

	// Responses format MUST timeout - tool takes 400ms, timeout is 150ms
	manager := setupMCPManager(t)

	slowToolSchema := schemas.ChatTool{
		Type: schemas.ChatToolTypeFunction,
		Function: &schemas.ChatToolFunction{
			Name:        "slow_responses_tool",
			Description: schemas.Ptr("Tool for responses format timeout test"),
			Parameters:  &schemas.ToolFunctionParameters{Type: "object", Properties: schemas.NewOrderedMap()},
		},
	}

	err := manager.RegisterTool("slow_responses_tool", "Tool for timeout test",
		func(args any) (string, error) {
			time.Sleep(400 * time.Millisecond) // Longer than 150ms timeout
			return `{"status": "should not complete"}`, nil
		}, slowToolSchema)
	require.NoError(t, err)

	err = SetInternalClientAutoExecute(manager, []string{"slow_responses_tool"})
	require.NoError(t, err)

	ctx, cancel := createTestContextWithTimeout(150 * time.Millisecond)
	defer cancel()

	mockLLM := &MockLLMCaller{
		responsesResponses: []*schemas.BifrostResponsesResponse{
			CreateResponsesResponseWithToolCalls([]schemas.ResponsesToolMessage{
				{CallID: schemas.Ptr("call-1"), Name: schemas.Ptr("bifrostInternal-slow_responses_tool"), Arguments: schemas.Ptr(`{}`)},
			}),
		},
	}

	_, bifrostErr := manager.CheckAndExecuteAgentForResponsesRequest(ctx,
		&schemas.BifrostResponsesRequest{Provider: schemas.OpenAI, Model: "gpt-4o",
			Input: []schemas.ResponsesMessage{{Type: schemas.Ptr(schemas.ResponsesMessageTypeMessage), Role: schemas.Ptr(schemas.ResponsesInputMessageRoleUser),
				Content: &schemas.ResponsesMessageContent{ContentStr: schemas.Ptr("Test")}}}},
		mockLLM.responsesResponses[0], mockLLM.MakeResponsesRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		})

	require.NotNil(t, bifrostErr, "Responses format timeout not enforced!")
	t.Logf("✅ Responses format timeout enforced: %v", bifrostErr.Error)
}

// =============================================================================
// ERROR PROPAGATION TESTS
// =============================================================================

func TestAgent_ErrorPropagation(t *testing.T) {
	t.Parallel()

	clientConfig := GetSampleHTTPClientConfig(GetTestConfig(t).HTTPServerURL)
	clientConfig.ToolsToExecute = []string{"*"}
	clientConfig.ToolsToAutoExecute = []string{"*"}
	manager := setupMCPManager(t, clientConfig)

	ctx := createTestContext()

	// Mock a tool that will return an error
	mockLLM := &MockLLMCaller{
		chatResponses: []*schemas.BifrostChatResponse{
			// Call a non-existent tool
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				{
					ID:   schemas.Ptr("call-error"),
					Type: schemas.Ptr("function"),
					Function: schemas.ChatAssistantMessageToolCallFunction{
						Name: schemas.Ptr("bifrostInternal-nonexistent_tool"),
						Arguments: `{}`,
					},
				},
			}),
			CreateChatResponseWithText("Handled error"),
		},
	}

	initialResponse := mockLLM.chatResponses[0]
	mockLLM.chatCallCount = 1

	originalReq := &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: schemas.Ptr("Test error"),
				},
			},
		},
	}

	result, err := manager.CheckAndExecuteAgentForChatRequest(
		ctx,
		originalReq,
		initialResponse,
		mockLLM.MakeChatRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		},
	)

	// Error should be propagated, or tool result should contain error
	// The exact behavior depends on implementation
	if err != nil {
		t.Logf("Error propagated: %v", err)
	} else {
		require.NotNil(t, result)
		t.Logf("Error handled in response")
	}
}

func TestAgent_ErrorInMiddleOfLoop(t *testing.T) {
	t.Parallel()

	clientConfig := GetSampleHTTPClientConfig(GetTestConfig(t).HTTPServerURL)
	clientConfig.ToolsToExecute = []string{"*"}
	clientConfig.ToolsToAutoExecute = []string{"*"}
	manager := setupMCPManager(t, clientConfig)

	ctx := createTestContext()

	mockLLM := &MockLLMCaller{
		chatResponses: []*schemas.BifrostChatResponse{
			// First tool succeeds
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-1", "success"),
			}),
			// Second tool has error
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				{
					ID:   schemas.Ptr("call-2"),
					Type: schemas.Ptr("function"),
					Function: schemas.ChatAssistantMessageToolCallFunction{
						Name: schemas.Ptr("bifrostInternal-nonexistent"),
						Arguments: `{}`,
					},
				},
			}),
			CreateChatResponseWithText("Recovered"),
		},
	}

	initialResponse := mockLLM.chatResponses[0]
	mockLLM.chatCallCount = 1

	originalReq := &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: schemas.Ptr("Multi-step with error"),
				},
			},
		},
	}

	result, err := manager.CheckAndExecuteAgentForChatRequest(
		ctx,
		originalReq,
		initialResponse,
		mockLLM.MakeChatRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		},
	)

	// First tool should have executed successfully
	// Error in second tool should be handled
	if err != nil {
		t.Logf("Error in middle of loop: %v", err)
	} else {
		require.NotNil(t, result)
		t.Logf("Agent handled error in middle of loop")
	}
}

func TestAgent_LLMError(t *testing.T) {
	t.Parallel()

	clientConfig := GetSampleHTTPClientConfig(GetTestConfig(t).HTTPServerURL)
	clientConfig.ToolsToExecute = []string{"*"}
	clientConfig.ToolsToAutoExecute = []string{"*"}
	manager := setupMCPManager(t, clientConfig)

	ctx := createTestContext()

	mockLLM := &MockLLMCaller{
		chatResponses: []*schemas.BifrostChatResponse{
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-1", "first"),
			}),
			// Next call will error (no more responses)
		},
	}

	initialResponse := mockLLM.chatResponses[0]
	mockLLM.chatCallCount = 1

	originalReq := &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: schemas.Ptr("Test LLM error"),
				},
			},
		},
	}

	result, err := manager.CheckAndExecuteAgentForChatRequest(
		ctx,
		originalReq,
		initialResponse,
		mockLLM.MakeChatRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		},
	)

	// LLM error should be returned
	if err != nil {
		assert.Contains(t, err.Error.Message, "no more mock", "should get LLM error")
		t.Logf("LLM error correctly propagated: %s", err.Error.Message)
	} else {
		t.Logf("LLM error handled gracefully, result: %+v", result)
	}
}

// =============================================================================
// COMBINED LIMITS TESTS
// =============================================================================

func TestAgent_MaxDepthAndTimeout(t *testing.T) {
	t.Parallel()

	clientConfig := GetSampleHTTPClientConfig(GetTestConfig(t).HTTPServerURL)
	clientConfig.ToolsToExecute = []string{"*"}
	clientConfig.ToolsToAutoExecute = []string{"*"}
	manager := setupMCPManager(t, clientConfig)

	// Set both max depth and timeout
	manager.UpdateToolManagerConfig(&schemas.MCPToolManagerConfig{
		MaxAgentDepth:        3,
		ToolExecutionTimeout: 5 * time.Second,
	}) // Max depth = 3, timeout = 5 seconds

	ctx := createTestContext()

	mockLLM := &MockLLMCaller{
		chatResponses: []*schemas.BifrostChatResponse{
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-1", "iter 1"),
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-2", "iter 2"),
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-3", "iter 3"),
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-4", "should not reach"),
			}),
		},
	}

	initialResponse := mockLLM.chatResponses[0]
	mockLLM.chatCallCount = 1

	originalReq := &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: schemas.Ptr("Test combined limits"),
				},
			},
		},
	}

	result, err := manager.CheckAndExecuteAgentForChatRequest(
		ctx,
		originalReq,
		initialResponse,
		mockLLM.MakeChatRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		},
	)

	// Whichever limit hits first should stop the agent
	// In this case, max depth should hit first
	if err != nil {
		t.Logf("Agent stopped with error: %v", err)
	} else {
		require.NotNil(t, result)
		assert.LessOrEqual(t, mockLLM.chatCallCount, 2, "should stop at max depth 3")
		t.Logf("Agent stopped at %d calls (max depth: 3)", mockLLM.chatCallCount)
	}
}

// =============================================================================
// EDGE CASE TESTS
// =============================================================================

func TestAgent_MaxDepthZero(t *testing.T) {
	t.Parallel()

	clientConfig := GetSampleHTTPClientConfig(GetTestConfig(t).HTTPServerURL)
	clientConfig.ToolsToExecute = []string{"*"}
	clientConfig.ToolsToAutoExecute = []string{"*"}
	manager := setupMCPManager(t, clientConfig)

	// Set max depth = 0 (should not allow any iterations)
	manager.UpdateToolManagerConfig(&schemas.MCPToolManagerConfig{
		MaxAgentDepth:        0,
		ToolExecutionTimeout: 30 * time.Second,
	})

	ctx := createTestContext()

	mockLLM := &MockLLMCaller{
		chatResponses: []*schemas.BifrostChatResponse{
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-1", "should not execute"),
			}),
		},
	}

	initialResponse := mockLLM.chatResponses[0]
	mockLLM.chatCallCount = 0

	originalReq := &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: schemas.Ptr("Test zero depth"),
				},
			},
		},
	}

	result, err := manager.CheckAndExecuteAgentForChatRequest(
		ctx,
		originalReq,
		initialResponse,
		mockLLM.MakeChatRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		},
	)

	// Should return immediately with tool calls
	require.Nil(t, err)
	require.NotNil(t, result)
	assert.Equal(t, 0, mockLLM.chatCallCount, "should not make any LLM calls with max depth 0")

	// Should return the tools for approval
	finalMessage := result.Choices[0].ChatNonStreamResponseChoice.Message
	require.NotNil(t, finalMessage.ChatAssistantMessage)
	require.NotEmpty(t, finalMessage.ChatAssistantMessage.ToolCalls)
	t.Logf("Max depth 0: correctly returned tools without execution")
}

func TestAgent_ParallelToolExecution(t *testing.T) {
	t.Parallel()

	clientConfig := GetSampleHTTPClientConfig(GetTestConfig(t).HTTPServerURL)
	clientConfig.ToolsToExecute = []string{"*"}
	clientConfig.ToolsToAutoExecute = []string{"*"}
	manager := setupMCPManager(t, clientConfig)

	ctx := createTestContext()

	// LLM returns multiple tools in parallel
	mockLLM := &MockLLMCaller{
		chatResponses: []*schemas.BifrostChatResponse{
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-1", "parallel 1"),
				GetSampleEchoToolCall("call-2", "parallel 2"),
				GetSampleEchoToolCall("call-3", "parallel 3"),
			}),
			CreateChatResponseWithText("All parallel tools executed"),
		},
	}

	initialResponse := mockLLM.chatResponses[0]
	mockLLM.chatCallCount = 1

	originalReq := &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: schemas.Ptr("Parallel test"),
				},
			},
		},
	}

	result, err := manager.CheckAndExecuteAgentForChatRequest(
		ctx,
		originalReq,
		initialResponse,
		mockLLM.MakeChatRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		},
	)

	require.Nil(t, err)
	require.NotNil(t, result)

	// All 3 tools should be executed in parallel in one iteration
	assert.GreaterOrEqual(t, mockLLM.chatCallCount, 1, "should continue after parallel execution")
	t.Logf("Parallel tool execution completed successfully")
}

func TestAgent_IterationTracking(t *testing.T) {
	t.Parallel()

	clientConfig := GetSampleHTTPClientConfig(GetTestConfig(t).HTTPServerURL)
	clientConfig.ToolsToExecute = []string{"*"}
	clientConfig.ToolsToAutoExecute = []string{"*"}
	manager := setupMCPManager(t, clientConfig)

	manager.UpdateToolManagerConfig(&schemas.MCPToolManagerConfig{
		MaxAgentDepth:        10,
		ToolExecutionTimeout: 30 * time.Second,
	})

	ctx := createTestContext()

	iterationCount := 0
	trackingMockLLM := &MockLLMCaller{
		chatResponses: []*schemas.BifrostChatResponse{
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-1", fmt.Sprintf("iteration %d", iterationCount)),
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-2", fmt.Sprintf("iteration %d", iterationCount+1)),
			}),
			CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
				GetSampleEchoToolCall("call-3", fmt.Sprintf("iteration %d", iterationCount+2)),
			}),
			CreateChatResponseWithText("Done with iterations"),
		},
	}

	initialResponse := trackingMockLLM.chatResponses[0]
	trackingMockLLM.chatCallCount = 1

	originalReq := &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: schemas.Ptr("Track iterations"),
				},
			},
		},
	}

	result, err := manager.CheckAndExecuteAgentForChatRequest(
		ctx,
		originalReq,
		initialResponse,
		trackingMockLLM.MakeChatRequest,
		func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
			return manager.ExecuteToolCall(ctx, request)
		},
	)

	require.Nil(t, err)
	require.NotNil(t, result)

	// Track actual iterations
	actualIterations := trackingMockLLM.chatCallCount
	t.Logf("Agent completed with %d iterations", actualIterations)
	assert.GreaterOrEqual(t, actualIterations, 1, "should track iterations")
	assert.LessOrEqual(t, actualIterations, 3, "should not exceed expected iterations")
}