package llm_service

import (
	"ai_scheduler/internal/config"
	"ai_scheduler/internal/entitys"
	"ai_scheduler/internal/pkg"
	"ai_scheduler/internal/pkg/utils_vllm"
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"strings"

	"github.com/cloudwego/eino/schema"
	"github.com/ollama/ollama/api"
)

type VllmService struct {
	client *utils_vllm.Client
	config *config.Config
}

func NewVllmService(
	client *utils_vllm.Client,
	config *config.Config,
) *VllmService {
	return &VllmService{
		client: client,
		config: config,
	}
}

func (s *VllmService) IntentRecognize(ctx context.Context, req *entitys.ToolSelect) (msg string, err error) {
	msgs := s.convertMessages(req.Prompt)
	tools := s.convertTools(req.Tools)

	resp, err := s.client.ToolSelect(ctx, msgs, tools)
	if err != nil {
		return
	}

	if resp.Content == "" {
		if len(resp.ToolCalls) > 0 {
			call := resp.ToolCalls[0]
			var matchFromTools = &entitys.Match{
				Confidence: 1,
				Index:      call.Function.Name,
				Parameters: call.Function.Arguments,
				IsMatch:    true,
			}
			msg = pkg.JsonStringIgonErr(matchFromTools)
		} else {
			err = errors.New("不太明白你想表达的意思呢，可以在仔细描述一下您所需要的内容吗，感谢感谢")
			return
		}
	} else {
		msg = resp.Content
	}
	return
}

func (s *VllmService) convertMessages(prompts []api.Message) []*schema.Message {
	msgs := make([]*schema.Message, 0, len(prompts))
	for _, p := range prompts {
		msg := &schema.Message{
			Role:    schema.RoleType(p.Role),
			Content: p.Content,
		}

		// 这里实际应该不会走进来
		if len(p.Images) > 0 {
			parts := []schema.MessageInputPart{
				{Type: schema.ChatMessagePartTypeText, Text: p.Content},
			}
			for _, imgData := range p.Images {
				b64 := base64.StdEncoding.EncodeToString(imgData)
				mimeType := "image/jpeg"
				parts = append(parts, schema.MessageInputPart{
					Type: schema.ChatMessagePartTypeImageURL,
					Image: &schema.MessageInputImage{
						MessagePartCommon: schema.MessagePartCommon{
							MIMEType:   mimeType,
							Base64Data: &b64,
						},
					},
				})
			}
			msg.UserInputMultiContent = parts
		}
		msgs = append(msgs, msg)
	}
	return msgs
}

func (s *VllmService) convertTools(tasks []entitys.RegistrationTask) []*schema.ToolInfo {
	tools := make([]*schema.ToolInfo, 0, len(tasks))
	for _, task := range tasks {
		params := make(map[string]*schema.ParameterInfo)
		for k, v := range task.TaskConfigDetail.Param.Properties {
			dt := schema.String

			// Handle v.Type dynamically to support both string and []string (compiler suggests []string)
			// Using fmt.Sprint handles both cases safely without knowing exact type structure
			typeStr := fmt.Sprintf("%v", v.Type)
			typeStr = strings.Trim(typeStr, "[]") // normalize "[string]" -> "string"

			switch typeStr {
			case "string":
				dt = schema.String
			case "integer", "int":
				dt = schema.Integer
			case "number", "float":
				dt = schema.Number
			case "boolean", "bool":
				dt = schema.Boolean
			case "object":
				dt = schema.Object
			case "array":
				dt = schema.Array
			}

			required := false
			for _, r := range task.TaskConfigDetail.Param.Required {
				if r == k {
					required = true
					break
				}
			}

			desc := v.Description
			if len(v.Enum) > 0 {
				var enumStrs []string
				for _, e := range v.Enum {
					enumStrs = append(enumStrs, fmt.Sprintf("%v", e))
				}
				desc += " Enum: " + strings.Join(enumStrs, ", ")
			}

			params[k] = &schema.ParameterInfo{
				Type:     dt,
				Desc:     desc,
				Required: required,
			}
		}

		tools = append(tools, &schema.ToolInfo{
			Name:        task.Name,
			Desc:        task.Desc,
			ParamsOneOf: schema.NewParamsOneOfByParams(params),
		})
	}
	return tools
}