package utils_vllm

import (
	"ai_scheduler/internal/config"
	"context"

	"github.com/cloudwego/eino-ext/components/model/openai"
	"github.com/cloudwego/eino/schema"
)

type Client struct {
	model  *openai.ChatModel
	config *config.Config
}

func NewClient(config *config.Config) (*Client, func(), error) {
	m, err := openai.NewChatModel(context.Background(), &openai.ChatModelConfig{
		BaseURL: config.Vllm.BaseURL,
		Model:   config.Vllm.VlModel,
		Timeout: config.Vllm.Timeout,
	})
	if err != nil {
		return nil, nil, err
	}
	c := &Client{model: m, config: config}
	cleanup := func() {}
	return c, cleanup, nil
}

func (c *Client) Chat(ctx context.Context, msgs []*schema.Message) (*schema.Message, error) {
	return c.model.Generate(ctx, msgs)
}

func (c *Client) RecognizeWithImg(ctx context.Context, systemPrompt, userPrompt string, imgURLs []string) (*schema.Message, error) {
	in := []*schema.Message{
		{
			Role:    schema.System,
			Content: systemPrompt,
		},
		{
			Role: schema.User,
		},
	}
	parts := []schema.MessageInputPart{
		{Type: schema.ChatMessagePartTypeText, Text: userPrompt},
	}
	for i := range imgURLs {
		u := imgURLs[i]
		parts = append(parts, schema.MessageInputPart{
			Type: schema.ChatMessagePartTypeImageURL,
			Image: &schema.MessageInputImage{
				MessagePartCommon: schema.MessagePartCommon{URL: &u},
				Detail:            schema.ImageURLDetailHigh,
			},
		})
	}

	in[1].UserInputMultiContent = parts
	return c.model.Generate(ctx, in)
}