bifrost.go

// Package bifrost provides the core implementation of the Bifrost system.
// Bifrost is a unified interface for interacting with various AI model providers,
// managing concurrent requests, and handling provider-specific configurations.
package bifrost

import (
	"context"
	"fmt"
	"math/rand"
	"slices"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/google/uuid"
	"github.com/maximhq/bifrost/core/providers/anthropic"
	"github.com/maximhq/bifrost/core/providers/azure"
	"github.com/maximhq/bifrost/core/providers/bedrock"
	"github.com/maximhq/bifrost/core/providers/cerebras"
	"github.com/maximhq/bifrost/core/providers/cohere"
	"github.com/maximhq/bifrost/core/providers/elevenlabs"
	"github.com/maximhq/bifrost/core/providers/gemini"
	"github.com/maximhq/bifrost/core/providers/groq"
	"github.com/maximhq/bifrost/core/providers/mistral"
	"github.com/maximhq/bifrost/core/providers/ollama"
	"github.com/maximhq/bifrost/core/providers/openai"
	"github.com/maximhq/bifrost/core/providers/openrouter"
	"github.com/maximhq/bifrost/core/providers/parasail"
	"github.com/maximhq/bifrost/core/providers/perplexity"
	"github.com/maximhq/bifrost/core/providers/sgl"
	providerUtils "github.com/maximhq/bifrost/core/providers/utils"
	"github.com/maximhq/bifrost/core/providers/vertex"
	schemas "github.com/maximhq/bifrost/core/schemas"
)

// ChannelMessage represents a message passed through the request channel.
// It contains the request, response and error channels, and the request type.
type ChannelMessage struct {
	schemas.BifrostRequest
	Context        context.Context
	Response       chan *schemas.BifrostResponse
	ResponseStream chan chan *schemas.BifrostStream
	Err            chan schemas.BifrostError
}

// Bifrost manages providers and maintains specified open channels for concurrent processing.
// It handles request routing, provider management, and response processing.
type Bifrost struct {
	ctx                 context.Context
	cancel              context.CancelFunc
	account             schemas.Account                    // account interface
	plugins             atomic.Pointer[[]schemas.Plugin]   // list of plugins
	providers           atomic.Pointer[[]schemas.Provider] // list of providers
	requestQueues       sync.Map                           // provider request queues (thread-safe)
	waitGroups          sync.Map                           // wait groups for each provider (thread-safe)
	providerMutexes     sync.Map                           // mutexes for each provider to prevent concurrent updates (thread-safe)
	channelMessagePool  sync.Pool                          // Pool for ChannelMessage objects, initial pool size is set in Init
	responseChannelPool sync.Pool                          // Pool for response channels, initial pool size is set in Init
	errorChannelPool    sync.Pool                          // Pool for error channels, initial pool size is set in Init
	responseStreamPool  sync.Pool                          // Pool for response stream channels, initial pool size is set in Init
	pluginPipelinePool  sync.Pool                          // Pool for PluginPipeline objects
	bifrostRequestPool  sync.Pool                          // Pool for BifrostRequest objects
	logger              schemas.Logger                     // logger instance, default logger is used if not provided
	mcpManager          *MCPManager                        // MCP integration manager (nil if MCP not configured)
	dropExcessRequests  atomic.Bool                        // If true, in cases where the queue is full, requests will not wait for the queue to be empty and will be dropped instead.
	keySelector         schemas.KeySelector                // Custom key selector function
}

// PluginPipeline encapsulates the execution of plugin PreHooks and PostHooks, tracks how many plugins ran, and manages short-circuiting and error aggregation.
type PluginPipeline struct {
	plugins []schemas.Plugin
	logger  schemas.Logger

	// Number of PreHooks that were executed (used to determine which PostHooks to run in reverse order)
	executedPreHooks int
	// Errors from PreHooks and PostHooks
	preHookErrors  []error
	postHookErrors []error
}

// Global logger instance which is set in the Init function
var logger schemas.Logger

// INITIALIZATION

// Init initializes a new Bifrost instance with the given configuration.
// It sets up the account, plugins, object pools, and initializes providers.
// Returns an error if initialization fails.
// Initial Memory Allocations happens here as per the initial pool size.
func Init(ctx context.Context, config schemas.BifrostConfig) (*Bifrost, error) {
	if config.Account == nil {
		return nil, fmt.Errorf("account is required to initialize Bifrost")
	}

	if config.Logger == nil {
		config.Logger = NewDefaultLogger(schemas.LogLevelInfo)
	}

	providerUtils.SetLogger(config.Logger)
	bifrostCtx, cancel := context.WithCancel(ctx)
	bifrost := &Bifrost{
		ctx:           bifrostCtx,
		cancel:        cancel,
		account:       config.Account,
		plugins:       atomic.Pointer[[]schemas.Plugin]{},
		requestQueues: sync.Map{},
		waitGroups:    sync.Map{},
		keySelector:   config.KeySelector,
		logger:        config.Logger,
	}
	bifrost.plugins.Store(&config.Plugins)

	// Initialize providers slice
	bifrost.providers.Store(&[]schemas.Provider{})

	bifrost.dropExcessRequests.Store(config.DropExcessRequests)

	if bifrost.keySelector == nil {
		bifrost.keySelector = WeightedRandomKeySelector
	}

	// Initialize object pools
	bifrost.channelMessagePool = sync.Pool{
		New: func() interface{} {
			return &ChannelMessage{}
		},
	}
	bifrost.responseChannelPool = sync.Pool{
		New: func() interface{} {
			return make(chan *schemas.BifrostResponse, 1)
		},
	}
	bifrost.errorChannelPool = sync.Pool{
		New: func() interface{} {
			return make(chan schemas.BifrostError, 1)
		},
	}
	bifrost.responseStreamPool = sync.Pool{
		New: func() interface{} {
			return make(chan chan *schemas.BifrostStream, 1)
		},
	}
	bifrost.pluginPipelinePool = sync.Pool{
		New: func() interface{} {
			return &PluginPipeline{
				preHookErrors:  make([]error, 0),
				postHookErrors: make([]error, 0),
			}
		},
	}
	bifrost.bifrostRequestPool = sync.Pool{
		New: func() interface{} {
			return &schemas.BifrostRequest{}
		},
	}
	// Prewarm pools with multiple objects
	for range config.InitialPoolSize {
		// Create and put new objects directly into pools
		bifrost.channelMessagePool.Put(&ChannelMessage{})
		bifrost.responseChannelPool.Put(make(chan *schemas.BifrostResponse, 1))
		bifrost.errorChannelPool.Put(make(chan schemas.BifrostError, 1))
		bifrost.responseStreamPool.Put(make(chan chan *schemas.BifrostStream, 1))
		bifrost.pluginPipelinePool.Put(&PluginPipeline{
			preHookErrors:  make([]error, 0),
			postHookErrors: make([]error, 0),
		})
		bifrost.bifrostRequestPool.Put(&schemas.BifrostRequest{})
	}

	providerKeys, err := bifrost.account.GetConfiguredProviders()
	if err != nil {
		return nil, err
	}

	// Initialize MCP manager if configured
	if config.MCPConfig != nil {
		mcpManager, err := newMCPManager(bifrostCtx, *config.MCPConfig, bifrost.logger)
		if err != nil {
			bifrost.logger.Warn(fmt.Sprintf("failed to initialize MCP manager: %v", err))
		} else {
			bifrost.mcpManager = mcpManager
			bifrost.logger.Info("MCP integration initialized successfully")
		}
	}

	// Create buffered channels for each provider and start workers
	for _, providerKey := range providerKeys {
		if strings.TrimSpace(string(providerKey)) == "" {
			bifrost.logger.Warn("provider key is empty, skipping init")
			continue
		}

		config, err := bifrost.account.GetConfigForProvider(providerKey)
		if err != nil {
			bifrost.logger.Warn(fmt.Sprintf("failed to get config for provider, skipping init: %v", err))
			continue
		}
		if config == nil {
			bifrost.logger.Warn(fmt.Sprintf("config is nil for provider %s, skipping init", providerKey))
			continue
		}

		// Lock the provider mutex during initialization
		providerMutex := bifrost.getProviderMutex(providerKey)
		providerMutex.Lock()
		err = bifrost.prepareProvider(providerKey, config)
		providerMutex.Unlock()

		if err != nil {
			bifrost.logger.Warn(fmt.Sprintf("failed to prepare provider %s: %v", providerKey, err))
		}
	}

	// Set logger
	logger = bifrost.logger

	return bifrost, nil
}

// ReloadConfig reloads the config from DB
// Currently we only update account and drop excess requests
// We will keep on adding other aspects as required
func (bifrost *Bifrost) ReloadConfig(config schemas.BifrostConfig) error {
	bifrost.dropExcessRequests.Store(config.DropExcessRequests)
	return nil
}

// PUBLIC API METHODS

// ListModelsRequest sends a list models request to the specified provider.
func (bifrost *Bifrost) ListModelsRequest(ctx context.Context, req *schemas.BifrostListModelsRequest) (*schemas.BifrostListModelsResponse, *schemas.BifrostError) {
	if req == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "list models request is nil",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.ListModelsRequest,
			},
		}
	}
	if req.Provider == "" {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "provider is required for list models request",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.ListModelsRequest,
			},
		}
	}
	if ctx == nil {
		ctx = bifrost.ctx
	}
	// Preparing request
	request := &schemas.BifrostListModelsRequest{
		Provider:    req.Provider,
		PageSize:    req.PageSize,
		PageToken:   req.PageToken,
		ExtraParams: req.ExtraParams,
	}
	// Getting provider from the memory
	provider := bifrost.getProviderByKey(req.Provider)
	if provider == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "provider not found for list models request",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.ListModelsRequest,
				Provider:    req.Provider,
			},
		}
	}

	// Determine the base provider type for key requirement checks
	baseProvider := req.Provider
	config, err := bifrost.account.GetConfigForProvider(req.Provider)
	if err != nil {
		bifrostErr := newBifrostErrorFromMsg(fmt.Sprintf("failed to get config for provider %s: %v", req.Provider, err.Error()))
		bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
			RequestType: schemas.ListModelsRequest,
			Provider:    req.Provider,
		}
		return nil, bifrostErr
	}
	if config == nil {
		bifrostErr := newBifrostErrorFromMsg(fmt.Sprintf("config is nil for provider %s", req.Provider))
		bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
			RequestType: schemas.ListModelsRequest,
			Provider:    req.Provider,
		}
		return nil, bifrostErr
	}
	if config.CustomProviderConfig != nil && config.CustomProviderConfig.BaseProviderType != "" {
		baseProvider = config.CustomProviderConfig.BaseProviderType
	}

	var keys []schemas.Key
	if providerRequiresKey(baseProvider, config.CustomProviderConfig) {
		keys, err = bifrost.getAllSupportedKeys(&ctx, req.Provider, baseProvider)
		if err != nil {
			bifrostErr := newBifrostError(err)
			bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
				RequestType: schemas.ListModelsRequest,
				Provider:    req.Provider,
			}
			return nil, bifrostErr
		}
	}

	response, bifrostErr := executeRequestWithRetries(&ctx, config, func() (*schemas.BifrostListModelsResponse, *schemas.BifrostError) {
		return provider.ListModels(ctx, keys, request)
	}, schemas.ListModelsRequest, req.Provider, "")
	if bifrostErr != nil {
		bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
			RequestType: schemas.ListModelsRequest,
			Provider:    req.Provider,
		}
		return nil, bifrostErr
	}
	return response, nil
}

// ListAllModels lists all models from all configured providers.
// It accumulates responses from all providers with a limit of 1000 per provider to get all results.
func (bifrost *Bifrost) ListAllModels(ctx context.Context, request *schemas.BifrostListModelsRequest) (*schemas.BifrostListModelsResponse, *schemas.BifrostError) {
	if request == nil {
		request = &schemas.BifrostListModelsRequest{}
	}

	providerKeys, err := bifrost.GetConfiguredProviders()
	if err != nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: err.Error(),
				Error:   err,
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.ListModelsRequest,
			},
		}
	}

	startTime := time.Now()

	// Result structure for collecting provider responses
	type providerResult struct {
		models []schemas.Model
		err    *schemas.BifrostError
	}

	results := make(chan providerResult, len(providerKeys))
	var wg sync.WaitGroup

	// Launch concurrent requests for all providers
	for _, providerKey := range providerKeys {
		if strings.TrimSpace(string(providerKey)) == "" {
			continue
		}

		wg.Add(1)
		go func(providerKey schemas.ModelProvider) {
			defer wg.Done()

			providerModels := make([]schemas.Model, 0)
			var providerErr *schemas.BifrostError

			// Create request for this provider with limit of 1000
			providerRequest := &schemas.BifrostListModelsRequest{
				Provider: providerKey,
				PageSize: schemas.DefaultPageSize,
			}

			iterations := 0
			for {
				// check for context cancellation
				select {
				case <-ctx.Done():
					bifrost.logger.Warn(fmt.Sprintf("context cancelled for provider %s", providerKey))
					return
				default:
				}

				iterations++
				if iterations > schemas.MaxPaginationRequests {
					bifrost.logger.Warn(fmt.Sprintf("reached maximum pagination requests (%d) for provider %s, please increase the page size", schemas.MaxPaginationRequests, providerKey))
					break
				}

				response, bifrostErr := bifrost.ListModelsRequest(ctx, providerRequest)
				if bifrostErr != nil {
					// Skip logging "no keys found" and "not supported" errors as they are expected when a provider is not configured
					if !strings.Contains(bifrostErr.Error.Message, "no keys found") &&
						!strings.Contains(bifrostErr.Error.Message, "not supported") {
						providerErr = bifrostErr
						bifrost.logger.Warn(fmt.Sprintf("failed to list models for provider %s: %s", providerKey, GetErrorMessage(bifrostErr)))
					}
					break
				}

				if response == nil || len(response.Data) == 0 {
					break
				}

				providerModels = append(providerModels, response.Data...)

				// Check if there are more pages
				if response.NextPageToken == "" {
					break
				}

				// Set the page token for the next request
				providerRequest.PageToken = response.NextPageToken
			}

			results <- providerResult{models: providerModels, err: providerErr}
		}(providerKey)
	}

	// Wait for all goroutines to complete
	wg.Wait()
	close(results)

	// Accumulate all models from all providers
	allModels := make([]schemas.Model, 0)
	var firstError *schemas.BifrostError

	for result := range results {
		if len(result.models) > 0 {
			allModels = append(allModels, result.models...)
		}
		if result.err != nil && firstError == nil {
			firstError = result.err
		}
	}

	// If we couldn't get any models from any provider, return the first error
	if len(allModels) == 0 && firstError != nil {
		return nil, firstError
	}

	// Sort models alphabetically by ID
	sort.Slice(allModels, func(i, j int) bool {
		return allModels[i].ID < allModels[j].ID
	})

	// Return aggregated response with accumulated latency
	response := &schemas.BifrostListModelsResponse{
		Data: allModels,
		ExtraFields: schemas.BifrostResponseExtraFields{
			RequestType: schemas.ListModelsRequest,
			Latency:     time.Since(startTime).Milliseconds(),
		},
	}

	response = response.ApplyPagination(request.PageSize, request.PageToken)

	return response, nil
}

// TextCompletionRequest sends a text completion request to the specified provider.
func (bifrost *Bifrost) TextCompletionRequest(ctx context.Context, req *schemas.BifrostTextCompletionRequest) (*schemas.BifrostTextCompletionResponse, *schemas.BifrostError) {
	if req == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "text completion request is nil",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.TextCompletionRequest,
			},
		}
	}
	if req.Input == nil || (req.Input.PromptStr == nil && req.Input.PromptArray == nil) {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "prompt not provided for text completion request",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType:    schemas.TextCompletionRequest,
				Provider:       req.Provider,
				ModelRequested: req.Model,
			},
		}
	}
	// Preparing request
	bifrostReq := bifrost.getBifrostRequest()
	bifrostReq.RequestType = schemas.TextCompletionRequest
	bifrostReq.TextCompletionRequest = req

	response, err := bifrost.handleRequest(ctx, bifrostReq)
	if err != nil {
		return nil, err
	}
	//TODO: Release the response
	return response.TextCompletionResponse, nil
}

// TextCompletionStreamRequest sends a streaming text completion request to the specified provider.
func (bifrost *Bifrost) TextCompletionStreamRequest(ctx context.Context, req *schemas.BifrostTextCompletionRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) {
	if req == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "text completion stream request is nil",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.TextCompletionStreamRequest,
			},
		}
	}
	if req.Input == nil || (req.Input.PromptStr == nil && req.Input.PromptArray == nil) {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "text not provided for text completion stream request",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType:    schemas.TextCompletionStreamRequest,
				Provider:       req.Provider,
				ModelRequested: req.Model,
			},
		}
	}
	bifrostReq := bifrost.getBifrostRequest()
	bifrostReq.RequestType = schemas.TextCompletionStreamRequest
	bifrostReq.TextCompletionRequest = req
	return bifrost.handleStreamRequest(ctx, bifrostReq)
}

// ChatCompletionRequest sends a chat completion request to the specified provider.
func (bifrost *Bifrost) ChatCompletionRequest(ctx context.Context, req *schemas.BifrostChatRequest) (*schemas.BifrostChatResponse, *schemas.BifrostError) {
	if req == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "chat completion request is nil",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.ChatCompletionRequest,
			},
		}
	}
	if req.Input == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "chats not provided for chat completion request",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType:    schemas.ChatCompletionRequest,
				Provider:       req.Provider,
				ModelRequested: req.Model,
			},
		}
	}

	bifrostReq := bifrost.getBifrostRequest()
	bifrostReq.RequestType = schemas.ChatCompletionRequest
	bifrostReq.ChatRequest = req

	response, err := bifrost.handleRequest(ctx, bifrostReq)
	if err != nil {
		return nil, err
	}
	//TODO: Release the response
	return response.ChatResponse, nil
}

// ChatCompletionStreamRequest sends a chat completion stream request to the specified provider.
func (bifrost *Bifrost) ChatCompletionStreamRequest(ctx context.Context, req *schemas.BifrostChatRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) {
	if req == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "chat completion stream request is nil",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.ChatCompletionStreamRequest,
			},
		}
	}
	if req.Input == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "chats not provided for chat completion request",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType:    schemas.ChatCompletionStreamRequest,
				Provider:       req.Provider,
				ModelRequested: req.Model,
			},
		}
	}

	bifrostReq := bifrost.getBifrostRequest()
	bifrostReq.RequestType = schemas.ChatCompletionStreamRequest
	bifrostReq.ChatRequest = req

	return bifrost.handleStreamRequest(ctx, bifrostReq)
}

// ResponsesRequest sends a responses request to the specified provider.
func (bifrost *Bifrost) ResponsesRequest(ctx context.Context, req *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) {
	if req == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "responses request is nil",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.ResponsesRequest,
			},
		}
	}
	if req.Input == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "responses not provided for responses request",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType:    schemas.ResponsesRequest,
				Provider:       req.Provider,
				ModelRequested: req.Model,
			},
		}
	}

	bifrostReq := bifrost.getBifrostRequest()
	bifrostReq.RequestType = schemas.ResponsesRequest
	bifrostReq.ResponsesRequest = req

	response, err := bifrost.handleRequest(ctx, bifrostReq)
	if err != nil {
		return nil, err
	}
	//TODO: Release the response
	return response.ResponsesResponse, nil
}

// ResponsesStreamRequest sends a responses stream request to the specified provider.
func (bifrost *Bifrost) ResponsesStreamRequest(ctx context.Context, req *schemas.BifrostResponsesRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) {
	if req == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "responses stream request is nil",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.ResponsesStreamRequest,
			},
		}
	}
	if req.Input == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "responses not provided for responses stream request",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType:    schemas.ResponsesStreamRequest,
				Provider:       req.Provider,
				ModelRequested: req.Model,
			},
		}
	}

	bifrostReq := bifrost.getBifrostRequest()
	bifrostReq.RequestType = schemas.ResponsesStreamRequest
	bifrostReq.ResponsesRequest = req

	return bifrost.handleStreamRequest(ctx, bifrostReq)
}

// EmbeddingRequest sends an embedding request to the specified provider.
func (bifrost *Bifrost) EmbeddingRequest(ctx context.Context, req *schemas.BifrostEmbeddingRequest) (*schemas.BifrostEmbeddingResponse, *schemas.BifrostError) {
	if req == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "embedding request is nil",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.EmbeddingRequest,
			},
		}
	}
	if req.Input == nil || (req.Input.Text == nil && req.Input.Texts == nil && req.Input.Embedding == nil && req.Input.Embeddings == nil) {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "embedding input not provided for embedding request",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType:    schemas.EmbeddingRequest,
				Provider:       req.Provider,
				ModelRequested: req.Model,
			},
		}
	}

	bifrostReq := bifrost.getBifrostRequest()
	bifrostReq.RequestType = schemas.EmbeddingRequest
	bifrostReq.EmbeddingRequest = req

	response, err := bifrost.handleRequest(ctx, bifrostReq)
	if err != nil {
		return nil, err
	}
	//TODO: Release the response
	return response.EmbeddingResponse, nil
}

// SpeechRequest sends a speech request to the specified provider.
func (bifrost *Bifrost) SpeechRequest(ctx context.Context, req *schemas.BifrostSpeechRequest) (*schemas.BifrostSpeechResponse, *schemas.BifrostError) {
	if req == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "speech request is nil",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.SpeechRequest,
			},
		}
	}
	if req.Input == nil || req.Input.Input == "" {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "speech input not provided for speech request",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType:    schemas.SpeechRequest,
				Provider:       req.Provider,
				ModelRequested: req.Model,
			},
		}
	}

	bifrostReq := bifrost.getBifrostRequest()
	bifrostReq.RequestType = schemas.SpeechRequest
	bifrostReq.SpeechRequest = req

	response, err := bifrost.handleRequest(ctx, bifrostReq)
	if err != nil {
		return nil, err
	}
	//TODO: Release the response
	return response.SpeechResponse, nil
}

// SpeechStreamRequest sends a speech stream request to the specified provider.
func (bifrost *Bifrost) SpeechStreamRequest(ctx context.Context, req *schemas.BifrostSpeechRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) {
	if req == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "speech stream request is nil",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.SpeechStreamRequest,
			},
		}
	}
	if req.Input == nil || req.Input.Input == "" {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "speech input not provided for speech stream request",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType:    schemas.SpeechStreamRequest,
				Provider:       req.Provider,
				ModelRequested: req.Model,
			},
		}
	}

	bifrostReq := bifrost.getBifrostRequest()
	bifrostReq.RequestType = schemas.SpeechStreamRequest
	bifrostReq.SpeechRequest = req

	return bifrost.handleStreamRequest(ctx, bifrostReq)
}

// TranscriptionRequest sends a transcription request to the specified provider.
func (bifrost *Bifrost) TranscriptionRequest(ctx context.Context, req *schemas.BifrostTranscriptionRequest) (*schemas.BifrostTranscriptionResponse, *schemas.BifrostError) {
	if req == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "transcription request is nil",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.TranscriptionRequest,
			},
		}
	}
	if req.Input == nil || req.Input.File == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "transcription input not provided for transcription request",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType:    schemas.TranscriptionRequest,
				Provider:       req.Provider,
				ModelRequested: req.Model,
			},
		}
	}

	bifrostReq := bifrost.getBifrostRequest()
	bifrostReq.RequestType = schemas.TranscriptionRequest
	bifrostReq.TranscriptionRequest = req

	response, err := bifrost.handleRequest(ctx, bifrostReq)
	if err != nil {
		return nil, err
	}
	//TODO: Release the response
	return response.TranscriptionResponse, nil
}

// TranscriptionStreamRequest sends a transcription stream request to the specified provider.
func (bifrost *Bifrost) TranscriptionStreamRequest(ctx context.Context, req *schemas.BifrostTranscriptionRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) {
	if req == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "transcription stream request is nil",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.TranscriptionStreamRequest,
			},
		}
	}
	if req.Input == nil || req.Input.File == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "transcription input not provided for transcription stream request",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType:    schemas.TranscriptionStreamRequest,
				Provider:       req.Provider,
				ModelRequested: req.Model,
			},
		}
	}

	bifrostReq := bifrost.getBifrostRequest()
	bifrostReq.RequestType = schemas.TranscriptionStreamRequest
	bifrostReq.TranscriptionRequest = req

	return bifrost.handleStreamRequest(ctx, bifrostReq)
}

// ImageGenerationRequest sends a image generation request to the specified provider.
func (bifrost *Bifrost) ImageGenerationRequest(ctx context.Context,
	req *schemas.BifrostImageGenerationRequest) (*schemas.BifrostImageGenerationResponse, *schemas.BifrostError) {
	if req == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "image generation request is nil",
			},
		}
	}
	if req.Input == nil || req.Input.Prompt == "" {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "prompt not provided for image generation request",
			},
		}
	}

	bifrostReq := bifrost.getBifrostRequest()
	bifrostReq.RequestType = schemas.ImageGenerationRequest
	bifrostReq.ImageGenerationRequest = req

	response, err := bifrost.handleRequest(ctx, bifrostReq)
	if err != nil {
		return nil, err
	}
	if response == nil || response.ImageGenerationResponse == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "received nil response from provider",
			},
		}
	}

	return response.ImageGenerationResponse, nil
}

// ImageGenerationStreamRequest sends a image generation stream request to the specified provider.
func (bifrost *Bifrost) ImageGenerationStreamRequest(ctx context.Context,
	req *schemas.BifrostImageGenerationRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) {
	if req == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "image generation stream request is nil",
			},
		}
	}
	if req.Input == nil || req.Input.Prompt == "" {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "prompt not provided for image generation stream request",
			},
		}
	}

	bifrostReq := bifrost.getBifrostRequest()
	bifrostReq.RequestType = schemas.ImageGenerationStreamRequest
	bifrostReq.ImageGenerationRequest = req

	return bifrost.handleStreamRequest(ctx, bifrostReq)
}

// RemovePlugin removes a plugin from the server.
func (bifrost *Bifrost) RemovePlugin(name string) error {

	for {
		oldPlugins := bifrost.plugins.Load()
		if oldPlugins == nil {
			return nil
		}
		var pluginToCleanup schemas.Plugin
		found := false
		// Create new slice with replaced plugin
		newPlugins := make([]schemas.Plugin, len(*oldPlugins))
		copy(newPlugins, *oldPlugins)
		for i, p := range newPlugins {
			if p.GetName() == name {
				pluginToCleanup = p
				bifrost.logger.Debug("removing plugin %s", name)
				newPlugins = append(newPlugins[:i], newPlugins[i+1:]...)
				found = true
				break
			}
		}
		if !found {
			return nil
		}
		if pluginToCleanup != nil {
			// Atomic compare-and-swap
			if bifrost.plugins.CompareAndSwap(oldPlugins, &newPlugins) {
				// Cleanup the old plugin
				err := pluginToCleanup.Cleanup()
				if err != nil {
					bifrost.logger.Warn("failed to cleanup old plugin %s: %v", pluginToCleanup.GetName(), err)
				}
				return nil
			}
		}
		// Retrying as swapping did not work
	}
}

// ReloadPlugin reloads a plugin with new instance
// During the reload - it's stop the world phase where we take a global lock on the plugin mutex
func (bifrost *Bifrost) ReloadPlugin(plugin schemas.Plugin) error {
	for {
		var pluginToCleanup schemas.Plugin
		found := false
		oldPlugins := bifrost.plugins.Load()
		if oldPlugins == nil {
			return nil
		}
		// Create new slice with replaced plugin
		newPlugins := make([]schemas.Plugin, len(*oldPlugins))
		copy(newPlugins, *oldPlugins)
		for i, p := range newPlugins {
			if p.GetName() == plugin.GetName() {
				// Cleaning up old plugin before replacing it
				pluginToCleanup = p
				bifrost.logger.Debug("replacing plugin %s with new instance", plugin.GetName())
				newPlugins[i] = plugin
				found = true
				break
			}
		}
		if !found {
			// This means that user is adding a new plugin
			bifrost.logger.Debug("adding new plugin %s", plugin.GetName())
			newPlugins = append(newPlugins, plugin)
		}
		// Atomic compare-and-swap
		if bifrost.plugins.CompareAndSwap(oldPlugins, &newPlugins) {
			// Cleanup the old plugin
			if found && pluginToCleanup != nil {
				err := pluginToCleanup.Cleanup()
				if err != nil {
					bifrost.logger.Warn("failed to cleanup old plugin %s: %v", pluginToCleanup.GetName(), err)
				}
			}
			return nil
		}
		// Retrying as swapping did not work
	}
}

// GetConfiguredProviders returns the configured providers.
//
// Returns:
//   - []schemas.ModelProvider: List of configured providers
//   - error: Any error that occurred during the retrieval process
//
// Example:
//
//	providers, err := bifrost.GetConfiguredProviders()
//	if err != nil {
//		return nil, err
//	}
//	fmt.Println(providers)
func (bifrost *Bifrost) GetConfiguredProviders() ([]schemas.ModelProvider, error) {
	providers := bifrost.providers.Load()
	if providers == nil {
		return nil, fmt.Errorf("no providers configured")
	}
	modelProviders := make([]schemas.ModelProvider, len(*providers))
	for i, provider := range *providers {
		modelProviders[i] = provider.GetProviderKey()
	}
	return modelProviders, nil
}

// UpdateProvider dynamically updates a provider with new configuration.
// This method gracefully recreates the provider instance with updated settings,
// stops existing workers, creates a new queue with updated settings,
// and starts new workers with the updated provider and concurrency configuration.
//
// Parameters:
//   - providerKey: The provider to update
//
// Returns:
//   - error: Any error that occurred during the update process
//
// Note: This operation will temporarily pause request processing for the specified provider
// while the transition occurs. In-flight requests will complete before workers are stopped.
// Buffered requests in the old queue will be transferred to the new queue to prevent loss.
func (bifrost *Bifrost) UpdateProvider(providerKey schemas.ModelProvider) error {
	bifrost.logger.Info(fmt.Sprintf("Updating provider configuration for provider %s", providerKey))

	// Get the updated configuration from the account
	providerConfig, err := bifrost.account.GetConfigForProvider(providerKey)
	if err != nil {
		return fmt.Errorf("failed to get updated config for provider %s: %v", providerKey, err)
	}
	if providerConfig == nil {
		return fmt.Errorf("config is nil for provider %s", providerKey)
	}

	// Lock the provider to prevent concurrent access during update
	providerMutex := bifrost.getProviderMutex(providerKey)
	providerMutex.Lock()
	defer providerMutex.Unlock()

	// Check if provider currently exists
	oldQueueValue, exists := bifrost.requestQueues.Load(providerKey)
	if !exists {
		bifrost.logger.Debug("provider %s not currently active, initializing with new configuration", providerKey)
		// If provider doesn't exist, just prepare it with new configuration
		return bifrost.prepareProvider(providerKey, providerConfig)
	}

	oldQueue := oldQueueValue.(chan *ChannelMessage)

	bifrost.logger.Debug("gracefully stopping existing workers for provider %s", providerKey)

	// Step 1: Create new queue with updated buffer size
	newQueue := make(chan *ChannelMessage, providerConfig.ConcurrencyAndBufferSize.BufferSize)

	// Step 2: Transfer any buffered requests from old queue to new queue
	// This prevents request loss during the transition
	transferredCount := 0
	var transferWaitGroup sync.WaitGroup
	for {
		select {
		case msg := <-oldQueue:
			select {
			case newQueue <- msg:
				transferredCount++
			default:
				// New queue is full, handle this request in a goroutine
				// This is unlikely with proper buffer sizing but provides safety
				transferWaitGroup.Add(1)
				go func(m *ChannelMessage) {
					defer transferWaitGroup.Done()
					select {
					case newQueue <- m:
						// Message successfully transferred
					case <-time.After(5 * time.Second):
						bifrost.logger.Warn("Failed to transfer buffered request to new queue within timeout")
						// Send error response to avoid hanging the client
						provider, model, _ := m.BifrostRequest.GetRequestFields()
						select {
						case m.Err <- schemas.BifrostError{
							IsBifrostError: false,
							Error: &schemas.ErrorField{
								Message: "request failed during provider concurrency update",
							},
							ExtraFields: schemas.BifrostErrorExtraFields{
								RequestType:    m.RequestType,
								Provider:       provider,
								ModelRequested: model,
							},
						}:
						case <-time.After(1 * time.Second):
							// If we can't send the error either, just log and continue
							bifrost.logger.Warn("Failed to send error response during transfer timeout")
						}
					}
				}(msg)
				goto transferComplete
			}
		default:
			// No more buffered messages
			goto transferComplete
		}
	}

transferComplete:
	// Wait for all transfer goroutines to complete
	transferWaitGroup.Wait()
	if transferredCount > 0 {
		bifrost.logger.Info("transferred %d buffered requests to new queue for provider %s", transferredCount, providerKey)
	}

	// Step 3: Close the old queue to signal workers to stop
	close(oldQueue)

	// Step 4: Atomically replace the queue
	bifrost.requestQueues.Store(providerKey, newQueue)

	// Step 5: Wait for all existing workers to finish processing in-flight requests
	waitGroup, exists := bifrost.waitGroups.Load(providerKey)
	if exists {
		waitGroup.(*sync.WaitGroup).Wait()
		bifrost.logger.Debug("all workers for provider %s have stopped", providerKey)
	}

	// Step 6: Create new wait group for the updated workers
	bifrost.waitGroups.Store(providerKey, &sync.WaitGroup{})

	// Step 7: Create provider instance
	provider, err := bifrost.createBaseProvider(providerKey, providerConfig)
	if err != nil {
		return fmt.Errorf("failed to create provider instance for %s: %v", providerKey, err)
	}

	// Step 7.5: Atomically replace the provider in the providers slice
	// This must happen before starting new workers to prevent stale reads
	bifrost.logger.Debug("atomically replacing provider instance in providers slice for %s", providerKey)

	replacementAttempts := 0
	maxReplacementAttempts := 100 // Prevent infinite loops in high-contention scenarios

	for {
		replacementAttempts++
		if replacementAttempts > maxReplacementAttempts {
			return fmt.Errorf("failed to replace provider %s in providers slice after %d attempts", providerKey, maxReplacementAttempts)
		}

		oldPtr := bifrost.providers.Load()
		var oldSlice []schemas.Provider
		if oldPtr != nil {
			oldSlice = *oldPtr
		}

		// Create new slice without the old provider of this key
		// Use exact capacity to avoid allocations
		newSlice := make([]schemas.Provider, 0, len(oldSlice))
		oldProviderFound := false

		for _, existingProvider := range oldSlice {
			if existingProvider.GetProviderKey() != providerKey {
				newSlice = append(newSlice, existingProvider)
			} else {
				oldProviderFound = true
			}
		}

		// Add the new provider
		newSlice = append(newSlice, provider)

		if bifrost.providers.CompareAndSwap(oldPtr, &newSlice) {
			if oldProviderFound {
				bifrost.logger.Debug("successfully replaced existing provider instance for %s in providers slice", providerKey)
			} else {
				bifrost.logger.Debug("successfully added new provider instance for %s to providers slice", providerKey)
			}
			break
		}
		// Retrying as swapping did not work (likely due to concurrent modification)
	}

	// Step 8: Start new workers with updated concurrency
	bifrost.logger.Debug("starting %d new workers for provider %s with buffer size %d",
		providerConfig.ConcurrencyAndBufferSize.Concurrency,
		providerKey,
		providerConfig.ConcurrencyAndBufferSize.BufferSize)

	waitGroupValue, _ := bifrost.waitGroups.Load(providerKey)
	currentWaitGroup := waitGroupValue.(*sync.WaitGroup)

	for range providerConfig.ConcurrencyAndBufferSize.Concurrency {
		currentWaitGroup.Add(1)
		go bifrost.requestWorker(provider, providerConfig, newQueue)
	}

	bifrost.logger.Info("successfully updated provider configuration for provider %s", providerKey)
	return nil
}

// GetDropExcessRequests returns the current value of DropExcessRequests
func (bifrost *Bifrost) GetDropExcessRequests() bool {
	return bifrost.dropExcessRequests.Load()
}

// UpdateDropExcessRequests updates the DropExcessRequests setting at runtime.
// This allows for hot-reloading of this configuration value.
func (bifrost *Bifrost) UpdateDropExcessRequests(value bool) {
	bifrost.dropExcessRequests.Store(value)
	bifrost.logger.Info("drop_excess_requests updated to: %v", value)
}

// getProviderMutex gets or creates a mutex for the given provider
func (bifrost *Bifrost) getProviderMutex(providerKey schemas.ModelProvider) *sync.RWMutex {
	mutexValue, _ := bifrost.providerMutexes.LoadOrStore(providerKey, &sync.RWMutex{})
	return mutexValue.(*sync.RWMutex)
}

// MCP PUBLIC API

// RegisterMCPTool registers a typed tool handler with the MCP integration.
// This allows developers to easily add custom tools that will be available
// to all LLM requests processed by this Bifrost instance.
//
// Parameters:
//   - name: Unique tool name
//   - description: Human-readable tool description
//   - handler: Function that handles tool execution
//   - toolSchema: Bifrost tool schema for function calling
//
// Returns:
//   - error: Any registration error
//
// Example:
//
//	type EchoArgs struct {
//	    Message string `json:"message"`
//	}
//
//	err := bifrost.RegisterMCPTool("echo", "Echo a message",
//	    func(args EchoArgs) (string, error) {
//	        return args.Message, nil
//	    }, toolSchema)
func (bifrost *Bifrost) RegisterMCPTool(name, description string, handler func(args any) (string, error), toolSchema schemas.ChatTool) error {
	if bifrost.mcpManager == nil {
		return fmt.Errorf("MCP is not configured in this Bifrost instance")
	}

	return bifrost.mcpManager.registerTool(name, description, handler, toolSchema)
}

// ExecuteMCPTool executes an MCP tool call and returns the result as a tool message.
// This is the main public API for manual MCP tool execution.
//
// Parameters:
//   - ctx: Execution context
//   - toolCall: The tool call to execute (from assistant message)
//
// Returns:
//   - schemas.ChatMessage: Tool message with execution result
//   - schemas.BifrostError: Any execution error
func (bifrost *Bifrost) ExecuteMCPTool(ctx context.Context, toolCall schemas.ChatAssistantMessageToolCall) (*schemas.ChatMessage, *schemas.BifrostError) {
	if bifrost.mcpManager == nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: "MCP is not configured in this Bifrost instance",
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.ChatCompletionRequest, // MCP tools are used with chat completions
			},
		}
	}

	result, err := bifrost.mcpManager.executeTool(ctx, toolCall)
	if err != nil {
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: err.Error(),
				Error:   err,
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType: schemas.ChatCompletionRequest, // MCP tools are used with chat completions
			},
		}
	}

	return result, nil
}

// IMPORTANT: Running the MCP client management operations (GetMCPClients, AddMCPClient, RemoveMCPClient, EditMCPClientTools)
// may temporarily increase latency for incoming requests while the operations are being processed.
// These operations involve network I/O and connection management that require mutex locks
// which can block briefly during execution.

// GetMCPClients returns all MCP clients managed by the Bifrost instance.
//
// Returns:
//   - []schemas.MCPClient: List of all MCP clients
//   - error: Any retrieval error
func (bifrost *Bifrost) GetMCPClients() ([]schemas.MCPClient, error) {
	if bifrost.mcpManager == nil {
		return nil, fmt.Errorf("MCP is not configured in this Bifrost instance")
	}

	clients, err := bifrost.mcpManager.GetClients()
	if err != nil {
		return nil, err
	}

	clientsInConfig := make([]schemas.MCPClient, 0, len(clients))
	for _, client := range clients {
		tools := make([]schemas.ChatToolFunction, 0, len(client.ToolMap))
		for _, tool := range client.ToolMap {
			if tool.Function != nil {
				tools = append(tools, *tool.Function)
			}
		}

		sort.Slice(tools, func(i, j int) bool {
			return tools[i].Name < tools[j].Name
		})

		state := schemas.MCPConnectionStateConnected
		if client.Conn == nil {
			state = schemas.MCPConnectionStateDisconnected
		}

		clientsInConfig = append(clientsInConfig, schemas.MCPClient{
			Config: client.ExecutionConfig,
			Tools:  tools,
			State:  state,
		})
	}

	return clientsInConfig, nil
}

// AddMCPClient adds a new MCP client to the Bifrost instance.
// This allows for dynamic MCP client management at runtime.
//
// Parameters:
//   - config: MCP client configuration
//
// Returns:
//   - error: Any registration error
//
// Example:
//
//	err := bifrost.AddMCPClient(schemas.MCPClientConfig{
//	    Name: "my-mcp-client",
//	    ConnectionType: schemas.MCPConnectionTypeHTTP,
//	    ConnectionString: &url,
//	})
func (bifrost *Bifrost) AddMCPClient(config schemas.MCPClientConfig) error {
	if bifrost.mcpManager == nil {
		manager := &MCPManager{
			ctx:       bifrost.ctx,
			clientMap: make(map[string]*MCPClient),
			logger:    bifrost.logger,
		}

		bifrost.mcpManager = manager
	}

	return bifrost.mcpManager.AddClient(config)
}

// RemoveMCPClient removes an MCP client from the Bifrost instance.
// This allows for dynamic MCP client management at runtime.
//
// Parameters:
//   - id: ID of the client to remove
//
// Returns:
//   - error: Any removal error
//
// Example:
//
//	err := bifrost.RemoveMCPClient("my-mcp-client-id")
//	if err != nil {
//	    log.Fatalf("Failed to remove MCP client: %v", err)
//	}
func (bifrost *Bifrost) RemoveMCPClient(id string) error {
	if bifrost.mcpManager == nil {
		return fmt.Errorf("MCP is not configured in this Bifrost instance")
	}

	return bifrost.mcpManager.RemoveClient(id)
}

// EditMCPClient edits the tools of an MCP client.
// This allows for dynamic MCP client tool management at runtime.
//
// Parameters:
//   - id: ID of the client to edit
//   - updatedConfig: Updated MCP client configuration
//
// Returns:
//   - error: Any edit error
//
// Example:
//
//	err := bifrost.EditMCPClient("my-mcp-client-id", schemas.MCPClientConfig{
//	    Name:           "my-mcp-client-name",
//	    ToolsToExecute: []string{"tool1", "tool2"},
//	})
func (bifrost *Bifrost) EditMCPClient(id string, updatedConfig schemas.MCPClientConfig) error {
	if bifrost.mcpManager == nil {
		return fmt.Errorf("MCP is not configured in this Bifrost instance")
	}

	return bifrost.mcpManager.EditClient(id, updatedConfig)
}

// ReconnectMCPClient attempts to reconnect an MCP client if it is disconnected.
//
// Parameters:
//   - id: ID of the client to reconnect
//
// Returns:
//   - error: Any reconnection error
func (bifrost *Bifrost) ReconnectMCPClient(id string) error {
	if bifrost.mcpManager == nil {
		return fmt.Errorf("MCP is not configured in this Bifrost instance")
	}

	return bifrost.mcpManager.ReconnectClient(id)
}

// PROVIDER MANAGEMENT

// createBaseProvider creates a provider based on the base provider type
func (bifrost *Bifrost) createBaseProvider(providerKey schemas.ModelProvider, config *schemas.ProviderConfig) (schemas.Provider, error) {
	// Determine which provider type to create
	targetProviderKey := providerKey

	if config.CustomProviderConfig != nil {
		// Validate custom provider config
		if config.CustomProviderConfig.BaseProviderType == "" {
			return nil, fmt.Errorf("custom provider config missing base provider type")
		}

		// Validate that base provider type is supported
		if !IsSupportedBaseProvider(config.CustomProviderConfig.BaseProviderType) {
			return nil, fmt.Errorf("unsupported base provider type: %s", config.CustomProviderConfig.BaseProviderType)
		}

		// Automatically set the custom provider key to the provider name
		config.CustomProviderConfig.CustomProviderKey = string(providerKey)

		targetProviderKey = config.CustomProviderConfig.BaseProviderType
	}

	switch targetProviderKey {
	case schemas.OpenAI:
		return openai.NewOpenAIProvider(config, bifrost.logger), nil
	case schemas.Anthropic:
		return anthropic.NewAnthropicProvider(config, bifrost.logger), nil
	case schemas.Bedrock:
		return bedrock.NewBedrockProvider(config, bifrost.logger)
	case schemas.Cohere:
		return cohere.NewCohereProvider(config, bifrost.logger)
	case schemas.Azure:
		return azure.NewAzureProvider(config, bifrost.logger)
	case schemas.Vertex:
		return vertex.NewVertexProvider(config, bifrost.logger)
	case schemas.Mistral:
		return mistral.NewMistralProvider(config, bifrost.logger), nil
	case schemas.Ollama:
		return ollama.NewOllamaProvider(config, bifrost.logger)
	case schemas.Groq:
		return groq.NewGroqProvider(config, bifrost.logger)
	case schemas.SGL:
		return sgl.NewSGLProvider(config, bifrost.logger)
	case schemas.Parasail:
		return parasail.NewParasailProvider(config, bifrost.logger)
	case schemas.Perplexity:
		return perplexity.NewPerplexityProvider(config, bifrost.logger)
	case schemas.Cerebras:
		return cerebras.NewCerebrasProvider(config, bifrost.logger)
	case schemas.Gemini:
		return gemini.NewGeminiProvider(config, bifrost.logger), nil
	case schemas.OpenRouter:
		return openrouter.NewOpenRouterProvider(config, bifrost.logger), nil
	case schemas.Elevenlabs:
		return elevenlabs.NewElevenlabsProvider(config, bifrost.logger), nil
	default:
		return nil, fmt.Errorf("unsupported provider: %s", targetProviderKey)
	}
}

// prepareProvider sets up a provider with its configuration, keys, and worker channels.
// It initializes the request queue and starts worker goroutines for processing requests.
// Note: This function assumes the caller has already acquired the appropriate mutex for the provider.
func (bifrost *Bifrost) prepareProvider(providerKey schemas.ModelProvider, config *schemas.ProviderConfig) error {
	providerConfig, err := bifrost.account.GetConfigForProvider(providerKey)
	if err != nil {
		return fmt.Errorf("failed to get config for provider: %v", err)
	}
	if providerConfig == nil {
		return fmt.Errorf("config is nil for provider %s", providerKey)
	}

	queue := make(chan *ChannelMessage, providerConfig.ConcurrencyAndBufferSize.BufferSize) // Buffered channel per provider

	bifrost.requestQueues.Store(providerKey, queue)

	// Start specified number of workers
	bifrost.waitGroups.Store(providerKey, &sync.WaitGroup{})

	provider, err := bifrost.createBaseProvider(providerKey, config)
	if err != nil {
		return fmt.Errorf("failed to create provider for the given key: %v", err)
	}

	waitGroupValue, _ := bifrost.waitGroups.Load(providerKey)
	currentWaitGroup := waitGroupValue.(*sync.WaitGroup)

	// Atomically append provider to the providers slice
	for {
		oldPtr := bifrost.providers.Load()
		var oldSlice []schemas.Provider
		if oldPtr != nil {
			oldSlice = *oldPtr
		}
		newSlice := make([]schemas.Provider, len(oldSlice)+1)
		copy(newSlice, oldSlice)
		newSlice[len(oldSlice)] = provider
		if bifrost.providers.CompareAndSwap(oldPtr, &newSlice) {
			break
		}
	}

	for range providerConfig.ConcurrencyAndBufferSize.Concurrency {
		currentWaitGroup.Add(1)
		go bifrost.requestWorker(provider, providerConfig, queue)
	}

	return nil
}

// getProviderQueue returns the request queue for a given provider key.
// If the queue doesn't exist, it creates one at runtime and initializes the provider,
// given the provider config is provided in the account interface implementation.
// This function uses read locks to prevent race conditions during provider updates.
func (bifrost *Bifrost) getProviderQueue(providerKey schemas.ModelProvider) (chan *ChannelMessage, error) {
	// Use read lock to allow concurrent reads but prevent concurrent updates
	providerMutex := bifrost.getProviderMutex(providerKey)
	providerMutex.RLock()

	if queueValue, exists := bifrost.requestQueues.Load(providerKey); exists {
		queue := queueValue.(chan *ChannelMessage)
		providerMutex.RUnlock()
		return queue, nil
	}

	// Provider doesn't exist, need to create it
	// Upgrade to write lock for creation
	providerMutex.RUnlock()
	providerMutex.Lock()
	defer providerMutex.Unlock()

	// Double-check after acquiring write lock (another goroutine might have created it)
	if queueValue, exists := bifrost.requestQueues.Load(providerKey); exists {
		queue := queueValue.(chan *ChannelMessage)
		return queue, nil
	}

	bifrost.logger.Debug(fmt.Sprintf("Creating new request queue for provider %s at runtime", providerKey))

	config, err := bifrost.account.GetConfigForProvider(providerKey)
	if err != nil {
		return nil, fmt.Errorf("failed to get config for provider: %v", err)
	}
	if config == nil {
		return nil, fmt.Errorf("config is nil for provider %s", providerKey)
	}

	if err := bifrost.prepareProvider(providerKey, config); err != nil {
		return nil, err
	}

	queueValue, _ := bifrost.requestQueues.Load(providerKey)
	queue := queueValue.(chan *ChannelMessage)

	return queue, nil
}

// getProviderByKey retrieves a provider instance from the providers array by its provider key.
// Returns the provider if found, or nil if no provider with the given key exists.
func (bifrost *Bifrost) getProviderByKey(providerKey schemas.ModelProvider) schemas.Provider {
	providers := bifrost.providers.Load()
	if providers == nil {
		return nil
	}
	// Checking if provider is in the memory
	for _, provider := range *providers {
		if provider.GetProviderKey() == providerKey {
			return provider
		}
	}
	// Could happen when provider is not initialized yet, check if provider config exists in account and if so, initialize it
	config, err := bifrost.account.GetConfigForProvider(providerKey)
	if err != nil || config == nil {
		return nil
	}
	// Lock the provider mutex to avoid races
	providerMutex := bifrost.getProviderMutex(providerKey)
	providerMutex.Lock()
	defer providerMutex.Unlock()
	// Double-check after acquiring the lock
	providers = bifrost.providers.Load()
	if providers != nil {
		for _, p := range *providers {
			if p.GetProviderKey() == providerKey {
				return p
			}
		}
	}
	// Preparing provider
	if err := bifrost.prepareProvider(providerKey, config); err != nil {
		return nil
	}
	// Return newly prepared provider without recursion
	providers = bifrost.providers.Load()
	if providers != nil {
		for _, p := range *providers {
			if p.GetProviderKey() == providerKey {
				return p
			}
		}
	}
	return nil
}

// CORE INTERNAL LOGIC

// shouldTryFallbacks handles the primary error and returns true if we should proceed with fallbacks, false if we should return immediately
func (bifrost *Bifrost) shouldTryFallbacks(req *schemas.BifrostRequest, primaryErr *schemas.BifrostError) bool {
	// If no primary error, we succeeded
	if primaryErr == nil {
		bifrost.logger.Debug("No primary error, we should not try fallbacks")
		return false
	}

	// Handle request cancellation
	if primaryErr.Error != nil && primaryErr.Error.Type != nil && *primaryErr.Error.Type == schemas.RequestCancelled {
		bifrost.logger.Debug("Request cancelled, we should not try fallbacks")
		return false
	}

	// Check if this is a short-circuit error that doesn't allow fallbacks
	// Note: AllowFallbacks = nil is treated as true (allow fallbacks by default)
	if primaryErr.AllowFallbacks != nil && !*primaryErr.AllowFallbacks {
		bifrost.logger.Debug("AllowFallbacks is false, we should not try fallbacks")
		return false
	}

	// If no fallbacks configured, return primary error
	_, _, fallbacks := req.GetRequestFields()
	if len(fallbacks) == 0 {
		bifrost.logger.Debug("No fallbacks configured, we should not try fallbacks")
		return false
	}

	// Should proceed with fallbacks
	return true
}

// prepareFallbackRequest creates a fallback request and validates the provider config
// Returns the fallback request or nil if this fallback should be skipped
func (bifrost *Bifrost) prepareFallbackRequest(req *schemas.BifrostRequest, fallback schemas.Fallback) *schemas.BifrostRequest {
	// Check if we have config for this fallback provider
	_, err := bifrost.account.GetConfigForProvider(fallback.Provider)
	if err != nil {
		bifrost.logger.Warn(fmt.Sprintf("Config not found for provider %s, skipping fallback: %v", fallback.Provider, err))
		return nil
	}

	// Create a new request with the fallback provider and model
	fallbackReq := *req

	if req.TextCompletionRequest != nil {
		tmp := *req.TextCompletionRequest
		tmp.Provider = fallback.Provider
		tmp.Model = fallback.Model
		fallbackReq.TextCompletionRequest = &tmp
	}

	if req.ChatRequest != nil {
		tmp := *req.ChatRequest
		tmp.Provider = fallback.Provider
		tmp.Model = fallback.Model
		fallbackReq.ChatRequest = &tmp
	}

	if req.ResponsesRequest != nil {
		tmp := *req.ResponsesRequest
		tmp.Provider = fallback.Provider
		tmp.Model = fallback.Model
		fallbackReq.ResponsesRequest = &tmp
	}

	if req.EmbeddingRequest != nil {
		tmp := *req.EmbeddingRequest
		tmp.Provider = fallback.Provider
		tmp.Model = fallback.Model
		fallbackReq.EmbeddingRequest = &tmp
	}

	if req.SpeechRequest != nil {
		tmp := *req.SpeechRequest
		tmp.Provider = fallback.Provider
		tmp.Model = fallback.Model
		fallbackReq.SpeechRequest = &tmp
	}

	if req.TranscriptionRequest != nil {
		tmp := *req.TranscriptionRequest
		tmp.Provider = fallback.Provider
		tmp.Model = fallback.Model
		fallbackReq.TranscriptionRequest = &tmp
	}
	if req.ImageGenerationRequest != nil {
		tmp := *req.ImageGenerationRequest
		tmp.Provider = fallback.Provider
		tmp.Model = fallback.Model
		fallbackReq.ImageGenerationRequest = &tmp
	}

	return &fallbackReq
}

// shouldContinueWithFallbacks processes errors from fallback attempts
// Returns true if we should continue with more fallbacks, false if we should stop
func (bifrost *Bifrost) shouldContinueWithFallbacks(fallback schemas.Fallback, fallbackErr *schemas.BifrostError) bool {
	if fallbackErr.Error.Type != nil && *fallbackErr.Error.Type == schemas.RequestCancelled {
		return false
	}

	// Check if it was a short-circuit error that doesn't allow fallbacks
	if fallbackErr.AllowFallbacks != nil && !*fallbackErr.AllowFallbacks {
		return false
	}

	bifrost.logger.Debug(fmt.Sprintf("Fallback provider %s failed: %s", fallback.Provider, fallbackErr.Error.Message))
	return true
}

// handleRequest handles the request to the provider based on the request type
// It handles plugin hooks, request validation, response processing, and fallback providers.
// If the primary provider fails, it will try each fallback provider in order until one succeeds.
// It is the wrapper for all non-streaming public API methods.
func (bifrost *Bifrost) handleRequest(ctx context.Context, req *schemas.BifrostRequest) (*schemas.BifrostResponse, *schemas.BifrostError) {
	defer bifrost.releaseBifrostRequest(req)

	provider, model, fallbacks := req.GetRequestFields()

	if err := validateRequest(req); err != nil {
		err.ExtraFields = schemas.BifrostErrorExtraFields{
			RequestType:    req.RequestType,
			Provider:       provider,
			ModelRequested: model,
		}
		return nil, err
	}

	// Handle nil context early to prevent blocking
	if ctx == nil {
		ctx = bifrost.ctx
	}

	bifrost.logger.Debug(fmt.Sprintf("Primary provider %s with model %s and %d fallbacks", provider, model, len(fallbacks)))

	// Try the primary provider first
	ctx = context.WithValue(ctx, schemas.BifrostContextKeyFallbackIndex, 0)
	primaryResult, primaryErr := bifrost.tryRequest(ctx, req)
	if primaryErr != nil {
		if primaryErr.Error != nil {
			bifrost.logger.Debug(fmt.Sprintf("Primary provider %s with model %s returned error: %s", provider, model, primaryErr.Error.Message))
		} else {
			bifrost.logger.Debug(fmt.Sprintf("Primary provider %s with model %s returned error: %v", provider, model, primaryErr))
		}
		if len(fallbacks) > 0 {
			bifrost.logger.Debug(fmt.Sprintf("Check if we should try %d fallbacks", len(fallbacks)))
		}
	}

	// Check if we should proceed with fallbacks
	shouldTryFallbacks := bifrost.shouldTryFallbacks(req, primaryErr)
	if !shouldTryFallbacks {
		if primaryErr != nil {
			primaryErr.ExtraFields = schemas.BifrostErrorExtraFields{
				RequestType:    req.RequestType,
				Provider:       provider,
				ModelRequested: model,
			}
		}
		return primaryResult, primaryErr
	}

	// Try fallbacks in order
	for i, fallback := range fallbacks {
		ctx = context.WithValue(ctx, schemas.BifrostContextKeyFallbackIndex, i+1)
		bifrost.logger.Debug(fmt.Sprintf("Trying fallback provider %s with model %s", fallback.Provider, fallback.Model))
		ctx = context.WithValue(ctx, schemas.BifrostContextKeyFallbackRequestID, uuid.New().String())

		fallbackReq := bifrost.prepareFallbackRequest(req, fallback)
		if fallbackReq == nil {
			bifrost.logger.Debug(fmt.Sprintf("Fallback provider %s with model %s is nil", fallback.Provider, fallback.Model))
			continue
		}

		// Try the fallback provider
		result, fallbackErr := bifrost.tryRequest(ctx, fallbackReq)
		if fallbackErr == nil {
			bifrost.logger.Debug(fmt.Sprintf("Successfully used fallback provider %s with model %s", fallback.Provider, fallback.Model))
			return result, nil
		}

		// Check if we should continue with more fallbacks
		if !bifrost.shouldContinueWithFallbacks(fallback, fallbackErr) {
			fallbackErr.ExtraFields = schemas.BifrostErrorExtraFields{
				RequestType:    req.RequestType,
				Provider:       fallback.Provider,
				ModelRequested: fallback.Model,
			}
			return nil, fallbackErr
		}
	}

	if primaryErr != nil {
		primaryErr.ExtraFields = schemas.BifrostErrorExtraFields{
			RequestType:    req.RequestType,
			Provider:       provider,
			ModelRequested: model,
		}
	}

	// All providers failed, return the original error
	return nil, primaryErr
}

// handleStreamRequest handles the stream request to the provider based on the request type
// It handles plugin hooks, request validation, response processing, and fallback providers.
// If the primary provider fails, it will try each fallback provider in order until one succeeds.
// It is the wrapper for all streaming public API methods.
func (bifrost *Bifrost) handleStreamRequest(ctx context.Context, req *schemas.BifrostRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) {
	defer bifrost.releaseBifrostRequest(req)

	provider, model, fallbacks := req.GetRequestFields()

	if err := validateRequest(req); err != nil {
		err.ExtraFields = schemas.BifrostErrorExtraFields{
			RequestType:    req.RequestType,
			Provider:       provider,
			ModelRequested: model,
		}
		return nil, err
	}

	// Handle nil context early to prevent blocking
	if ctx == nil {
		ctx = bifrost.ctx
	}

	// Try the primary provider first
	ctx = context.WithValue(ctx, schemas.BifrostContextKeyFallbackIndex, 0)
	primaryResult, primaryErr := bifrost.tryStreamRequest(ctx, req)

	// Check if we should proceed with fallbacks
	shouldTryFallbacks := bifrost.shouldTryFallbacks(req, primaryErr)
	if !shouldTryFallbacks {
		if primaryErr != nil {
			primaryErr.ExtraFields = schemas.BifrostErrorExtraFields{
				RequestType:    req.RequestType,
				Provider:       provider,
				ModelRequested: model,
			}
		}
		return primaryResult, primaryErr
	}

	// Try fallbacks in order
	for i, fallback := range fallbacks {
		ctx = context.WithValue(ctx, schemas.BifrostContextKeyFallbackIndex, i+1)
		ctx = context.WithValue(ctx, schemas.BifrostContextKeyFallbackRequestID, uuid.New().String())

		fallbackReq := bifrost.prepareFallbackRequest(req, fallback)
		if fallbackReq == nil {
			continue
		}

		// Try the fallback provider
		result, fallbackErr := bifrost.tryStreamRequest(ctx, fallbackReq)
		if fallbackErr == nil {
			bifrost.logger.Debug(fmt.Sprintf("Successfully used fallback provider %s with model %s", fallback.Provider, fallback.Model))
			return result, nil
		}

		// Check if we should continue with more fallbacks
		if !bifrost.shouldContinueWithFallbacks(fallback, fallbackErr) {
			fallbackErr.ExtraFields = schemas.BifrostErrorExtraFields{
				RequestType:    req.RequestType,
				Provider:       fallback.Provider,
				ModelRequested: fallback.Model,
			}
			return nil, fallbackErr
		}
	}

	if primaryErr != nil {
		primaryErr.ExtraFields = schemas.BifrostErrorExtraFields{
			RequestType:    req.RequestType,
			Provider:       provider,
			ModelRequested: model,
		}
	}

	// All providers failed, return the original error
	return nil, primaryErr
}

// tryRequest is a generic function that handles common request processing logic
// It consolidates queue setup, plugin pipeline execution, enqueue logic, and response handling
func (bifrost *Bifrost) tryRequest(ctx context.Context, req *schemas.BifrostRequest) (*schemas.BifrostResponse, *schemas.BifrostError) {
	provider, model, _ := req.GetRequestFields()
	queue, err := bifrost.getProviderQueue(provider)
	if err != nil {
		bifrostErr := newBifrostError(err)
		bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
			RequestType:    req.RequestType,
			Provider:       provider,
			ModelRequested: model,
		}
		return nil, bifrostErr
	}

	// Add MCP tools to request if MCP is configured and requested
	if req.RequestType != schemas.EmbeddingRequest &&
		req.RequestType != schemas.SpeechRequest &&
		req.RequestType != schemas.TranscriptionRequest &&
		bifrost.mcpManager != nil {
		req = bifrost.mcpManager.addMCPToolsToBifrostRequest(ctx, req)
	}

	pipeline := bifrost.getPluginPipeline()
	defer bifrost.releasePluginPipeline(pipeline)

	preReq, shortCircuit, preCount := pipeline.RunPreHooks(&ctx, req)
	if shortCircuit != nil {
		// Handle short-circuit with response (success case)
		if shortCircuit.Response != nil {
			resp, bifrostErr := pipeline.RunPostHooks(&ctx, shortCircuit.Response, nil, preCount)
			if bifrostErr != nil {
				return nil, bifrostErr
			}
			return resp, nil
		}
		// Handle short-circuit with error
		if shortCircuit.Error != nil {
			resp, bifrostErr := pipeline.RunPostHooks(&ctx, nil, shortCircuit.Error, preCount)
			if bifrostErr != nil {
				return nil, bifrostErr
			}
			return resp, nil
		}
	}
	if preReq == nil {
		bifrostErr := newBifrostErrorFromMsg("bifrost request after plugin hooks cannot be nil")
		bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
			RequestType:    req.RequestType,
			Provider:       provider,
			ModelRequested: model,
		}
		return nil, bifrostErr
	}

	msg := bifrost.getChannelMessage(*preReq)
	msg.Context = ctx
	select {
	case queue <- msg:
		// Message was sent successfully
	case <-ctx.Done():
		bifrost.releaseChannelMessage(msg)
		bifrostErr := newBifrostErrorFromMsg("request cancelled while waiting for queue space")
		bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
			RequestType:    req.RequestType,
			Provider:       provider,
			ModelRequested: model,
		}
		return nil, bifrostErr
	default:
		if bifrost.dropExcessRequests.Load() {
			bifrost.releaseChannelMessage(msg)
			bifrost.logger.Warn("Request dropped: queue is full, please increase the queue size or set dropExcessRequests to false")
			bifrostErr := newBifrostErrorFromMsg("request dropped: queue is full")
			bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
				RequestType:    req.RequestType,
				Provider:       provider,
				ModelRequested: model,
			}
			return nil, bifrostErr
		}
		select {
		case queue <- msg:
			// Message was sent successfully
		case <-ctx.Done():
			bifrost.releaseChannelMessage(msg)
			bifrostErr := newBifrostErrorFromMsg("request cancelled while waiting for queue space")
			bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
				RequestType:    req.RequestType,
				Provider:       provider,
				ModelRequested: model,
			}
			return nil, bifrostErr
		}
	}

	var result *schemas.BifrostResponse
	var resp *schemas.BifrostResponse
	pluginCount := len(*bifrost.plugins.Load())
	select {
	case result = <-msg.Response:
		resp, bifrostErr := pipeline.RunPostHooks(&msg.Context, result, nil, pluginCount)
		if bifrostErr != nil {
			bifrost.releaseChannelMessage(msg)
			return nil, bifrostErr
		}
		bifrost.releaseChannelMessage(msg)
		return resp, nil
	case bifrostErrVal := <-msg.Err:
		bifrostErrPtr := &bifrostErrVal
		resp, bifrostErrPtr = pipeline.RunPostHooks(&msg.Context, nil, bifrostErrPtr, pluginCount)
		bifrost.releaseChannelMessage(msg)
		if bifrostErrPtr != nil {
			return nil, bifrostErrPtr
		}
		return resp, nil
	}
}

// tryStreamRequest is a generic function that handles common request processing logic
// It consolidates queue setup, plugin pipeline execution, enqueue logic, and response handling
func (bifrost *Bifrost) tryStreamRequest(ctx context.Context, req *schemas.BifrostRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) {
	provider, model, _ := req.GetRequestFields()
	queue, err := bifrost.getProviderQueue(provider)
	if err != nil {
		bifrostErr := newBifrostError(err)
		bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
			RequestType:    req.RequestType,
			Provider:       provider,
			ModelRequested: model,
		}
		return nil, bifrostErr
	}

	// Add MCP tools to request if MCP is configured and requested
	if req.RequestType != schemas.SpeechStreamRequest && req.RequestType != schemas.TranscriptionStreamRequest && bifrost.mcpManager != nil {
		req = bifrost.mcpManager.addMCPToolsToBifrostRequest(ctx, req)
	}

	pipeline := bifrost.getPluginPipeline()
	defer bifrost.releasePluginPipeline(pipeline)

	preReq, shortCircuit, preCount := pipeline.RunPreHooks(&ctx, req)
	if shortCircuit != nil {
		// Handle short-circuit with response (success case)
		if shortCircuit.Response != nil {
			resp, bifrostErr := pipeline.RunPostHooks(&ctx, shortCircuit.Response, nil, preCount)
			if bifrostErr != nil {
				return nil, bifrostErr
			}
			return newBifrostMessageChan(resp), nil
		}
		// Handle short-circuit with stream
		if shortCircuit.Stream != nil {
			outputStream := make(chan *schemas.BifrostStream)

			// Create a post hook runner cause pipeline object is put back in the pool on defer
			pipelinePostHookRunner := func(ctx *context.Context, result *schemas.BifrostResponse, err *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError) {
				return pipeline.RunPostHooks(ctx, result, err, preCount)
			}

			go func() {
				defer close(outputStream)

				for streamMsg := range shortCircuit.Stream {
					if streamMsg == nil {
						continue
					}

					bifrostResponse := &schemas.BifrostResponse{}
					if streamMsg.BifrostTextCompletionResponse != nil {
						bifrostResponse.TextCompletionResponse = streamMsg.BifrostTextCompletionResponse
					}
					if streamMsg.BifrostChatResponse != nil {
						bifrostResponse.ChatResponse = streamMsg.BifrostChatResponse
					}
					if streamMsg.BifrostResponsesStreamResponse != nil {
						bifrostResponse.ResponsesStreamResponse = streamMsg.BifrostResponsesStreamResponse
					}
					if streamMsg.BifrostSpeechStreamResponse != nil {
						bifrostResponse.SpeechStreamResponse = streamMsg.BifrostSpeechStreamResponse
					}
					if streamMsg.BifrostTranscriptionStreamResponse != nil {
						bifrostResponse.TranscriptionStreamResponse = streamMsg.BifrostTranscriptionStreamResponse
					}

					// Run post hooks on the stream message
					processedResponse, processedError := pipelinePostHookRunner(&ctx, bifrostResponse, streamMsg.BifrostError)

					streamResponse := &schemas.BifrostStream{}
					if processedResponse != nil {
						streamResponse.BifrostTextCompletionResponse = processedResponse.TextCompletionResponse
						streamResponse.BifrostChatResponse = processedResponse.ChatResponse
						streamResponse.BifrostResponsesStreamResponse = processedResponse.ResponsesStreamResponse
						streamResponse.BifrostSpeechStreamResponse = processedResponse.SpeechStreamResponse
						streamResponse.BifrostTranscriptionStreamResponse = processedResponse.TranscriptionStreamResponse
					}
					if processedError != nil {
						streamResponse.BifrostError = processedError
					}

					// Send the processed message to the output stream
					outputStream <- streamResponse

					//TODO: Release the processed response immediately after use
				}
			}()

			return outputStream, nil
		}
		// Handle short-circuit with error
		if shortCircuit.Error != nil {
			resp, bifrostErr := pipeline.RunPostHooks(&ctx, nil, shortCircuit.Error, preCount)
			if bifrostErr != nil {
				return nil, bifrostErr
			}
			return newBifrostMessageChan(resp), nil
		}
	}
	if preReq == nil {
		bifrostErr := newBifrostErrorFromMsg("bifrost request after plugin hooks cannot be nil")
		bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
			RequestType:    req.RequestType,
			Provider:       provider,
			ModelRequested: model,
		}
		return nil, bifrostErr
	}

	msg := bifrost.getChannelMessage(*preReq)
	msg.Context = ctx

	select {
	case queue <- msg:
		// Message was sent successfully
	case <-ctx.Done():
		bifrost.releaseChannelMessage(msg)
		bifrostErr := newBifrostErrorFromMsg("request cancelled while waiting for queue space")
		bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
			RequestType:    req.RequestType,
			Provider:       provider,
			ModelRequested: model,
		}
		return nil, bifrostErr
	default:
		if bifrost.dropExcessRequests.Load() {
			bifrost.releaseChannelMessage(msg)
			bifrost.logger.Warn("Request dropped: queue is full, please increase the queue size or set dropExcessRequests to false")
			bifrostErr := newBifrostErrorFromMsg("request dropped: queue is full")
			bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
				RequestType:    req.RequestType,
				Provider:       provider,
				ModelRequested: model,
			}
			return nil, bifrostErr
		}
		select {
		case queue <- msg:
			// Message was sent successfully
		case <-ctx.Done():
			bifrost.releaseChannelMessage(msg)
			bifrostErr := newBifrostErrorFromMsg("request cancelled while waiting for queue space")
			bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
				RequestType:    req.RequestType,
				Provider:       provider,
				ModelRequested: model,
			}
			return nil, bifrostErr
		}
	}

	select {
	case stream := <-msg.ResponseStream:
		bifrost.releaseChannelMessage(msg)
		return stream, nil
	case bifrostErrVal := <-msg.Err:
		if bifrostErrVal.Error != nil {
			bifrost.logger.Debug("error while executing stream request: %s", bifrostErrVal.Error.Message)
		} else {
			bifrost.logger.Debug("error while executing stream request: %+v", bifrostErrVal)
		}
		// Marking final chunk
		ctx = context.WithValue(ctx, schemas.BifrostContextKeyStreamEndIndicator, true)
		// On error we will complete post-hooks
		recoveredResp, recoveredErr := pipeline.RunPostHooks(&ctx, nil, &bifrostErrVal, len(*bifrost.plugins.Load()))
		bifrost.releaseChannelMessage(msg)
		if recoveredErr != nil {
			return nil, recoveredErr
		}
		if recoveredResp != nil {
			return newBifrostMessageChan(recoveredResp), nil
		}
		return nil, &bifrostErrVal
	}
}

// executeRequestWithRetries is a generic function that handles common request processing logic
// It consolidates retry logic, backoff calculation, and error handling
// It is not a bifrost method because interface methods in go cannot be generic
func executeRequestWithRetries[T any](
	ctx *context.Context,
	config *schemas.ProviderConfig,
	requestHandler func() (T, *schemas.BifrostError),
	requestType schemas.RequestType,
	providerKey schemas.ModelProvider,
	model string,
) (T, *schemas.BifrostError) {
	var result T
	var bifrostError *schemas.BifrostError
	var attempts int

	for attempts = 0; attempts <= config.NetworkConfig.MaxRetries; attempts++ {
		*ctx = context.WithValue(*ctx, schemas.BifrostContextKeyNumberOfRetries, attempts)
		if attempts > 0 {
			// Log retry attempt
			var retryMsg string
			if bifrostError != nil && bifrostError.Error != nil {
				retryMsg = bifrostError.Error.Message
			} else if bifrostError != nil && bifrostError.StatusCode != nil {
				retryMsg = fmt.Sprintf("status=%d", *bifrostError.StatusCode)
				if bifrostError.Type != nil {
					retryMsg += ", type=" + *bifrostError.Type
				}
			}
			logger.Debug("retrying request (attempt %d/%d) for model %s: %s", attempts, config.NetworkConfig.MaxRetries, model, retryMsg)

			// Calculate and apply backoff
			backoff := calculateBackoff(attempts-1, config)
			logger.Debug("sleeping for %s", backoff)
			time.Sleep(backoff)
		}

		logger.Debug("attempting %s request for provider %s", requestType, providerKey)

		// Attempt the request
		result, bifrostError = requestHandler()

		logger.Debug("request %s for provider %s completed", requestType, providerKey)

		// Check if successful or if we should retry
		if bifrostError == nil ||
			bifrostError.IsBifrostError ||
			(bifrostError.Error != nil && bifrostError.Error.Type != nil && *bifrostError.Error.Type == schemas.RequestCancelled) {
			break
		}

		// Check if we should retry based on status code or error message
		shouldRetry := false

		if bifrostError.Error != nil && bifrostError.Error.Message == schemas.ErrProviderDoRequest {
			shouldRetry = true
			logger.Debug("detected request HTTP error, will retry: %s", bifrostError.Error.Message)
		}

		// Retry if status code or error object indicates rate limiting
		if (bifrostError.StatusCode != nil && retryableStatusCodes[*bifrostError.StatusCode]) ||
			(bifrostError.Error != nil &&
				(IsRateLimitErrorMessage(bifrostError.Error.Message) ||
					(bifrostError.Error.Type != nil && IsRateLimitErrorMessage(*bifrostError.Error.Type)))) {
			shouldRetry = true
			logger.Debug("detected rate limit error in message, will retry: %s", bifrostError.Error.Message)
		}

		if !shouldRetry {
			break
		}
	}

	// Add retry information to error
	if attempts > 0 {
		logger.Debug("request failed after %d %s", attempts, map[bool]string{true: "attempts", false: "attempt"}[attempts > 1])
	}

	return result, bifrostError
}

// requestWorker handles incoming requests from the queue for a specific provider.
// It manages retries, error handling, and response processing.
func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas.ProviderConfig, queue chan *ChannelMessage) {
	defer func() {
		if waitGroupValue, ok := bifrost.waitGroups.Load(provider.GetProviderKey()); ok {
			waitGroup := waitGroupValue.(*sync.WaitGroup)
			waitGroup.Done()
		}
	}()

	for req := range queue {
		_, model, _ := req.BifrostRequest.GetRequestFields()

		var result *schemas.BifrostResponse
		var stream chan *schemas.BifrostStream
		var bifrostError *schemas.BifrostError
		var err error

		// Determine the base provider type for key requirement checks
		baseProvider := provider.GetProviderKey()
		if cfg := config.CustomProviderConfig; cfg != nil && cfg.BaseProviderType != "" {
			baseProvider = cfg.BaseProviderType
		}

		key := schemas.Key{}
		if providerRequiresKey(baseProvider, config.CustomProviderConfig) {
			// Use the custom provider name for actual key selection, but pass base provider type for key validation
			key, err = bifrost.selectKeyFromProviderForModel(&req.Context, req.RequestType, provider.GetProviderKey(), model, baseProvider)
			if err != nil {
				bifrost.logger.Debug("error selecting key for model %s: %v", model, err)
				req.Err <- schemas.BifrostError{
					IsBifrostError: false,
					Error: &schemas.ErrorField{
						Message: err.Error(),
						Error:   err,
					},
					ExtraFields: schemas.BifrostErrorExtraFields{
						Provider:       provider.GetProviderKey(),
						ModelRequested: model,
						RequestType:    req.RequestType,
					},
				}
				continue
			}
			req.Context = context.WithValue(req.Context, schemas.BifrostContextKeySelectedKeyID, key.ID)
			req.Context = context.WithValue(req.Context, schemas.BifrostContextKeySelectedKeyName, key.Name)
		}
		// Create plugin pipeline for streaming requests outside retry loop to prevent leaks
		var postHookRunner schemas.PostHookRunner
		var pipeline *PluginPipeline
		if IsStreamRequestType(req.RequestType) {
			pipeline = bifrost.getPluginPipeline()
			postHookRunner = func(ctx *context.Context, result *schemas.BifrostResponse, err *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError) {
				resp, bifrostErr := pipeline.RunPostHooks(ctx, result, err, len(*bifrost.plugins.Load()))
				if bifrostErr != nil {
					return nil, bifrostErr
				}
				return resp, nil
			}
		}

		// Execute request with retries
		if IsStreamRequestType(req.RequestType) {
			stream, bifrostError = executeRequestWithRetries(&req.Context, config, func() (chan *schemas.BifrostStream, *schemas.BifrostError) {
				return bifrost.handleProviderStreamRequest(provider, req, key, postHookRunner)
			}, req.RequestType, provider.GetProviderKey(), model)
		} else {
			result, bifrostError = executeRequestWithRetries(&req.Context, config, func() (*schemas.BifrostResponse, *schemas.BifrostError) {
				return bifrost.handleProviderRequest(provider, req, key)
			}, req.RequestType, provider.GetProviderKey(), model)
		}

		if pipeline != nil {
			bifrost.releasePluginPipeline(pipeline)
		}

		if bifrostError != nil {
			bifrostError.ExtraFields = schemas.BifrostErrorExtraFields{
				Provider:       provider.GetProviderKey(),
				ModelRequested: model,
				RequestType:    req.RequestType,
			}

			// Send error with context awareness to prevent deadlock
			select {
			case req.Err <- *bifrostError:
				// Error sent successfully
			case <-req.Context.Done():
				// Client no longer listening, log and continue
				bifrost.logger.Debug("Client context cancelled while sending error response")
			case <-time.After(5 * time.Second):
				// Timeout to prevent indefinite blocking
				bifrost.logger.Warn("Timeout while sending error response, client may have disconnected")
			}
		} else {
			if IsStreamRequestType(req.RequestType) {
				// Send stream with context awareness to prevent deadlock
				select {
				case req.ResponseStream <- stream:
					// Stream sent successfully
				case <-req.Context.Done():
					// Client no longer listening, log and continue
					bifrost.logger.Debug("Client context cancelled while sending stream response")
				case <-time.After(5 * time.Second):
					// Timeout to prevent indefinite blocking
					bifrost.logger.Warn("Timeout while sending stream response, client may have disconnected")
				}
			} else {
				// Send response with context awareness to prevent deadlock
				select {
				case req.Response <- result:
					// Response sent successfully
				case <-req.Context.Done():
					// Client no longer listening, log and continue
					bifrost.logger.Debug("Client context cancelled while sending response")
				case <-time.After(5 * time.Second):
					// Timeout to prevent indefinite blocking
					bifrost.logger.Warn("Timeout while sending response, client may have disconnected")
				}
			}
		}
	}

	// bifrost.logger.Debug("worker for provider %s exiting...", provider.GetProviderKey())
}

// handleProviderRequest handles the request to the provider based on the request type
func (bifrost *Bifrost) handleProviderRequest(provider schemas.Provider, req *ChannelMessage, key schemas.Key) (*schemas.BifrostResponse, *schemas.BifrostError) {
	response := &schemas.BifrostResponse{}
	switch req.RequestType {
	case schemas.TextCompletionRequest:
		textCompletionResponse, bifrostError := provider.TextCompletion(req.Context, key, req.BifrostRequest.TextCompletionRequest)
		if bifrostError != nil {
			return nil, bifrostError
		}
		response.TextCompletionResponse = textCompletionResponse
	case schemas.ChatCompletionRequest:
		chatCompletionResponse, bifrostError := provider.ChatCompletion(req.Context, key, req.BifrostRequest.ChatRequest)
		if bifrostError != nil {
			return nil, bifrostError
		}
		response.ChatResponse = chatCompletionResponse
	case schemas.ResponsesRequest:
		responsesResponse, bifrostError := provider.Responses(req.Context, key, req.BifrostRequest.ResponsesRequest)
		if bifrostError != nil {
			return nil, bifrostError
		}
		response.ResponsesResponse = responsesResponse
	case schemas.EmbeddingRequest:
		embeddingResponse, bifrostError := provider.Embedding(req.Context, key, req.BifrostRequest.EmbeddingRequest)
		if bifrostError != nil {
			return nil, bifrostError
		}
		response.EmbeddingResponse = embeddingResponse
	case schemas.SpeechRequest:
		speechResponse, bifrostError := provider.Speech(req.Context, key, req.BifrostRequest.SpeechRequest)
		if bifrostError != nil {
			return nil, bifrostError
		}
		response.SpeechResponse = speechResponse
	case schemas.TranscriptionRequest:
		transcriptionResponse, bifrostError := provider.Transcription(req.Context, key, req.BifrostRequest.TranscriptionRequest)
		if bifrostError != nil {
			return nil, bifrostError
		}
		response.TranscriptionResponse = transcriptionResponse
	case schemas.ImageGenerationRequest:
		imageResponse, bifrostError := provider.ImageGeneration(req.Context, key, req.BifrostRequest.ImageGenerationRequest)
		if bifrostError != nil {
			return nil, bifrostError
		}
		response.ImageGenerationResponse = imageResponse
	default:
		_, model, _ := req.BifrostRequest.GetRequestFields()
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: fmt.Sprintf("unsupported request type: %s", req.RequestType),
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType:    req.RequestType,
				Provider:       provider.GetProviderKey(),
				ModelRequested: model,
			},
		}
	}
	return response, nil
}

// handleProviderStreamRequest handles the stream request to the provider based on the request type
func (bifrost *Bifrost) handleProviderStreamRequest(provider schemas.Provider, req *ChannelMessage, key schemas.Key, postHookRunner schemas.PostHookRunner) (chan *schemas.BifrostStream, *schemas.BifrostError) {
	switch req.RequestType {
	case schemas.TextCompletionStreamRequest:
		return provider.TextCompletionStream(req.Context, postHookRunner, key, req.BifrostRequest.TextCompletionRequest)
	case schemas.ChatCompletionStreamRequest:
		return provider.ChatCompletionStream(req.Context, postHookRunner, key, req.BifrostRequest.ChatRequest)
	case schemas.ResponsesStreamRequest:
		return provider.ResponsesStream(req.Context, postHookRunner, key, req.BifrostRequest.ResponsesRequest)
	case schemas.SpeechStreamRequest:
		return provider.SpeechStream(req.Context, postHookRunner, key, req.BifrostRequest.SpeechRequest)
	case schemas.TranscriptionStreamRequest:
		return provider.TranscriptionStream(req.Context, postHookRunner, key, req.BifrostRequest.TranscriptionRequest)
	case schemas.ImageGenerationStreamRequest:
		return provider.ImageGenerationStream(req.Context, postHookRunner, key, req.BifrostRequest.ImageGenerationRequest)
	default:
		_, model, _ := req.BifrostRequest.GetRequestFields()
		return nil, &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: fmt.Sprintf("unsupported request type: %s", req.RequestType),
			},
			ExtraFields: schemas.BifrostErrorExtraFields{
				RequestType:    req.RequestType,
				Provider:       provider.GetProviderKey(),
				ModelRequested: model,
			},
		}
	}
}

// PLUGIN MANAGEMENT

// RunPreHooks executes PreHooks in order, tracks how many ran, and returns the final request, any short-circuit decision, and the count.
func (p *PluginPipeline) RunPreHooks(ctx *context.Context, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.PluginShortCircuit, int) {
	var shortCircuit *schemas.PluginShortCircuit
	var err error
	pluginCtx, cancel := schemas.NewBifrostContextWithTimeout(*ctx, 10*time.Second)
	defer cancel()
	defer func() {
		*ctx = pluginCtx.GetParentCtxWithUserValues()
	}()
	for i, plugin := range p.plugins {
		p.logger.Debug("running pre-hook for plugin %s", plugin.GetName())
		req, shortCircuit, err = plugin.PreHook(pluginCtx, req)
		if err != nil {
			p.preHookErrors = append(p.preHookErrors, err)
			p.logger.Warn("error in PreHook for plugin %s: %v", plugin.GetName(), err)
		}
		p.executedPreHooks = i + 1
		if shortCircuit != nil {
			return req, shortCircuit, p.executedPreHooks // short-circuit: only plugins up to and including i ran
		}
	}
	return req, nil, p.executedPreHooks
}

// RunPostHooks executes PostHooks in reverse order for the plugins whose PreHook ran.
// Accepts the response and error, and allows plugins to transform either (e.g., recover from error, or invalidate a response).
// Returns the final response and error after all hooks. If both are set, error takes precedence unless error is nil.
// runFrom is the count of plugins whose PreHooks ran; PostHooks will run in reverse from index (runFrom - 1) down to 0
func (p *PluginPipeline) RunPostHooks(ctx *context.Context, resp *schemas.BifrostResponse, bifrostErr *schemas.BifrostError, runFrom int) (*schemas.BifrostResponse, *schemas.BifrostError) {
	// Defensive: ensure count is within valid bounds
	if runFrom < 0 {
		runFrom = 0
	}
	if runFrom > len(p.plugins) {
		runFrom = len(p.plugins)
	}
	var err error
	pluginCtx, cancel := schemas.NewBifrostContextWithTimeout(*ctx, 10*time.Second)
	defer cancel()
	for i := runFrom - 1; i >= 0; i-- {
		plugin := p.plugins[i]
		p.logger.Debug("running post-hook for plugin %s", plugin.GetName())
		resp, bifrostErr, err = plugin.PostHook(pluginCtx, resp, bifrostErr)
		if err != nil {
			p.postHookErrors = append(p.postHookErrors, err)
			p.logger.Warn("error in PostHook for plugin %s: %v", plugin.GetName(), err)
		}
		// If a plugin recovers from an error (sets bifrostErr to nil and sets resp), allow that
		// If a plugin invalidates a response (sets resp to nil and sets bifrostErr), allow that
	}
	// Capturing plugin ctx values and putting them in the request context
	*ctx = pluginCtx.GetParentCtxWithUserValues()
	// Final logic: if both are set, error takes precedence, unless error is nil
	if bifrostErr != nil {
		if resp != nil && bifrostErr.StatusCode == nil && bifrostErr.Error != nil && bifrostErr.Error.Type == nil &&
			bifrostErr.Error.Message == "" && bifrostErr.Error.Error == nil {
			// Defensive: treat as recovery if error is empty
			return resp, nil
		}
		return resp, bifrostErr
	}
	return resp, nil
}

// resetPluginPipeline resets a PluginPipeline instance for reuse
func (p *PluginPipeline) resetPluginPipeline() {
	p.executedPreHooks = 0
	p.preHookErrors = p.preHookErrors[:0]
	p.postHookErrors = p.postHookErrors[:0]
}

// getPluginPipeline gets a PluginPipeline from the pool and configures it
func (bifrost *Bifrost) getPluginPipeline() *PluginPipeline {
	pipeline := bifrost.pluginPipelinePool.Get().(*PluginPipeline)
	pipeline.plugins = *bifrost.plugins.Load()
	pipeline.logger = bifrost.logger
	return pipeline
}

// releasePluginPipeline returns a PluginPipeline to the pool
func (bifrost *Bifrost) releasePluginPipeline(pipeline *PluginPipeline) {
	pipeline.resetPluginPipeline()
	bifrost.pluginPipelinePool.Put(pipeline)
}

// POOL & RESOURCE MANAGEMENT

// getChannelMessage gets a ChannelMessage from the pool and configures it with the request.
// It also gets response and error channels from their respective pools.
func (bifrost *Bifrost) getChannelMessage(req schemas.BifrostRequest) *ChannelMessage {
	// Get channels from pool
	responseChan := bifrost.responseChannelPool.Get().(chan *schemas.BifrostResponse)
	errorChan := bifrost.errorChannelPool.Get().(chan schemas.BifrostError)

	// Clear any previous values to avoid leaking between requests
	select {
	case <-responseChan:
	default:
	}
	select {
	case <-errorChan:
	default:
	}

	// Get message from pool and configure it
	msg := bifrost.channelMessagePool.Get().(*ChannelMessage)
	msg.BifrostRequest = req
	msg.Response = responseChan
	msg.Err = errorChan

	// Conditionally allocate ResponseStream for streaming requests only
	if IsStreamRequestType(req.RequestType) {
		responseStreamChan := bifrost.responseStreamPool.Get().(chan chan *schemas.BifrostStream)
		// Clear any previous values to avoid leaking between requests
		select {
		case <-responseStreamChan:
		default:
		}
		msg.ResponseStream = responseStreamChan
	}

	return msg
}

// releaseChannelMessage returns a ChannelMessage and its channels to their respective pools.
func (bifrost *Bifrost) releaseChannelMessage(msg *ChannelMessage) {
	// Put channels back in pools
	bifrost.responseChannelPool.Put(msg.Response)
	bifrost.errorChannelPool.Put(msg.Err)

	// Return ResponseStream to pool if it was used
	if msg.ResponseStream != nil {
		// Drain any remaining channels to prevent memory leaks
		select {
		case <-msg.ResponseStream:
		default:
		}
		bifrost.responseStreamPool.Put(msg.ResponseStream)
	}

	// Release of Bifrost Request is handled in handle methods as they are required for fallbacks

	// Clear references and return to pool
	msg.Response = nil
	msg.ResponseStream = nil
	msg.Err = nil
	bifrost.channelMessagePool.Put(msg)
}

// resetBifrostRequest resets a BifrostRequest instance for reuse
func resetBifrostRequest(req *schemas.BifrostRequest) {
	req.RequestType = ""
	req.TextCompletionRequest = nil
	req.ChatRequest = nil
	req.ResponsesRequest = nil
	req.EmbeddingRequest = nil
	req.SpeechRequest = nil
	req.TranscriptionRequest = nil
	req.ImageGenerationRequest = nil
}

// getBifrostRequest gets a BifrostRequest from the pool
func (bifrost *Bifrost) getBifrostRequest() *schemas.BifrostRequest {
	req := bifrost.bifrostRequestPool.Get().(*schemas.BifrostRequest)
	return req
}

// releaseBifrostRequest returns a BifrostRequest to the pool
func (bifrost *Bifrost) releaseBifrostRequest(req *schemas.BifrostRequest) {
	resetBifrostRequest(req)
	bifrost.bifrostRequestPool.Put(req)
}

// getAllSupportedKeys retrieves all valid keys for a ListModels request.
// allowing the provider to aggregate results from multiple keys.
func (bifrost *Bifrost) getAllSupportedKeys(ctx *context.Context, providerKey schemas.ModelProvider, baseProviderType schemas.ModelProvider) ([]schemas.Key, error) {
	// Check if key has been set in the context explicitly
	if ctx != nil {
		key, ok := (*ctx).Value(schemas.BifrostContextKeyDirectKey).(schemas.Key)
		if ok {
			// If a direct key is specified, return it as a single-element slice
			return []schemas.Key{key}, nil
		}
	}

	keys, err := bifrost.account.GetKeysForProvider(ctx, providerKey)
	if err != nil {
		return nil, err
	}

	if len(keys) == 0 {
		return nil, fmt.Errorf("no keys found for provider: %v", providerKey)
	}

	// Filter keys for ListModels - only check if key has a value
	var supportedKeys []schemas.Key
	for _, k := range keys {
		if strings.TrimSpace(k.Value) != "" || canProviderKeyValueBeEmpty(baseProviderType) {
			supportedKeys = append(supportedKeys, k)
		}
	}

	if len(supportedKeys) == 0 {
		return nil, fmt.Errorf("no valid keys found for provider: %v", providerKey)
	}

	return supportedKeys, nil
}

// selectKeyFromProviderForModel selects an appropriate API key for a given provider and model.
// It uses weighted random selection if multiple keys are available.
func (bifrost *Bifrost) selectKeyFromProviderForModel(ctx *context.Context, requestType schemas.RequestType, providerKey schemas.ModelProvider, model string, baseProviderType schemas.ModelProvider) (schemas.Key, error) {
	// Check if key has been set in the context explicitly
	if ctx != nil {
		key, ok := (*ctx).Value(schemas.BifrostContextKeyDirectKey).(schemas.Key)
		if ok {
			return key, nil
		}
	}

	if skipKeySelection, ok := (*ctx).Value(schemas.BifrostContextKeySkipKeySelection).(bool); ok && skipKeySelection && isKeySkippingAllowed(providerKey) {
		return schemas.Key{}, nil
	}

	keys, err := bifrost.account.GetKeysForProvider(ctx, providerKey)
	if err != nil {
		return schemas.Key{}, err
	}

	if len(keys) == 0 {
		return schemas.Key{}, fmt.Errorf("no keys found for provider: %v and model: %s", providerKey, model)
	}

	// filter out keys which dont support the model, if the key has no models, it is supported for all models
	var supportedKeys []schemas.Key
	if requestType == schemas.ListModelsRequest {
		// Skip deployment check but still check if the key has a value
		for _, k := range keys {
			if strings.TrimSpace(k.Value) != "" || canProviderKeyValueBeEmpty(baseProviderType) {
				supportedKeys = append(supportedKeys, k)
			}
		}
	} else {
		for _, key := range keys {
			modelSupported := (slices.Contains(key.Models, model) && (strings.TrimSpace(key.Value) != "" || canProviderKeyValueBeEmpty(baseProviderType))) || len(key.Models) == 0

			// Additional deployment checks for Azure, Bedrock and Vertex
			deploymentSupported := true
			if baseProviderType == schemas.Azure && key.AzureKeyConfig != nil {
				// For Azure, check if deployment exists for this model
				if len(key.AzureKeyConfig.Deployments) > 0 {
					_, deploymentSupported = key.AzureKeyConfig.Deployments[model]
				}
			} else if baseProviderType == schemas.Bedrock && key.BedrockKeyConfig != nil {
				// For Bedrock, check if deployment exists for this model
				if len(key.BedrockKeyConfig.Deployments) > 0 {
					_, deploymentSupported = key.BedrockKeyConfig.Deployments[model]
				}
			} else if baseProviderType == schemas.Vertex && key.VertexKeyConfig != nil {
				// For Vertex, check if deployment exists for this model
				if len(key.VertexKeyConfig.Deployments) > 0 {
					_, deploymentSupported = key.VertexKeyConfig.Deployments[model]
				}
			}

			if modelSupported && deploymentSupported {
				supportedKeys = append(supportedKeys, key)
			}
		}
	}
	if len(supportedKeys) == 0 {
		if baseProviderType == schemas.Azure || baseProviderType == schemas.Bedrock || baseProviderType == schemas.Vertex {
			return schemas.Key{}, fmt.Errorf("no keys found that support model/deployment: %s", model)
		}
		return schemas.Key{}, fmt.Errorf("no keys found that support model: %s", model)
	}

	if len(supportedKeys) == 1 {
		return supportedKeys[0], nil
	}

	selectedKey, err := bifrost.keySelector(ctx, supportedKeys, providerKey, model)
	if err != nil {
		return schemas.Key{}, err
	}

	return selectedKey, nil

}

func WeightedRandomKeySelector(ctx *context.Context, keys []schemas.Key, providerKey schemas.ModelProvider, model string) (schemas.Key, error) {
	// Use a weighted random selection based on key weights
	totalWeight := 0
	for _, key := range keys {
		totalWeight += int(key.Weight * 100) // Convert float to int for better performance
	}

	// Use a fast random number generator
	randomSource := rand.New(rand.NewSource(time.Now().UnixNano()))
	randomValue := randomSource.Intn(totalWeight)

	// Select key based on weight
	currentWeight := 0
	for _, key := range keys {
		currentWeight += int(key.Weight * 100)
		if randomValue < currentWeight {
			return key, nil
		}
	}

	// Fallback to first key if something goes wrong
	return keys[0], nil
}

// Shutdown gracefully stops all workers when triggered.
// It closes all request channels and waits for workers to exit.
func (bifrost *Bifrost) Shutdown() {
	bifrost.logger.Info("closing all request channels...")
	// Check if the context is done
	if bifrost.ctx.Err() != nil {
		bifrost.logger.Warn("context already done, skipping cancel")
		return
	} else if bifrost.cancel != nil {
		bifrost.cancel()
	}
	// Close all provider queues to signal workers to stop
	bifrost.requestQueues.Range(func(key, value interface{}) bool {
		close(value.(chan *ChannelMessage))
		return true
	})

	// Wait for all workers to exit
	bifrost.waitGroups.Range(func(key, value interface{}) bool {
		waitGroup := value.(*sync.WaitGroup)
		waitGroup.Wait()
		return true
	})

	// Cleanup MCP manager
	if bifrost.mcpManager != nil {
		err := bifrost.mcpManager.cleanup()
		if err != nil {
			bifrost.logger.Warn(fmt.Sprintf("Error cleaning up MCP manager: %s", err.Error()))
		}
	}

	// Cleanup plugins
	for _, plugin := range *bifrost.plugins.Load() {
		err := plugin.Cleanup()
		if err != nil {
			bifrost.logger.Warn(fmt.Sprintf("Error cleaning up plugin: %s", err.Error()))
		}
	}
	bifrost.logger.Info("all request channels closed")
}