package go_llm

import (
	"github.com/openai/openai-go"
)

type ResponseChoice struct {
	Index   int
	Role    Role
	Content string
	Refusal string
	Name    string
	Calls   []ToolCall
}

func (r ResponseChoice) toRaw() map[string]any {
	res := map[string]any{
		"index":   r.Index,
		"role":    r.Role,
		"content": r.Content,
		"refusal": r.Refusal,
		"name":    r.Name,
	}

	calls := make([]map[string]any, 0, len(r.Calls))
	for _, call := range r.Calls {
		calls = append(calls, call.toRaw())
	}

	res["tool_calls"] = calls

	return res
}

func (r ResponseChoice) toChatCompletionMessages(_ string) []openai.ChatCompletionMessageParamUnion {
	var as openai.ChatCompletionAssistantMessageParam

	if r.Name != "" {
		as.Name = openai.String(r.Name)
	}
	if r.Refusal != "" {
		as.Refusal = openai.String(r.Refusal)
	}

	if r.Content != "" {
		as.Content.OfString = openai.String(r.Content)
	}

	for _, call := range r.Calls {
		as.ToolCalls = append(as.ToolCalls, openai.ChatCompletionMessageToolCallParam{
			ID: call.ID,
			Function: openai.ChatCompletionMessageToolCallFunctionParam{
				Name:      call.FunctionCall.Name,
				Arguments: call.FunctionCall.Arguments,
			},
		})
	}
	return []openai.ChatCompletionMessageParamUnion{
		{
			OfAssistant: &as,
		},
	}
}

func (r ResponseChoice) toInput() []Input {
	var res []Input

	for _, call := range r.Calls {
		res = append(res, call)
	}

	if r.Content != "" || r.Refusal != "" {
		res = append(res, Message{
			Role: RoleAssistant,
			Text: r.Content,
		})
	}

	return res
}

type Response struct {
	Choices []ResponseChoice
}