package agent

import (
	"context"
	"time"

	"gitea.stevedudenhoeffer.com/steve/answer/pkg/toolbox"

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

type Context struct {
	context.Context

	Messages []llms.MessageContent
	Agent    *Agent
}

func From(ctx context.Context, a *Agent) *Context {
	if ctx == nil {
		ctx = context.Background()
	}

	return &Context{
		Context: ctx,
		Agent:   a,
	}
}

func (c *Context) WithMessage(m llms.MessageContent) *Context {
	c.Messages = append(c.Messages, m)
	return c
}

func (c *Context) WithMessages(m ...llms.MessageContent) *Context {
	c.Messages = append(c.Messages, m...)
	return c
}

func (c *Context) WithToolResults(r ...toolbox.ToolResult) *Context {
	msg := llms.MessageContent{
		Role: llms.ChatMessageTypeTool,
	}

	for _, v := range r {
		res := v.Result
		if v.Error != nil {
			res = "error executing: " + v.Error.Error()
		}
		msg.Parts = append(msg.Parts, llms.ToolCallResponse{
			ToolCallID: v.ID,
			Name:       v.Name,
			Content:    res,
		})
	}

	return c.WithMessage(msg)
}

func (c *Context) WithAgent(a *Agent) *Context {
	return &Context{
		Context: c.Context,
		Agent:   a,
	}
}

func (c *Context) WithCancel() (*Context, func()) {
	ctx, cancel := context.WithCancel(c.Context)
	return &Context{
		Context: ctx,
		Agent:   c.Agent,
	}, cancel
}

func (c *Context) WithDeadline(deadline time.Time) (*Context, func()) {
	ctx, cancel := context.WithDeadline(c.Context, deadline)
	return &Context{
		Context: ctx,
		Agent:   c.Agent,
	}, cancel
}

func (c *Context) WithTimeout(timeout time.Duration) (*Context, func()) {
	ctx, cancel := context.WithTimeout(c.Context, timeout)
	return &Context{
		Context: ctx,
		Agent:   c.Agent,
	}, cancel
}

func (c *Context) WithValue(key, value interface{}) *Context {
	return &Context{
		Context: context.WithValue(c.Context, key, value),
		Agent:   c.Agent,
	}
}

func (c *Context) Done() <-chan struct{} {
	return c.Context.Done()
}

func (c *Context) Err() error {
	return c.Context.Err()
}

func (c *Context) Value(key interface{}) interface{} {
	return c.Context.Value(key)
}

func (c *Context) Deadline() (time.Time, bool) {
	return c.Context.Deadline()
}