package capacity

import (
	"context"

	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
	"gitea.maximumdirect.net/eric/promptkit/internal/llm"
)

type client struct {
	manager *Manager
	next    llm.Client
}

// NewClient wraps next with configured active-generation limits. A nil manager
// leaves next unchanged.
func NewClient(manager *Manager, next llm.Client) llm.Client {
	if manager == nil {
		return next
	}
	return &client{
		manager: manager,
		next:    next,
	}
}

func (c *client) Generate(
	ctx context.Context,
	req domain.GenerateRequest,
) (*domain.GenerateResponse, error) {
	if c.manager.getPool(req.Target.BackendID) == nil {
		return c.next.Generate(ctx, req)
	}
	release, err := c.manager.Acquire(ctx, req.Target.BackendID)
	if err != nil {
		return nil, err
	}
	defer release()
	return c.next.Generate(ctx, req)
}