// Package bridge relays events between the ACP agent session and the
// browsers connected over WebSocket.
//
// It is the meeting point of two asynchronous worlds: the agent streams
// updates and blocks on permission requests; browsers connect, disconnect and
// answer at human speed. The bridge broadcasts agent events to every client,
// replays the session history to late joiners, and routes each permission
// decision back to the exact request waiting for it.
package bridge

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"path/filepath"
	"strconv"
	"sync"

	acp "github.com/coder/acp-go-sdk"
)

// Prompter is the slice of the agent session the bridge drives. It is
// satisfied by *agent.Session and by test fakes.
type Prompter interface {
	ID() string
	Prompt(ctx context.Context, blocks []acp.ContentBlock) (acp.StopReason, error)
	Cancel(ctx context.Context) error
}

// maxHistory bounds the number of events replayed to late-joining clients;
// beyond it the oldest events are dropped (the live broadcast is unaffected).
const maxHistory = 4096

// clientBuffer is each subscriber's outgoing queue; a client too slow to
// drain it is disconnected rather than allowed to stall the whole bridge.
const clientBuffer = 256

type permissionDecision struct {
	optionId  string
	cancelled bool
}

// Bridge broadcasts agent events to WebSocket clients and routes their
// responses back to the agent. The zero value is not usable; call New.
//
// Example:
//
//	b := bridge.New(slog.Default())
//	session, _ := agent.Start(ctx, agent.Options{Handler: b, ...})
//	b.SetSession(session)
//	http.Handle("GET /ws", b.WebSocketHandler())
type Bridge struct {
	logger *slog.Logger

	mu           sync.Mutex
	session      Prompter
	root         string
	subscribers  map[chan []byte]struct{}
	history      [][]byte
	turnActive   bool
	nextRequest  int
	pendingPerms map[string]chan permissionDecision
	// pendingRequests keeps the encoded permission_request events so they
	// can be replayed to clients that connect while a request is open.
	pendingRequests map[string][]byte
}

// New creates an empty bridge; attach the agent session with SetSession.
func New(logger *slog.Logger) *Bridge {
	if logger == nil {
		logger = slog.Default()
	}
	return &Bridge{
		logger:          logger,
		subscribers:     make(map[chan []byte]struct{}),
		pendingPerms:    make(map[string]chan permissionDecision),
		pendingRequests: make(map[string][]byte),
	}
}

// SetSession attaches the agent session the bridge drives. It must be called
// before the first client message arrives.
func (b *Bridge) SetSession(session Prompter) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.session = session
}

// SetWorkspaceRoot sets the directory relative attachment paths resolve
// against (the agent's cwd). Without it, relative paths are sent as given.
func (b *Bridge) SetWorkspaceRoot(root string) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.root = root
}

// Subscribe registers a new client. It returns the client's event channel,
// the events to replay so the client catches up (history, then any pending
// permission requests), and an unsubscribe function.
func (b *Bridge) Subscribe() (events <-chan []byte, replay [][]byte, unsubscribe func()) {
	ch := make(chan []byte, clientBuffer)
	b.mu.Lock()
	defer b.mu.Unlock()

	b.subscribers[ch] = struct{}{}

	sessionId := ""
	if b.session != nil {
		sessionId = b.session.ID()
	}
	replay = append(replay, encode(Outgoing{Type: OutgoingHello, SessionId: sessionId, TurnActive: b.turnActive}))
	replay = append(replay, b.history...)
	for _, raw := range b.pendingRequests {
		replay = append(replay, raw)
	}

	return ch, replay, func() {
		b.mu.Lock()
		defer b.mu.Unlock()
		if _, ok := b.subscribers[ch]; ok {
			delete(b.subscribers, ch)
			close(ch)
		}
	}
}

// HandleIncoming processes one message from a client. Errors are reported to
// the clients as OutgoingError events, never returned: a malformed message
// from one browser must not tear down the connection handling.
func (b *Bridge) HandleIncoming(ctx context.Context, raw []byte) {
	var msg Incoming
	if err := json.Unmarshal(raw, &msg); err != nil {
		b.broadcast(encode(Outgoing{Type: OutgoingError, Message: "malformed message: " + err.Error()}), false)
		return
	}

	switch msg.Type {
	case IncomingPrompt:
		b.startTurn(msg.Text, msg.Attachments)
	case IncomingCancel:
		b.cancelTurn(ctx)
	case IncomingPermissionResponse:
		b.resolvePermission(msg)
	default:
		b.broadcast(encode(Outgoing{Type: OutgoingError, Message: fmt.Sprintf("unknown message type %q", msg.Type)}), false)
	}
}

// startTurn launches a prompt turn in the background; updates stream through
// HandleSessionUpdate while it runs.
func (b *Bridge) startTurn(text string, attachments []Attachment) {
	b.mu.Lock()
	session := b.session
	root := b.root
	if session == nil {
		b.mu.Unlock()
		b.broadcast(encode(Outgoing{Type: OutgoingError, Message: "no agent session"}), false)
		return
	}
	if b.turnActive {
		b.mu.Unlock()
		b.broadcast(encode(Outgoing{Type: OutgoingError, Message: "a turn is already running"}), false)
		return
	}
	b.turnActive = true
	b.mu.Unlock()

	b.broadcast(encode(Outgoing{Type: OutgoingUserMessage, Text: text, Attachments: attachments}), true)
	b.broadcast(encode(Outgoing{Type: OutgoingTurnStarted}), true)

	blocks := promptBlocks(text, attachments, root)
	go func() {
		// The turn outlives the WebSocket message that started it, so it
		// runs under its own context, ended by session/cancel only.
		stop, err := session.Prompt(context.Background(), blocks)

		b.mu.Lock()
		b.turnActive = false
		b.mu.Unlock()

		if err != nil {
			b.logger.Error("prompt failed", "error", err)
			b.broadcast(encode(Outgoing{Type: OutgoingError, Message: "prompt failed: " + err.Error()}), true)
			b.broadcast(encode(Outgoing{Type: OutgoingTurnEnded, StopReason: string(acp.StopReasonRefusal)}), true)
			return
		}
		b.broadcast(encode(Outgoing{Type: OutgoingTurnEnded, StopReason: string(stop)}), true)
	}()
}

// promptBlocks builds the ACP prompt: the text first, then one resource_link
// per attachment (file:// URI, named after the mention), which every ACP
// agent must accept.
//
// Example:
//
//	promptBlocks("Explain @src/main.go", []Attachment{{Path: "src/main.go"}}, "/work")
//	// [text "Explain @src/main.go", resource_link name=src/main.go uri=file:///work/src/main.go]
func promptBlocks(text string, attachments []Attachment, root string) []acp.ContentBlock {
	blocks := []acp.ContentBlock{acp.TextBlock(text)}
	for _, attachment := range attachments {
		if attachment.Path == "" {
			continue
		}
		path := attachment.Path
		if !filepath.IsAbs(path) && root != "" {
			path = filepath.Join(root, path)
		}
		name := attachment.Name
		if name == "" {
			name = attachment.Path
		}
		blocks = append(blocks, acp.ResourceLinkBlock(name, "file://"+filepath.ToSlash(path)))
	}
	return blocks
}

func (b *Bridge) cancelTurn(ctx context.Context) {
	b.mu.Lock()
	session := b.session
	b.mu.Unlock()
	if session == nil {
		return
	}
	if err := session.Cancel(ctx); err != nil {
		b.logger.Error("cancel failed", "error", err)
		b.broadcast(encode(Outgoing{Type: OutgoingError, Message: "cancel failed: " + err.Error()}), false)
	}
}

func (b *Bridge) resolvePermission(msg Incoming) {
	b.mu.Lock()
	waiting, ok := b.pendingPerms[msg.RequestId]
	if ok {
		delete(b.pendingPerms, msg.RequestId)
		delete(b.pendingRequests, msg.RequestId)
	}
	b.mu.Unlock()
	if !ok {
		// Either already answered by another client or unknown: ignore.
		return
	}
	waiting <- permissionDecision{optionId: msg.OptionId, cancelled: msg.Cancelled}
	b.broadcast(encode(Outgoing{Type: OutgoingPermissionResolved, RequestId: msg.RequestId}), true)
}

// HandleSessionUpdate implements agent.Handler: every ACP update is relayed
// verbatim to the clients and recorded for replay.
func (b *Bridge) HandleSessionUpdate(_ context.Context, notification acp.SessionNotification) {
	update, err := json.Marshal(notification.Update)
	if err != nil {
		b.logger.Error("marshal session update", "error", err)
		return
	}
	b.broadcast(encode(Outgoing{Type: OutgoingSessionUpdate, Update: update}), true)
}

// HandlePermissionRequest implements agent.Handler: the request is broadcast
// to the clients and the call blocks until one of them answers or the agent
// cancels the turn.
func (b *Bridge) HandlePermissionRequest(ctx context.Context, request acp.RequestPermissionRequest) (acp.RequestPermissionResponse, error) {
	rawRequest, err := json.Marshal(request)
	if err != nil {
		return acp.RequestPermissionResponse{}, fmt.Errorf("marshal permission request: %w", err)
	}

	decisionCh := make(chan permissionDecision, 1)
	b.mu.Lock()
	b.nextRequest++
	requestId := "perm-" + strconv.Itoa(b.nextRequest)
	event := encode(Outgoing{Type: OutgoingPermissionRequest, RequestId: requestId, Request: rawRequest})
	b.pendingPerms[requestId] = decisionCh
	b.pendingRequests[requestId] = event
	b.mu.Unlock()

	b.broadcast(event, false)

	select {
	case decision := <-decisionCh:
		if decision.cancelled {
			return acp.RequestPermissionResponse{Outcome: acp.NewRequestPermissionOutcomeCancelled()}, nil
		}
		return acp.RequestPermissionResponse{Outcome: acp.NewRequestPermissionOutcomeSelected(acp.PermissionOptionId(decision.optionId))}, nil
	case <-ctx.Done():
		// The agent cancelled the turn (or disconnected) while waiting.
		b.mu.Lock()
		delete(b.pendingPerms, requestId)
		delete(b.pendingRequests, requestId)
		b.mu.Unlock()
		b.broadcast(encode(Outgoing{Type: OutgoingPermissionResolved, RequestId: requestId}), true)
		return acp.RequestPermissionResponse{Outcome: acp.NewRequestPermissionOutcomeCancelled()}, nil
	}
}

// broadcast sends an event to every connected client, disconnecting any
// client whose buffer is full. With record=true the event also joins the
// replay history.
func (b *Bridge) broadcast(event []byte, record bool) {
	b.mu.Lock()
	defer b.mu.Unlock()

	if record {
		b.history = append(b.history, event)
		if len(b.history) > maxHistory {
			b.history = b.history[len(b.history)-maxHistory:]
		}
	}

	for ch := range b.subscribers {
		select {
		case ch <- event:
		default:
			// The client stopped draining; drop it instead of blocking
			// the agent's event stream.
			delete(b.subscribers, ch)
			close(ch)
			b.logger.Warn("dropped a slow websocket client")
		}
	}
}