package acp

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"sync"
)

const maxUnsent = 4 * 1024 * 1024

// Limit concatenation and JSON encoding work while the queue lock is held.
const maxCoalescedBytes = 32 * 1024

// The writer owns each textChunk and coalesces only assistant output.
type textChunk struct {
	SessionID string
	Kind      string
	MessageID string
	Text      string
}

func (chunk textChunk) message() map[string]any {
	return map[string]any{
		"jsonrpc": "2.0", "method": "session/update",
		"params": map[string]any{
			"sessionId": chunk.SessionID,
			"update": map[string]any{
				"sessionUpdate": chunk.Kind,
				"messageId":     chunk.MessageID,
				"content":       map[string]string{"type": "text", "text": chunk.Text},
			},
		},
	}
}

func (chunk textChunk) sameMessage(other textChunk) bool {
	return chunk.SessionID == other.SessionID && chunk.Kind == other.Kind && chunk.MessageID == other.MessageID
}

type queuedMessage struct {
	chunk   *textChunk
	barrier chan struct{}
	data    []byte
}

// writer accounts for the in-flight message until each byte has been written.
type writer struct {
	wake   chan struct{}
	done   chan struct{}
	err    error
	output io.Writer
	queue  []queuedMessage
	mu     sync.Mutex
	unsent int
}

func newWriter(output io.Writer) *writer {
	return &writer{output: output, wake: make(chan struct{}, 1), done: make(chan struct{})}
}
func (w *writer) enqueue(value any, barrier bool) (<-chan struct{}, error) {
	return w.enqueueMessage(value, barrier, nil)
}

func (w *writer) enqueueText(chunk textChunk) error {
	if chunk.Kind != "agent_message_chunk" && chunk.Kind != "agent_thought_chunk" {
		_, err := w.enqueue(chunk.message(), false)
		return err
	}
	_, err := w.enqueueMessage(chunk.message(), false, &chunk)
	return err
}

func (w *writer) enqueueMessage(value any, barrier bool, chunk *textChunk) (<-chan struct{}, error) {
	data, err := json.Marshal(value)
	if err != nil {
		return nil, err
	}
	data = append(data, '\n')
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.err != nil {
		return nil, w.err
	}

	index := len(w.queue)
	addedBytes := len(data)
	if chunk != nil && index > 0 {
		last := &w.queue[index-1]
		if last.chunk != nil && last.chunk.sameMessage(*chunk) && last.barrier == nil && len(last.data)+len(data) <= maxCoalescedBytes {
			combined := *last.chunk
			combined.Text += chunk.Text
			// A text envelope contains only strings and cannot fail encoding.
			data, _ = json.Marshal(combined.message())
			data = append(data, '\n')
			chunk = &combined
			addedBytes = len(data) - len(last.data)
			index--
		}
	}
	if w.unsent+addedBytes > maxUnsent {
		w.err = errors.New("ACP stdout queue exceeds 4 MiB")
		select {
		case w.wake <- struct{}{}:
		default:
		}
		return nil, w.err
	}
	item := queuedMessage{data: data, chunk: chunk}
	if barrier {
		item.barrier = make(chan struct{})
	}
	if index == len(w.queue) {
		w.queue = append(w.queue, item)
	} else {
		w.queue[index] = item
	}
	w.unsent += addedBytes
	select {
	case w.wake <- struct{}{}:
	default:
	}
	return item.barrier, nil
}
func (w *writer) failure() error { w.mu.Lock(); defer w.mu.Unlock(); return w.err }
func (w *writer) run(ctx context.Context) {
	defer close(w.done)
	for {
		w.mu.Lock()
		if w.err != nil {
			w.mu.Unlock()
			return
		}
		if len(w.queue) == 0 {
			w.mu.Unlock()
			select {
			case <-ctx.Done():
				return
			case <-w.wake:
				continue
			}
		}
		item := w.queue[0]
		w.queue[0] = queuedMessage{}
		w.queue = w.queue[1:]
		w.mu.Unlock()
		for len(item.data) > 0 {
			count, err := w.output.Write(item.data)
			if count < 0 || count > len(item.data) {
				count = 0
				err = io.ErrShortWrite
			}
			w.mu.Lock()
			w.unsent -= count
			if err == nil && count == 0 {
				err = io.ErrShortWrite
			}
			if err != nil {
				w.err = err
			}
			w.mu.Unlock()
			if err != nil {
				return
			}
			item.data = item.data[count:]
		}
		if item.barrier != nil {
			close(item.barrier)
		}
	}
}

func (w *writer) flush() <-chan struct{} {
	barrier := make(chan struct{})
	w.mu.Lock()
	defer w.mu.Unlock()
	w.queue = append(w.queue, queuedMessage{barrier: barrier})
	select {
	case w.wake <- struct{}{}:
	default:
	}
	return barrier
}