package main

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

// openDaemonMessages connects the root and its children. It restores queued
// messages before the loops start. A queue entry stays until its recipient
// finishes the turn that takes it, so a stop can interrupt delivery or a turn.
// A routed batch with a removed child goes to the root.
func openDaemonMessages(loops []*TurnLoop) (*DaemonMessages, error) {
	routes, err := openRoutes(filepath.Join(loops[0].agent.dir, "routes.json"), loops[0].logger)
	if err != nil {
		return nil, fmt.Errorf("open agent routes: %w", err)
	}

	m := &DaemonMessages{
		routes:  routes,
		path:    filepath.Join(loops[0].agent.dir, "messages.json"),
		ordered: append([]*TurnLoop(nil), loops...),
		loops:   make(map[string]*TurnLoop),
	}
	for _, loop := range loops {
		m.loops[loop.agent.config.Name] = loop
		loop.messages = m
		loop.tools.messages = m
		if loop.agent.config.Parent == "" {
			loop.logger = loop.logger.With("name", loop.agent.config.Name)
			loop.tools.routes = routes
			if loop.tools.whatsapp != nil {
				loop.tools.whatsapp.mutex.Lock()

				loop.tools.whatsapp.routes = routes
				loop.tools.whatsapp.mutex.Unlock()
			}
		}
	}

	for _, child := range loops[1:] {
		err := loops[0].runs.subagents.Reserve(child.agent.config.Name)
		if err != nil {
			return nil, fmt.Errorf("reserve child name: %w", err)
		}
	}

	data, err := os.ReadFile(m.path)
	if err != nil && !errors.Is(err, os.ErrNotExist) {
		return nil, fmt.Errorf("read agent messages: %w", err)
	}

	if err == nil {
		err = json.Unmarshal(data, &m.state)
		if err != nil {
			return nil, fmt.Errorf("decode agent messages: %w", err)
		}
	}

	if m.state.Next < 0 {
		return nil, errors.New("invalid agent message sequence")
	}

	for i, number := range m.state.WhatsAppBatches {
		if number <= 0 || slices.Contains(m.state.WhatsAppBatches[:i], number) {
			return nil, errors.New("invalid WhatsApp batch receipt")
		}
	}

	previous := 0
	for i := range m.state.Pending {
		message := &m.state.Pending[i]
		source := m.loops[message.Source]
		target := m.loops[message.Target]
		if message.WhatsAppBatch > 0 && source == loops[0] && (target == nil || (target != source && target.agent.config.Parent != message.Source)) {
			message.Target = message.Source
			message.LogID = nil
			target = source
		}

		if source == nil || target == nil || message.ID <= previous || message.ID > m.state.Next || (message.LogID != nil && *message.LogID < 0) {
			return nil, errors.New("invalid queued agent message")
		}

		if message.WhatsAppBatch != 0 {
			if message.WhatsAppBatch < 0 || message.Kind != "work" || source != loops[0] || (target != source && target.agent.config.Parent != message.Source) {
				return nil, errors.New("invalid queued WhatsApp batch")
			}
		} else if (message.Kind != "user" || target.agent.config.Parent != message.Source) && (message.Kind != "work" || source.agent.config.Parent != message.Target) {
			return nil, errors.New("queued agents are not parent and child")
		}

		previous = message.ID
	}

	for i := range m.state.Pending {
		err = m.deliver(&m.state.Pending[i])
		if err != nil {
			return nil, fmt.Errorf("restore agent message: %w", err)
		}
	}

	return m, nil
}

// DaemonMessages owns the durable queue between the root and its children and
// the routed inputs. The mutex serializes queue writes and delivery. Backups
// take it before the locks of the loops, memory, and subagents. A stored log ID
// makes delivery idempotent, including a stop between the log append and the
// wake.
type DaemonMessages struct {
	mutex   sync.Mutex
	path    string
	loops   map[string]*TurnLoop
	ordered []*TurnLoop
	state   DaemonMessageFile
	failure error
	routes  *Routes
}

// deliver records one queued message and wakes its recipient. The caller holds
// the mutex, or owns the queue before the loops start.
func (m *DaemonMessages) deliver(message *DaemonMessage) error {
	loop := m.loops[message.Target]
	logged, err := loop.memory.AppendOnce(message.Kind, message.Text, message.LogID, func(id int) error {
		message.LogID = &id
		return m.save()
	})
	if err != nil {
		return fmt.Errorf("record agent message: %w", err)
	}

	status := ""
	source := message.Source
	if message.WhatsAppBatch > 0 {
		source = whatsappSource
	} else if message.Kind == "work" {
		status = subagentFinished
	}

	loop.notify(Event{
		Kind:   message.Kind,
		ID:     logged.ID,
		Text:   message.Text,
		Source: source,
		Status: status,
	})
	loop.Wake(logged)

	if message.WhatsAppBatch > 0 {
		loop.logger.Info("batch routed", "kind", whatsappSource, "id", message.WhatsAppBatch, "to", message.Target)
	} else {
		loop.logger.Info("message received", "from", message.Source, "id", message.ID)
	}

	return nil
}

// save replaces the queue file under the mutex, or before the loops start.
func (m *DaemonMessages) save() error {
	data, err := json.Marshal(m.state)
	if err != nil {
		return fmt.Errorf("encode agent messages: %w", err)
	}

	err = writeFileAtomic(m.path, append(data, '\n'))
	if err != nil {
		return fmt.Errorf("save agent messages: %w", err)
	}

	return nil
}

// AddChild holds the queue lock while open makes a child's files and loop. The
// loop joins backups before its first task can be delivered.
func (m *DaemonMessages) AddChild(parent, name string, open func() (*TurnLoop, error)) (*TurnLoop, error) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	if m.loops[name] != nil {
		return nil, fmt.Errorf("name %q is already used", name)
	}

	root := m.loops[parent]
	if root == nil || root.agent.config.Parent != "" {
		return nil, errors.New("only the root can create children")
	}

	if m.failure != nil {
		return nil, fmt.Errorf("agent message queue stopped: %w", m.failure)
	}

	loop, err := open()
	if err != nil {
		return nil, fmt.Errorf("open child: %w", err)
	}

	loop.messages = m
	loop.tools.messages = m
	m.loops[name] = loop
	m.ordered = append(m.ordered, loop)

	return loop, nil
}

// Child returns only a direct child of parent. Loops live until the daemon
// stops, so the returned loop stays open for the caller's turn.
func (m *DaemonMessages) Child(parent, name string) *TurnLoop {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	loop := m.loops[name]
	if loop == nil || loop.agent.config.Parent != parent {
		return nil
	}

	return loop
}

// ChildState lists the names and responsibilities of a parent's children in
// name order. It returns no heading when the parent has no children.
func (m *DaemonMessages) ChildState(parent string) string {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	var lines []string

	for name, loop := range m.loops {
		if loop.agent.config.Parent == parent {
			lines = append(lines, name+": "+loop.agent.config.Responsibility)
		}
	}

	if len(lines) == 0 {
		return ""
	}

	slices.Sort(lines)

	return "\n\nChildren:\n" + strings.Join(lines, "\n")
}

// SnapshotAgents locks the queue and all registered loops, including children
// made during this daemon run, before snapshot copies their files.
func (m *DaemonMessages) SnapshotAgents(ctx context.Context, snapshot func() error) error {
	return m.Snapshot(func() error {
		return daemonSnapshot(ctx, m.ordered, snapshot)
	})
}

// Tell queues a user message from parent to its own child, then wakes the
// child. It is the daemon entry point for parent-to-child messages.
func (m *DaemonMessages) Tell(parent, child, text string) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	source := m.loops[parent]
	target := m.loops[child]
	if source == nil || target == nil || target.agent.config.Parent != parent {
		return errors.New("the recipient is not a child of the sender")
	}

	if m.failure != nil {
		return fmt.Errorf("agent message queue stopped: %w", m.failure)
	}

	m.add(parent, child, "user", text)

	err := m.save()
	if err == nil {
		err = m.deliver(&m.state.Pending[len(m.state.Pending)-1])
	}

	if err != nil {
		return fmt.Errorf("send agent message: %w", m.failed(err))
	}

	return nil
}

// RouteWhatsApp queues a routed batch as work and wakes its child. It returns
// false for an unrouted source or a missing child. A saved receipt survives
// turn acknowledgement until the inbox removes the batch. A repeat after a stop
// between queueing and inbox removal cannot append a second entry.
func (m *DaemonMessages) RouteWhatsApp(batch WhatsAppBatch, text string) (bool, error) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	if m.failure != nil {
		return false, fmt.Errorf("agent message queue stopped: %w", m.failure)
	}

	if batch.Number <= 0 {
		return false, errors.New("invalid WhatsApp batch number")
	}

	if slices.Contains(m.state.WhatsAppBatches, batch.Number) {
		return true, nil
	}

	root := m.ordered[0].agent.config.Name
	name := m.routes.Target(whatsappSource + ":" + batch.Chat.ID)
	child := m.loops[name]
	if child == nil || child.agent.config.Parent != root {
		return false, nil
	}

	m.add(root, name, "work", text)

	message := &m.state.Pending[len(m.state.Pending)-1]
	message.WhatsAppBatch = batch.Number
	m.state.WhatsAppBatches = append(m.state.WhatsAppBatches, batch.Number)
	err := m.save()
	if err == nil {
		err = m.deliver(message)
	}

	if err != nil {
		return false, fmt.Errorf("route WhatsApp batch: %w", m.failed(err))
	}

	return true, nil
}

// PruneWhatsApp drops receipts only for batches absent from the durable inbox.
// Queued turns keep their own log IDs until they finish. Receipts left by a
// stop after inbox removal are removed by the next inbox delivery.
func (m *DaemonMessages) PruneWhatsApp(batches []WhatsAppBatch) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	if m.failure != nil {
		return fmt.Errorf("agent message queue stopped: %w", m.failure)
	}

	held := make(map[int]bool, len(batches))
	for _, batch := range batches {
		held[batch.Number] = true
	}

	before := len(m.state.WhatsAppBatches)
	m.state.WhatsAppBatches = slices.DeleteFunc(m.state.WhatsAppBatches, func(number int) bool {
		return !held[number]
	})
	if len(m.state.WhatsAppBatches) == before {
		return nil
	}

	err := m.save()
	if err != nil {
		return fmt.Errorf("prune WhatsApp receipts: %w", m.failed(err))
	}

	return nil
}

// add appends a queue entry under the mutex. Its ID is never reused.
func (m *DaemonMessages) add(source, target, kind, text string) {
	m.state.Next++
	m.state.Pending = append(m.state.Pending, DaemonMessage{
		ID:     m.state.Next,
		Source: source,
		Target: target,
		Kind:   kind,
		Text:   text,
	})
}

// failed stops all loops after a queue failure. Further sends must not change
// the queue, because the outcome of the failed write is not known.
func (m *DaemonMessages) failed(err error) error {
	m.failure = err
	for _, loop := range m.loops {
		loop.fail(err)
	}

	return err
}

// Snapshot holds the queue lock while snapshot locks and copies agent files.
// Its callback must take the memory locks after this lock, as delivery does.
func (m *DaemonMessages) Snapshot(snapshot func() error) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	err := snapshot()
	if err != nil {
		return fmt.Errorf("copy agent message snapshot: %w", err)
	}

	return nil
}

// Finish acknowledges the queued messages taken by a turn and queues its final
// report in the same file replacement. A child with no final text sends
// nothing. Call after a final reply, a failure, or a user interrupt, but not
// when a daemon stop interrupts a model request.
func (m *DaemonMessages) Finish(loop *TurnLoop, active []Message, reply string) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	if m.failure != nil {
		return fmt.Errorf("agent message queue stopped: %w", m.failure)
	}

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

	name := loop.agent.config.Name
	before := len(m.state.Pending)
	m.state.Pending = slices.DeleteFunc(m.state.Pending, func(message DaemonMessage) bool {
		return message.Target == name && message.LogID != nil && ids[*message.LogID]
	})

	parent := loop.agent.config.Parent
	report := parent != "" && strings.TrimSpace(reply) != ""
	if report {
		m.add(name, parent, "work", "["+name+"] "+reply)
	}

	if !report && len(m.state.Pending) == before {
		return nil
	}

	err := m.save()
	if err == nil && report {
		err = m.deliver(&m.state.Pending[len(m.state.Pending)-1])
	}

	if err != nil {
		return fmt.Errorf("finish agent messages: %w", m.failed(err))
	}

	return nil
}

// DaemonMessageFile holds the sequence, unacknowledged deliveries, and receipts
// for WhatsApp batches until the inbox removes them. A receipt stays when its
// turn finishes before inbox removal.
type DaemonMessageFile struct {
	Next            int             `json:"next"`
	Pending         []DaemonMessage `json:"pending"`
	WhatsAppBatches []int           `json:"whatsapp_batches,omitempty"`
}

// DaemonMessage is one queued delivery. LogID is set before the log append. A
// nil LogID means that delivery has not started. WhatsAppBatch marks a routed
// external batch, whose event source is WhatsApp rather than its queue sender.
type DaemonMessage struct {
	ID            int    `json:"id"`
	Source        string `json:"source"`
	Target        string `json:"target"`
	Kind          string `json:"kind"`
	Text          string `json:"text"`
	LogID         *int   `json:"log_id,omitempty"`
	WhatsAppBatch int    `json:"whatsapp_batch,omitempty"`
}