package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
)

// errChainLogWrite marks an error of a log write, which ends the chain.
var errChainLogWrite = errors.New("log write failed")

// ChainClient sends model calls.
type ChainClient interface {
	Create(ctx context.Context, request ModelRequest) (ModelResponse, error)
}

// Chain is one request chain of a turn or of a subagent. It sends request with
// the input of the first request, runs the tool calls of each response, and
// sends their results back, until a response has no tool call: the final
// answer. The replies, the tool calls, and the results go to the log of tools
// as they happen. The thoughts never do.
//
// Each later request begins with the whole input of the request before it,
// then its output items, then the tool results. The output items hold the
// reasoning with its encrypted content, because the calls do not store it. So
// each request begins with all of the previous one and reads it from the prompt
// cache. The messages of take join each later request after the tool results.
type Chain struct {
	client  ChainClient
	tools   *Toolbox
	request ModelRequest
	take    func() []Message
}

// Run runs the chain from input, the input of its first request, and returns
// the text of the final answer. An error of a log write wraps
// errChainLogWrite.
func (c *Chain) Run(ctx context.Context, input []any) (string, error) {
	request := c.request
	for request.Turn = "first"; ; request.Turn = "next" {
		request.Input = input
		response, err := c.client.Create(ctx, request)
		if err != nil {
			return "", err
		}

		var text string
		input, text, err = c.handle(ctx, input, response)
		if err != nil {
			return "", err
		}

		if len(input) == 0 {
			return text, nil
		}
	}
}

// handle logs the replies of response and runs its tool calls. It returns the
// input of the next request: input, the output items, the tool results, and
// the messages of take. After the final answer it returns no input, and the
// replies of the answer as its text.
func (c *Chain) handle(ctx context.Context, input []any, response ModelResponse) ([]any, string, error) {
	var outputs []any
	var replies []string
	for _, raw := range response.Output {
		item, err := chainInputItem(raw)
		if err != nil {
			return nil, "", err
		}

		input = append(input, item)

		var call ChainFunctionCall
		err = json.Unmarshal(raw, &call)
		if err != nil {
			return nil, "", fmt.Errorf("decode output item: %w", err)
		}

		switch call.Type {
		case "message":
			text, err := modelOutputText([]json.RawMessage{raw})
			if err != nil {
				return nil, "", err
			}

			text = strings.TrimSpace(text)
			if text == "" {
				continue
			}

			_, err = c.tools.log.Append("reply", text)
			if err != nil {
				return nil, "", fmt.Errorf("%w: %w", errChainLogWrite, err)
			}

			replies = append(replies, text)
		case "function_call":
			output, err := c.tools.Call(ctx, call.Name, call.Arguments)
			if err != nil {
				return nil, "", fmt.Errorf("%w: %w", errChainLogWrite, err)
			}

			outputs = append(outputs, ChainFunctionCallOutput{
				Type:   "function_call_output",
				CallID: call.CallID,
				Output: output,
			})
		}
	}

	if len(outputs) == 0 {
		return nil, strings.Join(replies, "\n\n"), nil
	}

	input = append(input, outputs...)
	for _, message := range c.take() {
		input = append(input, modelUserMessage(message.Text))
	}

	return input, "", nil
}

// ChainFunctionCall holds the fields of a function call output item.
type ChainFunctionCall struct {
	Type      string `json:"type"`
	Name      string `json:"name"`
	Arguments string `json:"arguments"`
	CallID    string `json:"call_id"`
}

// ChainFunctionCallOutput is the input item with the result of a function
// call.
type ChainFunctionCallOutput struct {
	Type   string `json:"type"`
	CallID string `json:"call_id"`
	Output string `json:"output"`
}

// chainInputItem returns an output item as an input item for the next request.
// The calls do not store their items, so the item goes back whole, without its
// id.
func chainInputItem(raw json.RawMessage) (map[string]any, error) {
	var item map[string]any
	err := json.Unmarshal(raw, &item)
	if err != nil {
		return nil, fmt.Errorf("decode output item: %w", err)
	}

	delete(item, "id")

	return item, nil
}