package main

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

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

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

	l.runs = &SubagentRuns{
		agent:     agent,
		memory:    memory,
		subagents: subagents,
		layout:    layout,
		client:    client,
		state:     l.state,
		report:    l.Report,
		note:      l.note,
		notify:    notify,
		running:   make(map[string]*SubagentRun),
	}

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

	return l, nil
}

// Event tells the clients what the agent does as it happens. Kind is the kind
// of a logged message (user, reply, tool, echo, note, work), or one of ready,
// busy, idle, error, and subagent, which tells that the subagent Text starts a
// run. ID is the ID of a logged message. Source names the client that sent a
// user message, such as terminal or telegram, so that each client can show the
// messages of the others, or the subagent of a report. Status tells how the run
// of a report ended: finished, stopped, or failed.
type Event struct {
	Kind   string `json:"kind"`
	ID     int    `json:"id,omitempty"`
	Text   string `json:"text,omitempty"`
	Source string `json:"source,omitempty"`
	Status string `json:"status,omitempty"`
}

// TurnLoop 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 the turn runs a Chain whose first request holds the prompt,
// the tools, and the input of turnInput.
//
// 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. User messages
// that come during a turn join the next request of its chain after the tool
// results, and so do the reports of the subagents. The subagents run until the
// loop ends.
type TurnLoop struct {
	agent  *Agent
	memory *Memory
	client ChainClient
	tools  *Toolbox
	runs   *SubagentRuns
	layout *ViewLayout
	clock  func() time.Time
	notify func(Event)
	wake   chan struct{}

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

// Submit logs a user message from the client source 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 *TurnLoop) Submit(source, text string) (Message, error) {
	message, err := l.record("user", source, 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
}

// Report logs the report text of the subagent name as a work message and tells
// the clients, with the status of the run. Unless the subagent was stopped, the
// report joins the waiting messages and wakes the loop, as a user message of
// Submit does. A failed write stops the loop, because the log refuses all
// writes after it.
func (l *TurnLoop) Report(name, text, status string) {
	message, err := l.memory.Append("work", text)
	if err == nil {
		l.notify(Event{
			Kind:   "work",
			ID:     message.ID,
			Text:   text,
			Source: name,
			Status: status,
		})
	}

	l.mutex.Lock()
	switch {
	case err != nil:
		l.failure = fmt.Errorf("%w: %w", errChainLogWrite, err)
	case status != subagentStopped:
		l.pending = append(l.pending, message)
	}
	l.mutex.Unlock()

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

// StopSubagent ends the run of the subagent name and tells if it ran. Its
// report is logged, but it starts no turn.
func (l *TurnLoop) StopSubagent(name string) (bool, error) {
	return l.runs.Stop(name)
}

// Interrupt ends the running turn. A running tool is killed, and the turn sends
// no further request.
func (l *TurnLoop) 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 *TurnLoop) 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 *TurnLoop) Append(kind, text string) (Message, error) {
	return l.record(kind, "", text)
}

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

// Run runs a turn each time user messages or reports wait, until ctx ends. It
// returns the error of a failed log write, which stops the agent. At its end
// it ends the runs of the subagents and waits for them.
func (l *TurnLoop) Run(ctx context.Context) error {
	defer l.runs.Close()

	for {
		for l.waiting() {
			err := l.turn(ctx)
			if err != nil {
				return err
			}
		}

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

		err := l.failed()
		if err != nil {
			return err
		}
	}
}

// note logs a note and tells the clients. A note starts no turn.
func (l *TurnLoop) note(text string) error {
	_, err := l.record("note", "", text)

	return err
}

// record logs a message and tells the clients, with the client source of a user
// message.
func (l *TurnLoop) record(kind, source, 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,
		Source: source,
	})

	return message, nil
}

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

	return len(l.pending) > 0
}

// failed returns the error of a failed write of a report.
func (l *TurnLoop) failed() error {
	l.mutex.Lock()
	defer l.mutex.Unlock()

	return l.failure
}

// 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 *TurnLoop) 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, errChainLogWrite) {
		l.notify(Event{
			Kind: "error",
			Text: turnError(ctx, err),
		})

		err = nil
	}

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

	return err
}

// chain runs the request chain of a turn until the final answer.
func (l *TurnLoop) 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
	}

	chain := &Chain{
		client: l.client,
		tools:  l.tools,
		request: ModelRequest{
			Model:        l.agent.config.Model,
			Effort:       l.agent.config.Effort,
			Instructions: instructions,
			Tools:        promptTools(),
		},
		take: l.take,
	}

	_, err = chain.Run(ctx, turnInput(view, l.state(), messages))

	return err
}

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

	messages := l.pending
	l.pending = nil

	return messages
}

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

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

// 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, modelUserMessage(text))
	}

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

	return input
}