package go_llm

import (
	"context"
	"time"
)

type Context struct {
	context.Context
	request         Request
	response        *ResponseChoice
	toolcall        *ToolCall
	syntheticFields map[string]string
}

func (c *Context) ToNewRequest(toolResults ...ToolCallResponse) Request {
	var res Request

	res.Toolbox = c.request.Toolbox
	res.Temperature = c.request.Temperature

	res.Conversation = make([]Input, len(c.request.Conversation))
	copy(res.Conversation, c.request.Conversation)

	// now for every input message, convert those to an Input to add to the conversation
	for _, msg := range c.request.Messages {
		res.Conversation = append(res.Conversation, msg)
	}

	// if there are tool calls, then we need to add those to the conversation
	if c.response != nil {
		res.Conversation = append(res.Conversation, *c.response)
	}

	// if there are tool results, then we need to add those to the conversation
	for _, result := range toolResults {
		res.Conversation = append(res.Conversation, result)
	}

	return res
}

func NewContext(ctx context.Context, request Request, response *ResponseChoice, toolcall *ToolCall) *Context {
	return &Context{Context: ctx, request: request, response: response, toolcall: toolcall}
}

func (c *Context) Request() Request {
	return c.request
}

func (c *Context) Response() *ResponseChoice {
	return c.response
}

func (c *Context) ToolCall() *ToolCall {
	return c.toolcall
}

func (c *Context) SyntheticFields() map[string]string {
	if c.syntheticFields == nil {
		c.syntheticFields = map[string]string{}
	}

	return c.syntheticFields
}

func (c *Context) WithContext(ctx context.Context) *Context {
	return &Context{Context: ctx, request: c.request, response: c.response, toolcall: c.toolcall, syntheticFields: c.syntheticFields}
}

func (c *Context) WithRequest(request Request) *Context {
	return &Context{Context: c.Context, request: request, response: c.response, toolcall: c.toolcall, syntheticFields: c.syntheticFields}
}

func (c *Context) WithResponse(response *ResponseChoice) *Context {
	return &Context{Context: c.Context, request: c.request, response: response, toolcall: c.toolcall, syntheticFields: c.syntheticFields}
}

func (c *Context) WithToolCall(toolcall *ToolCall) *Context {
	return &Context{Context: c.Context, request: c.request, response: c.response, toolcall: toolcall, syntheticFields: c.syntheticFields}
}

func (c *Context) WithSyntheticFields(syntheticFields map[string]string) *Context {
	return &Context{Context: c.Context, request: c.request, response: c.response, toolcall: c.toolcall, syntheticFields: syntheticFields}
}

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

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

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

func (c *Context) Value(key any) any {
	switch key {
	case "request":
		return c.request

	case "response":
		return c.response

	case "toolcall":
		return c.toolcall

	case "syntheticFields":
		return c.syntheticFields

	}
	return c.Context.Value(key)
}

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