package llm

import (
	"agent/src/utils/log"
	"context"
	"strings"

	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/llms/ollama"
)

var LLMGroups map[string]*LLMGroup
var LLMNames []string

// const DefaultOllamaUrl = "http://192.168.1.5:11434"
const DefaultOllamaUrl = "http://ishangsf.com:11434"

func InitLocalModel(urls []string) error {
	model2url := make(map[string][]string)
	for _, url := range urls {
		models, err := listOllamaModels(url)
		if err != nil {
			return err
		}
		for _, m := range models {
			model2url[m] = append(model2url[m], url)
		}
	}
	
	LLMNames = make([]string, 0, len(model2url))
	LLMGroups = make(map[string]*LLMGroup, len(model2url))

	var models []string
	for model, urls := range model2url {
		localModel := "local " + model
		g := &LLMGroup{
			name: localModel,
			llm: make(chan llms.Model, len(urls)),
			local: true,
		}
		for _, url := range urls {
			llm, err := ollama.New(ollama.WithServerURL(url), ollama.WithModel(model))
			if err != nil {
				return err
			}
			g.llm <- llm

		}

		LLMGroups[localModel] = g
		LLMNames = append(LLMNames, localModel)
		models = append(models, model)
	}
	log.Info("load models", "models", strings.Join(models, ","))
	return nil
}


/*
多个LLM组成一个集群
*/
type LLMGroup struct {
	name  string
	llm   chan llms.Model
	local bool
}

func (g *LLMGroup) Name() string { return g.name }
func (g *LLMGroup) IsLocal() bool { return g.local }
func (g *LLMGroup) Cap() int { return cap(g.llm) }

func (g *LLMGroup) GenerateContent(ctx context.Context, messages []llms.MessageContent, options ...llms.CallOption) (*llms.ContentResponse, error) {
	m := <-g.llm
	defer func(){
		g.llm <- m
	}()
	return m.GenerateContent(ctx, messages, options...)
}

// use streaming if stream is not nil
func (g *LLMGroup) Chat(ctx context.Context, msg string, options ...llms.CallOption) (string, error) {
	contents := []llms.MessageContent{llms.TextParts(llms.ChatMessageTypeHuman, msg)}
	return g.ChatContents(ctx, contents, options...)
}

// use streaming if stream is not nil
func (g *LLMGroup) ChatContents(ctx context.Context, contents []llms.MessageContent, options ...llms.CallOption) (string, error) {
	m := <-g.llm
	defer func(){
		g.llm <- m
	}()

	options = append(options, llms.WithN(1))
	res, err := m.GenerateContent(ctx, contents, options...)
	if err != nil {
		return "", err
	}
	if len(res.Choices) == 0 {
		return "", nil
	}
	return res.Choices[0].Content, nil
}

// block util available
func (g LLMGroup) Get() llms.Model {
	return <-g.llm
}

func ContainLLM(name string) bool {
	_, ok := LLMGroups[name]
	return ok
}
