package chain

import (
	"context"
	"fmt"
	"gitee.com/kzangv/gsf-ai-agent/llms"
	"gitee.com/kzangv/gsf-ai-agent/utils"
	"reflect"
)

/**
Chain
*/

type Chain interface {
	Desc() string

	SetTraceLog(utils.TraceLog) error
	GetTraceLog() utils.TraceLog

	SetOption(Option) error
	GetOption() Option

	Init() error
	Run(context.Context, *llms.Input) (*llms.Output, error)
}

type SimpleChain struct {
	opt Option
	log utils.TraceLog
}

func (ch *SimpleChain) Desc() string {
	panic("SimpleChain no Desc() implement ")
}

func (ch *SimpleChain) SetTraceLog(log utils.TraceLog) error {
	if log == nil {
		ch.log = utils.DefaultTraceLog
	} else if rv := reflect.ValueOf(log); rv.Type().Kind() == reflect.Pointer && rv.IsNil() {
		ch.log = utils.DefaultTraceLog
	} else {
		ch.log = log
	}
	return nil
}

func (ch *SimpleChain) GetTraceLog() utils.TraceLog {
	return ch.log
}

func (ch *SimpleChain) SetOption(opt Option) error {
	ch.opt = opt
	return nil
}

func (ch *SimpleChain) GetOption() Option {
	return ch.opt
}

func (ch *SimpleChain) Init() error {
	if ch.opt == nil {
		ch.opt = NewOption()
	}
	if opt := ch.opt.GetParamManager(); opt == nil {
		ch.opt.SetParamManager(utils.NewSimpleChainParamsManager())
	} else if rv := reflect.ValueOf(opt); rv.Type().Kind() == reflect.Pointer && rv.IsNil() {
		ch.opt.SetParamManager(utils.NewSimpleChainParamsManager())
	}
	if ch.opt.GetOutput() == "" {
		ch.opt.SetOutput(DefaultOptionOutputKey)
	}
	return nil
}

func (ch *SimpleChain) Run(_ context.Context, _ *llms.Input) (*llms.Output, error) {
	panic("SimpleChain no Run implement")
}

func InitChain(ch Chain, log utils.TraceLog, opt Option) (err error) {
	for {
		if err = ch.SetTraceLog(log); err != nil {
			break
		}
		if err = ch.SetOption(opt); err != nil {
			break
		}
		if err = ch.Init(); err != nil {
			break
		}
		break
	}
	return
}

/**
LLMChain
*/

type LLMChain interface {
	Chain

	GetLLM() llms.LLM
	SetLLM(llms.LLM) error
}

type SimpleLLMChain struct {
	SimpleChain
	llm llms.LLM
}

func (ch *SimpleLLMChain) GetLLM() llms.LLM {
	return ch.llm
}

func (ch *SimpleLLMChain) SetLLM(llm llms.LLM) error {
	ch.llm = llm
	return nil
}

func (ch *SimpleLLMChain) Init() error {
	if ch.llm == nil {
		return fmt.Errorf("SimpleLLMChain.llm is nil")
	}
	if err := ch.SimpleChain.Init(); err != nil {
		return err
	}
	return nil
}

func InitLLMChain(ch LLMChain, llm llms.LLM, log utils.TraceLog, opt Option) (err error) {
	for {
		if err = ch.SetLLM(llm); err != nil {
			break
		}
		if err = InitChain(ch, log, opt); err != nil {
			break
		}
		break
	}
	return
}

/**
_ChainNode
*/

type _ChainNode struct {
	Chain Chain
	Desc  string
}

func NewChainNode(desc string, ch Chain) *_ChainNode {
	return &_ChainNode{
		Chain: ch,
		Desc:  desc,
	}
}