package workflow

import (
	"agent/src/llm"
	"agent/src/message"
	"agent/src/prompt"
	"agent/src/utils/log"
	"bytes"
	"context"
	"errors"
	"strings"

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


func (task *Task) chatWithStream(msg string, withThink bool) error {
	buf := new(bytes.Buffer)
	ch := task.client.GetChannel()
	var seqId uint64
	var isThink bool
	start := true
	
	option := llms.WithStreamingFunc(func(ctx context.Context, chunk []byte) error {
		select {
			case <-task.ctx.Done(): 
				return errors.New("task cancelled")
			default:
		}
		
		if !withThink && start {
			start = false
			if strings.HasPrefix(string(chunk), prompt.ThinkStart) {
				isThink = true
				return nil
			}
		}

		if !withThink && isThink {
			if strings.HasSuffix(string(chunk), prompt.ThinkEnd) {
				isThink = false
			}
			return nil
		}

		buf.Write(chunk)
		// 避免client网络太慢
		if len(ch) == 0 {
			err := task.client.WriteJson(&message.ResponseExecSuccess{
				RequestId    : task.request_id,
				Msg          : buf.String(),
				Stream_SeqId : &seqId,
			}, true)
			if err != nil {
				return err
			}
			buf.Reset()
			seqId++
		}
		return nil
	})

	_, err := task.pool.llmGroup.Chat(task.ctx, msg, option)
	if err != nil {
		return err
	}

	return task.client.WriteJson(&message.ResponseExecSuccess{
		RequestId    : task.request_id,
		Msg          : buf.String(),
		Stream_SeqId : &seqId,
		Stream_Finish: true,
	}, true)
}

func (task *Task) chat(msg string, stream bool) error {
	if stream {
		return task.chatWithStream(msg, false)
	} else {
		answer, err := task.pool.llmGroup.Chat(task.ctx, msg)
		if err != nil {
			return err
		}
		answer = prompt.RemoveThinkPart(answer)

		task.client.WriteJson(&message.ResponseExecSuccess{
			RequestId: task.request_id,
			Msg      : answer,
		}, true)
	}
	return nil
}

func (task *Task) exec_chat() error {
	log.Info("call chat")
	req, ok := task.request.(*message.RequestExecChat)
	if !ok {
		return errors.New("request type wrong")
	}
	return task.chat(req.Msg, req.Stream)
}

func (task *Task) listModel(requestId uint64) error {
	err := task.client.WriteJson(&message.ResponseListModel{
		RequestId: requestId,
		Models   : llm.LLMNames,
	}, false)
	return err
}