package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// modelClient sends model calls.
type modelClient interface {
	Create(ctx context.Context, request Request) (Response, error)
}

// Event tells the clients what the agent does as it happens. Kind is the kind
// of a logged message (user, reply, tool, echo), or one of ready, busy, idle,
// and error. ID is the ID of a logged message.
type Event struct {
	Kind string `json:"kind"`
	ID   int    `json:"id,omitempty"`
	Text string `json:"text,omitempty"`
}

// loop runs the turns of one agent. A user message goes to the log when it
// comes, and it wakes the loop. A turn waits until every message before it is
// summarized. Then its first request holds the prompt, the tools, and the
// input of turnInput. The loop runs the tool calls of each response and sends
// their results back, until a response has no tool call: the final answer.
//
// The view goes out as the messages of a viewLayout, so the first request of
// a turn reads the view of the turn before from the prompt cache.
//
// Each later request of a turn 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. User messages that come during a turn join the next
// request after the tool results. The replies, the tool calls, and the results
// go to the log as they happen. The thoughts never do.
type loop struct {
	agent  *Agent
	memory *Memory
	client modelClient
	tools  *toolbox
	layout *viewLayout
	clock  func() time.Time
	notify func(Event)
	wake   chan struct{}

	mutex   sync.Mutex
	pending []Message
	cancel  context.CancelFunc
}

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

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

// newLoop returns the loop of agent with its memory. Its calls go through
// client, and notify gets each event. The clock gives the date of a turn.
func newLoop(agent *Agent, memory *Memory, client modelClient, clock func() time.Time, notify func(Event)) (*loop, error) {
	layout, err := loadLayout(filepath.Join(agent.dir, layoutFileName))
	if err != nil {
		return nil, err
	}

	l := &loop{
		agent:  agent,
		memory: memory,
		client: client,
		layout: layout,
		clock:  clock,
		notify: notify,
		wake:   make(chan struct{}, 1),
	}

	l.tools = &toolbox{
		log:     l,
		chat:    memory,
		workdir: agent.config.Workdir,
	}

	return l, nil
}

// Submit logs a user message and wakes the loop. The message joins the next
// request: the first one of a new turn, or the next one of a running turn.
func (l *loop) Submit(text string) (Message, error) {
	message, err := l.Append("user", text)
	if err != nil {
		return Message{}, err
	}

	l.mutex.Lock()
	l.pending = append(l.pending, message)
	l.mutex.Unlock()

	select {
	case l.wake <- struct{}{}:
	default:
	}

	return message, nil
}

// Interrupt ends the running turn. A running tool is killed, and the turn sends
// no further request.
func (l *loop) Interrupt() bool {
	l.mutex.Lock()
	defer l.mutex.Unlock()

	if l.cancel == nil {
		return false
	}

	l.cancel()

	return true
}

// Busy tells if a turn runs.
func (l *loop) Busy() bool {
	l.mutex.Lock()
	defer l.mutex.Unlock()

	return l.cancel != nil
}

// Append logs a message and tells the clients. The toolbox logs its calls and
// results through it.
func (l *loop) Append(kind, text string) (Message, error) {
	message, err := l.memory.Append(kind, text)
	if err != nil {
		return Message{}, err
	}

	l.notify(Event{
		Kind: kind,
		ID:   message.ID,
		Text: text,
	})

	return message, nil
}

// Flush makes the log durable.
func (l *loop) Flush(ctx context.Context) error {
	return l.memory.Flush(ctx)
}

// Run runs a turn each time user messages wait, until ctx ends. It returns the
// error of a failed log write, which stops the agent.
func (l *loop) Run(ctx context.Context) error {
	for {
		for l.waiting() {
			err := l.turn(ctx)
			if err != nil {
				return err
			}
		}

		select {
		case <-l.wake:
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

// waiting tells if user messages wait for a request.
func (l *loop) waiting() bool {
	l.mutex.Lock()
	defer l.mutex.Unlock()

	return len(l.pending) > 0
}

// take returns the user messages that wait for a request, and removes them.
func (l *loop) take() []Message {
	l.mutex.Lock()
	defer l.mutex.Unlock()

	messages := l.pending
	l.pending = nil

	return messages
}

// turn runs one turn. A failed or interrupted model call ends the turn and
// tells the clients. The messages of the turn stay in the log and the view, and
// the next turn starts only with a new message, so a lasting failure such as a
// needed login does not repeat. Only a failed log write is an error of turn.
func (l *loop) turn(ctx context.Context) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	l.mutex.Lock()
	l.cancel = cancel
	l.mutex.Unlock()

	defer func() {
		l.mutex.Lock()
		l.cancel = nil
		l.mutex.Unlock()
	}()

	l.notify(Event{
		Kind: "busy",
	})

	err := l.chain(ctx)
	if err != nil && !errors.Is(err, errLogWrite) {
		l.notify(Event{
			Kind: "error",
			Text: turnError(ctx, err),
		})

		err = nil
	}

	l.notify(Event{
		Kind: "idle",
	})

	return err
}

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

// chain runs the request chain of a turn until the final answer.
func (l *loop) chain(ctx context.Context) error {
	messages := l.take()
	if len(messages) == 0 {
		return nil
	}

	err := l.memory.Wait(ctx)
	if err != nil {
		return err
	}

	instructions, err := l.agent.Prompt()
	if err != nil {
		return err
	}

	skip := make(map[int]bool, len(messages))
	for _, message := range messages {
		skip[message.ID] = true
	}

	view, err := l.layout.Split(l.memory.ViewLines(skip))
	if err != nil {
		return err
	}

	input := turnInput(view, l.state(), messages)
	for turn := "first"; ; turn = "next" {
		response, err := l.client.Create(ctx, Request{
			Model:        l.agent.config.Model,
			Effort:       l.agent.config.Effort,
			Instructions: instructions,
			Tools:        agentTools(),
			Input:        input,
			Turn:         turn,
		})
		if err != nil {
			return err
		}

		input, err = l.handle(ctx, input, response)
		if err != nil {
			return err
		}

		if len(input) == 0 {
			return 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 user messages that came in the meantime. After the final answer it
// returns no input.
func (l *loop) handle(ctx context.Context, input []any, response Response) ([]any, error) {
	var outputs []any
	for _, raw := range response.Output {
		item, err := inputItem(raw)
		if err != nil {
			return nil, err
		}

		input = append(input, item)

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

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

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

			_, err = l.Append("reply", text)
			if err != nil {
				return nil, fmt.Errorf("%w: %w", errLogWrite, err)
			}
		case "function_call":
			output, err := l.tools.Call(ctx, call.Name, call.Arguments)
			if err != nil {
				return nil, fmt.Errorf("%w: %w", errLogWrite, err)
			}

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

	if len(outputs) == 0 {
		return nil, nil
	}

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

	return input, nil
}

// state returns the per-turn state: the end of the view, and the date and time
// of the turn.
func (l *loop) state() string {
	return "</chat>\n\nNow: " + l.clock().Format("Monday 2006-01-02 15:04 -07:00")
}

// turnInput returns the input of the first request of a turn, stable parts
// first: a user message for each view message, one with the per-turn state,
// and one for each new message, whole. The view leaves out the leaf lines of
// the new messages.
func turnInput(view []string, state string, messages []Message) []any {
	input := make([]any, 0, len(view)+1+len(messages))
	for _, text := range view {
		input = append(input, userMessage(text))
	}

	input = append(input, userMessage(state))
	for _, message := range messages {
		input = append(input, userMessage(message.Text))
	}

	return input
}

// inputItem 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 inputItem(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
}

// turnError returns the text of the error that ended a turn.
func turnError(ctx context.Context, err error) string {
	if ctx.Err() != nil {
		return "interrupted"
	}

	return err.Error()
}