package llms

type OutputFinishType int

const (
	OutputFinishInvalid       OutputFinishType = 0
	OutputFinishNormal        OutputFinishType = 1 // 输出内容完全由大模型生成，未触发截断、替换
	OutputFinishStop          OutputFinishType = 2 // 输出结果命中入参 stop 中指定的字段后被截断
	OutputFinishLength        OutputFinishType = 3 // 达到了最大的 token 数
	OutputFinishContentFilter OutputFinishType = 4 // 输出内容被截断、兜底、替换为**等
	OutputFinishToolCalls     OutputFinishType = 5 // 调用了 tool call 功能
)

/**
Input
*/

type Input struct {
	sys  *SystemMessage
	prev *Input
	curr []Message
	deep int
}

func (i *Input) SetSysMsg(name, content string) {
	if content != "" {
		i.sys = NewSystemMessage(name, content)
	} else {
		i.sys = nil
	}
}

func (i *Input) GetSysMsg() *SystemMessage {
	return i.sys
}

func (i *Input) Push(fn ConvertMessage) error {
	if i.curr == nil {
		i.curr = make([]Message, 0, 8)
	}
	if cur, err := fn(); err != nil {
		return err
	} else {
		i.curr = append(i.curr, cur...)
	}
	return nil
}

func (i *Input) CurrValues() []Message {
	return i.curr
}

func (i *Input) TotalValues() []Message {
	inputChain, l := make([]*Input, i.deep+1), 0
	for curr, idx := i, i.deep; curr != nil && idx >= 0; curr, idx = i.prev, idx-1 {
		inputChain[idx], l = curr, l+len(curr.curr)
	}

	msgList := make([]Message, 0, l)
	for _, v := range inputChain {
		msgList = append(msgList, v.curr...)
	}
	return msgList
}

func (i *Input) LastMsg() Message {
	if l := len(i.curr); l > 0 {
		return i.curr[l-1]
	}
	return nil
}

func (i *Input) Next() *Input {
	return &Input{
		prev: i,
		deep: i.deep + 1,
		curr: nil,
	}
}

/**
Output
*/

type OutputMsg struct {
	msgType      MessageType
	name         string
	content      string
	fnList       []*FunctionCall
	finishReason OutputFinishType
}

func (om *OutputMsg) MsgType() MessageType {
	return om.msgType
}

func (om *OutputMsg) Content() string {
	if om.content == "" && len(om.fnList) > 0 {
		return om.fnList[len(om.fnList)-1].Result
	}
	return om.content
}

func (om *OutputMsg) FinishReason() OutputFinishType {
	return om.finishReason
}

func (om *OutputMsg) FunctionCallList() []FunctionCall {
	ret := make([]FunctionCall, 0, len(om.fnList))
	for _, v := range om.fnList {
		ret = append(ret, *v)
	}
	return ret
}

type Output struct {
	content  []*OutputMsg
	usage    TokenUsage
	children []*Output
}

func (o *Output) LastMsg() *OutputMsg {
	if l := len(o.content); l > 0 {
		v := *o.content[l-1]
		return &v
	}
	return nil
}

func (o *Output) SetChild(children ...*Output) {
	if len(children) > 0 {
		if o.children == nil {
			o.children = make([]*Output, 0, 4)
		}
		o.children = append(o.children, children...)
	}
}

func (o *Output) MergeMessage(oList ...*Output) {
	if len(oList) > 0 {
		if o.content == nil {
			o.content = make([]*OutputMsg, 0, 4)
		}

		for _, n := range oList {
			// 结果拼接
			o.content = append(o.content, n.content...)

			// token 统计
			o.usage.Prompt += n.usage.Prompt
			o.usage.Completion += n.usage.Completion
			o.usage.Total += n.usage.Total
		}
	}
	return
}

func (o *Output) GetCurrUsage() TokenUsage {
	return o.usage
}

func (o *Output) GetTotalUsage() TokenUsage {
	u := o.GetCurrUsage()
	for _, c := range o.children {
		cu := c.GetTotalUsage()

		u.Total += cu.Total
		u.Prompt += cu.Prompt
		u.Completion += cu.Completion
	}
	return u
}

func NewOutputMsg(msgType MessageType, name, content string, finishReason OutputFinishType, fnList ...*FunctionCall) *OutputMsg {
	return &OutputMsg{
		msgType:      msgType,
		name:         name,
		content:      content,
		finishReason: finishReason,
		fnList:       fnList,
	}
}

func NewOutput(usage TokenUsage, content []*OutputMsg) *Output {
	return &Output{
		content: content,
		usage:   usage,
	}
}