// Package acp is bob's second front end: the same engine, the same tools and
// the same history rules as the terminal REPL (internal/agent), exposed over
// the Agent Client Protocol — JSON-RPC 2.0 on stdio, one message per line —
// so that an editor (Zed, IntelliJ, Neovim…) can host the agent.
//
// The division of labour:
//
//   - the SDK (github.com/coder/acp-go-sdk) does the transport: framing,
//     routing initialize/session/new/session/prompt to the methods below,
//     correlating the ids of our own requests to the client (permissions);
//   - this package translates: a prompt turn becomes one engine.Generate call,
//     and the events the tools emit through ui.Sink become session/update
//     notifications — text chunks, tool_call lifecycles, diffs;
//   - nothing else changes: internal/engine and internal/tools are shared with
//     the terminal front end, which is the point of the exercise.
//
// stdout carries JSON-RPC only. Everything human — the banner, the [acp] trail,
// the SDK's own diagnostics — goes to stderr, which the spec leaves to logging
// and which Zed shows in `dev: open acp logs`.
package acp

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"
	"strings"
	"sync"

	"mm/internal/compact"
	"mm/internal/config"
	"mm/internal/engine"
	"mm/internal/mention"
	history "mm/internal/session"
	"mm/internal/ui"

	sdk "github.com/coder/acp-go-sdk"
	"github.com/firebase/genkit/go/ai"
)

// version is what initialize reports in agentInfo — the part number, so a
// trace says which chapter of the talk produced it.
const version = "0.11.0"

// session is one conversation with the client. Like the REPL, the history is
// plain []*ai.Message with the system prompt first; unlike the REPL there can
// be several of them, one per session/new.
type session struct {
	cwd      string
	messages []*ai.Message
	cancel   context.CancelFunc // cancels the turn in flight, if any
	// allowAlways remembers the tools the user granted for the whole session
	// ("allow always" in the permission dialog).
	allowAlways map[string]bool
}

type front struct {
	e      *engine.Engine
	system string
	tools  []ai.ToolRef
	conn   *sdk.AgentSideConnection
	out    *responseWriter // what the SDK writes to; see responseWriter

	mu       sync.Mutex // guards sessions and nextID
	sessions map[sdk.SessionId]*session
	nextID   int

	// turnMu serialises the turns. Two reasons, both structural: the active
	// ui.Sink is process-global, and engine.Generate was never designed for two
	// concurrent turns. The spec leaves this to the agent, and Zed does not
	// send a second prompt while one is running — a queued turn just waits.
	turnMu sync.Mutex
}

var _ sdk.Agent = (*front)(nil)

// Run wires the agent side of an ACP connection on stdio and blocks until the
// client closes it. Same arguments as agent.Run: main.go's wiring — engine,
// system prompt, tool list — is identical for the two front ends.
func Run(ctx context.Context, e *engine.Engine, system string, tools []ai.ToolRef) {
	f := newFront(e, system, tools, os.Stdout, os.Stdin)
	select {
	case <-f.conn.Done(): // the editor closed our stdin: we are done
	case <-ctx.Done():
	}
}

// newFront wires a front end on an arbitrary pair of streams — stdio in Run,
// pipes in the tests. w is where WE write (JSON-RPC to the client), r where
// we read the client's messages.
func newFront(e *engine.Engine, system string, tools []ai.ToolRef, w io.Writer, r io.Reader) *front {
	f := &front{e: e, system: system, tools: tools, sessions: map[sdk.SessionId]*session{}}
	f.out = &responseWriter{Writer: w, hooks: map[sdk.SessionId]func(){}}
	conn := sdk.NewAgentSideConnection(f, f.out, r)
	conn.SetLogger(slog.New(slog.NewTextHandler(os.Stderr, nil)))
	f.conn = conn
	return f
}

// responseWriter is the io.Writer handed to the SDK: the real output plus one
// hook per session, run once the session/new response carrying that session
// id has been written. It exists for available_commands_update. The spec says
// to announce the commands right after the session is created, and the client
// (Zed, verified 2026-09-15) drops any session/update for a session whose
// session/new response it has not received yet — "Available commands: none".
// But the SDK writes the response only when NewSession returns, after anything
// the method itself sent, and offers no "after the response" callback. Watching
// the bytes is the one place where the order is certain.
type responseWriter struct {
	io.Writer
	mu    sync.Mutex
	hooks map[sdk.SessionId]func()
}

// afterResponse registers hook to run — once, in its own goroutine — after the
// response that carries sid has been written.
func (w *responseWriter) afterResponse(sid sdk.SessionId, hook func()) {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.hooks[sid] = hook
}

// Write forwards to the real output, then fires the hook of the session whose
// response this line is, if any. The hook runs in a goroutine: the SDK holds
// its write lock here, and the hook is about to write a notification.
func (w *responseWriter) Write(p []byte) (int, error) {
	n, err := w.Writer.Write(p)
	if err != nil {
		return n, err
	}
	if sid, ok := responseSessionID(p); ok {
		w.mu.Lock()
		hook := w.hooks[sid]
		delete(w.hooks, sid)
		w.mu.Unlock()
		if hook != nil {
			go hook()
		}
	}
	return n, nil
}

// responseSessionID recognises a JSON-RPC response whose result carries a
// sessionId — session/new (and session/load, which this agent does not
// offer). Notifications have no id and prompt responses no sessionId, so
// neither matches.
func responseSessionID(line []byte) (sdk.SessionId, bool) {
	var m struct {
		ID     json.RawMessage `json:"id"`
		Result struct {
			SessionId sdk.SessionId `json:"sessionId"`
		} `json:"result"`
	}
	if json.Unmarshal(line, &m) != nil || len(m.ID) == 0 || m.Result.SessionId == "" {
		return "", false
	}
	return m.Result.SessionId, true
}

// --- initialize ---------------------------------------------------------------

func (f *front) Initialize(_ context.Context, p sdk.InitializeRequest) (sdk.InitializeResponse, error) {
	fmt.Fprintf(os.Stderr, "[acp] initialize: client %q protocolVersion=%d\n",
		clientName(p.ClientInfo), p.ProtocolVersion)
	return sdk.InitializeResponse{
		ProtocolVersion: sdk.ProtocolVersionNumber, // 1
		// Everything absent is false: no loadSession, no image/audio prompts,
		// no MCP over http. Honest capabilities are what lets a client adapt.
		AgentCapabilities: sdk.AgentCapabilities{},
		AgentInfo:         &sdk.Implementation{Name: "bob", Title: sdk.Ptr("Bob (bash-first agent)"), Version: version},
		AuthMethods:       []sdk.AuthMethod{}, // the engine is local: nothing to log into
	}, nil
}

func clientName(impl *sdk.Implementation) string {
	if impl == nil {
		return "unknown"
	}
	return impl.Name
}

func (f *front) Authenticate(context.Context, sdk.AuthenticateRequest) (sdk.AuthenticateResponse, error) {
	return sdk.AuthenticateResponse{}, nil
}

// --- session/new --------------------------------------------------------------

func (f *front) NewSession(ctx context.Context, p sdk.NewSessionRequest) (sdk.NewSessionResponse, error) {
	f.mu.Lock()
	f.nextID++
	sid := sdk.SessionId(fmt.Sprintf("bob-%d-%d", os.Getpid(), f.nextID))
	f.sessions[sid] = &session{
		cwd:         p.Cwd, // absolute, per the spec; where bash runs for this session
		messages:    history.Fresh(f.system),
		allowAlways: map[string]bool{},
	}
	f.mu.Unlock()
	// mcpServers is where the client offers extra tools. bob's whole point is
	// the opposite — one built-in tool — so they are acknowledged and ignored.
	fmt.Fprintf(os.Stderr, "[acp] session %s cwd=%s mcpServers=%d (ignored)\n", sid, p.Cwd, len(p.McpServers))
	// The slash commands, announced as the spec asks — right after the session
	// is created — so the editor can list and complete them. "Right after"
	// means after the response: Zed keys its sessions on the id that response
	// carries and drops updates for an id it does not know yet. Hence the hook
	// on the writer rather than a SessionUpdate call here (see responseWriter).
	// The request context dies with this method; the hook needs one that lives.
	f.out.afterResponse(sid, func() {
		_ = f.conn.SessionUpdate(context.WithoutCancel(ctx), sdk.SessionNotification{
			SessionId: sid,
			Update:    sdk.SessionUpdate{AvailableCommandsUpdate: &sdk.SessionAvailableCommandsUpdate{AvailableCommands: availableCommands()}},
		})
	})
	return sdk.NewSessionResponse{SessionId: sid}, nil
}

// compactCommand is the REPL's /compact, spelled here because ACP has no other
// place for it; /new is shared through internal/session. /quit and /abort are
// deliberately absent: the editor owns the process and ends it by closing
// stdin, and it cancels a turn with its own stop key through session/cancel.
const compactCommand = "/compact"

// isCommand says whether input, trimmed, is exactly the given command — the
// same rule as session.IsNewCommand: "/compact " is the command, "/compact
// now" is a question for the model.
func isCommand(input, command string) bool {
	return strings.TrimSpace(input) == command
}

// availableCommands lists the slash commands this front end handles itself,
// in the shape ACP advertises them: bare names, the client adds the slash.
// Only what session/prompt actually intercepts belongs here — announcing
// /quit or /abort, which the REPL has and this front end does not, would
// promise the editor something the next prompt would hand to the model.
func availableCommands() []sdk.AvailableCommand {
	return []sdk.AvailableCommand{
		{Name: history.NewCommandName, Description: "Clear the history and start a new session"},
		{Name: strings.TrimPrefix(compactCommand, "/"), Description: "Compress the history now, keeping the last turns"},
	}
}

// --- session/prompt -----------------------------------------------------------

func (f *front) Prompt(ctx context.Context, p sdk.PromptRequest) (sdk.PromptResponse, error) {
	f.mu.Lock()
	s, ok := f.sessions[p.SessionId]
	f.mu.Unlock()
	if !ok {
		return sdk.PromptResponse{}, sdk.NewInvalidParams(map[string]any{"sessionId": p.SessionId})
	}

	input := flatten(p.Prompt)

	f.turnMu.Lock()
	defer f.turnMu.Unlock()

	// A slash command is answered here, without a turn: it is the REPL's "only
	// between two generations" rule, met for free — turnMu guarantees no
	// generation is running on this session's history while we replace it.
	switch {
	case history.IsNewCommand(input):
		return f.startNewSession(ctx, p.SessionId, s), nil
	case isCommand(input, compactCommand):
		return f.compactSession(ctx, p.SessionId, s), nil
	}

	// One cancellable context per turn: session/cancel fires s.cancel, the
	// generation and any pending permission request fall with it, and the spec's
	// contract — still answer the prompt, with stopReason "cancelled" — is met
	// at the bottom of this function.
	turnCtx, cancel := context.WithCancel(ctx)
	f.mu.Lock()
	s.cancel = cancel
	f.mu.Unlock()
	defer func() {
		f.mu.Lock()
		s.cancel = nil
		f.mu.Unlock()
		cancel()
	}()

	// The sink is what turns the tools' events into session/update messages;
	// installing it is what switches engine+tools to "ACP mode" for this turn.
	ui.SetActive(&sink{f: f, ctx: turnCtx, sid: p.SessionId, s: s})
	defer ui.SetActive(nil)

	// A typed "@path" — the editor's picker bypassed, or an editor without
	// one — is understood like the resource_link the picker would have sent,
	// against the session's cwd, which is the project the editor opened.
	input, attached := mention.Expand(input, s.cwd)
	for _, a := range attached {
		fmt.Fprintf(os.Stderr, "[acp] session %s: attached %s\n", p.SessionId, a.Path)
	}
	s.messages = append(s.messages, ai.NewUserTextMessage(input))
	before := len(s.messages)

	resp, history, err := f.e.Generate(turnCtx, s.messages, f.tools)

	// Same rule as the REPL, for the same reason: keep the FULL history even
	// when the turn failed or was cancelled — the commands that already ran are
	// part of the conversation, and forgetting them makes the model re-run or
	// invent them on the next turn.
	switch {
	case len(history) > before:
		s.messages = history
	case resp != nil && resp.Message != nil:
		s.messages = append(s.messages, resp.Message)
	case err != nil:
		s.messages = s.messages[:before-1] // question without an answer: forget it
	}

	// The ⚙ recap of the REPL, one line on stderr: the editor already renders
	// every tool call, but the trail is what lets `dev: open acp logs` say at a
	// glance whether the model worked or improvised.
	turn := s.messages[min(before, len(s.messages)):]
	fmt.Fprintf(os.Stderr, "[acp] turn done: %d command(s) · %d file op(s) · %d skill(s)\n",
		len(engine.CommandList(turn)), len(engine.FileOps(turn)), engine.Skills(turn))

	switch {
	case turnCtx.Err() != nil:
		// The turn was cancelled — and both roads lead here: our Cancel method
		// fires s.cancel (the child), AND the SDK cancels the request's own
		// context when session/cancel arrives (the parent — measured: with the
		// parent-must-be-intact test the REPL uses, a cancelled turn answered
		// "Internal error: stream error: context canceled" instead of honouring
		// the spec's contract). Either way the answer is the same: respond to
		// session/prompt anyway, with stopReason "cancelled".
		return sdk.PromptResponse{StopReason: sdk.StopReasonCancelled}, nil
	case err != nil:
		// One line, in the provider's words — the same text the REPL prints
		// between brackets, delivered as a JSON-RPC error.
		return sdk.PromptResponse{}, sdk.NewInternalError(map[string]any{"error": f.e.Explain(err)})
	}
	return sdk.PromptResponse{StopReason: sdk.StopReasonEndTurn}, nil
}

// startNewSession is the ACP side of /new: the session keeps its id, its cwd
// and the "allow always" grants — the dialog promised them "for this session",
// and the session, as the editor sees it, is the same one — but its history
// goes back to the system prompt alone, and the server's token count of the
// last context is dropped with it, as after a compression. The one line the
// REPL prints travels as an agent message so the editor shows it in the
// thread, and the turn ends right there: the model is not consulted.
func (f *front) startNewSession(ctx context.Context, sid sdk.SessionId, s *session) sdk.PromptResponse {
	forgotten := resetSession(s, f.system)
	f.e.ForgetInputTokens()
	fmt.Fprintf(os.Stderr, "[acp] session %s: new session, %d message(s) forgotten\n", sid, forgotten)
	_ = f.conn.SessionUpdate(ctx, sdk.SessionNotification{
		SessionId: sid,
		Update:    sdk.UpdateAgentMessageText(fmt.Sprintf("🆕 New session: %d message(s) forgotten.", forgotten)),
	})
	return sdk.PromptResponse{StopReason: sdk.StopReasonEndTurn}
}

// compactSession is the ACP side of /compact: the REPL's compactHistory with
// the report sent as an agent message instead of printed. Forced, like the
// REPL's — the threshold is for the automatic path, which this front end does
// not have — so "nothing to compact" is said rather than swallowed. Runs under
// turnMu: Genkit is not holding this history.
func (f *front) compactSession(ctx context.Context, sid sdk.SessionId, s *session) sdk.PromptResponse {
	cfg := config.Cfg.Context
	line, compressed := compactHistory(ctx, s, cfg, func(ctx context.Context, request []*ai.Message) (string, error) {
		return f.e.Summarize(ctx, request, cfg.SummaryMaxTokens)
	}, f.e.Explain)
	if compressed {
		f.e.ForgetInputTokens() // the server's measure described the history just replaced
	}
	fmt.Fprintf(os.Stderr, "[acp] session %s: /compact — %s\n", sid, line)
	_ = f.conn.SessionUpdate(ctx, sdk.SessionNotification{
		SessionId: sid,
		Update:    sdk.UpdateAgentMessageText(line),
	})
	return sdk.PromptResponse{StopReason: sdk.StopReasonEndTurn}
}

// compactHistory compresses the session's history in place when there is
// something to compress, and returns the one line to show — the REPL's words,
// so the editor and the terminal agree — and whether the history changed.
// Kept apart from compactSession, with the summariser and the error explainer
// injected, so the three outcomes are testable without an engine or a
// connection.
func compactHistory(ctx context.Context, s *session, cfg config.ContextConfig, summarize compact.Summarizer, explain func(error) string) (line string, compressed bool) {
	res, err := compact.Compact(ctx, s.messages, cfg, summarize)
	switch {
	case errors.Is(err, compact.ErrNothingToCompact):
		return fmt.Sprintf("🗜️ nothing to compact: %d message(s), no turn older than the last %d", len(s.messages), cfg.KeepLastTurns), false
	case err != nil:
		return fmt.Sprintf("[compact: failed, history kept: %s]", explain(err)), false
	}
	s.messages = res.Messages
	return "🗜️ " + res.Report(), true
}

// resetSession replaces the session's history with a fresh one and returns how
// many messages were forgotten. Kept apart from startNewSession so the rule
// — history reset, permissions kept — is testable without a connection.
func resetSession(s *session, system string) int {
	forgotten := history.Forgotten(s.messages)
	s.messages = history.Fresh(system)
	return forgotten
}

// flatten turns the prompt's content blocks into the plain text bob reads.
// Text is taken as-is; a resource_link (a file the user @-mentioned through
// the editor's picker) becomes its path, in the same shape mention.Expand
// gives a typed @path — the model has tools to read it, that is the whole
// idea; anything else (image, audio, embedded resource) was not announced in
// our capabilities, so a compliant client never sends it.
func flatten(blocks []sdk.ContentBlock) string {
	var b strings.Builder
	for _, c := range blocks {
		switch {
		case c.Text != nil:
			b.WriteString(c.Text.Text)
		case c.ResourceLink != nil:
			fmt.Fprintf(&b, "\n[attached file: %s]", strings.TrimPrefix(c.ResourceLink.Uri, "file://"))
		}
	}
	return b.String()
}

// --- session/cancel -----------------------------------------------------------

func (f *front) Cancel(_ context.Context, p sdk.CancelNotification) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if s, ok := f.sessions[p.SessionId]; ok && s.cancel != nil {
		s.cancel()
	}
	return nil
}

// --- the rest of sdk.Agent: not supported, and saying so -----------------------
//
// The SDK requires the full interface; answering "method not found" is the
// protocol's way of saying a capability we never announced is indeed absent.

func (f *front) SetSessionMode(context.Context, sdk.SetSessionModeRequest) (sdk.SetSessionModeResponse, error) {
	return sdk.SetSessionModeResponse{}, sdk.NewMethodNotFound(sdk.AgentMethodSessionSetMode)
}

func (f *front) SetSessionConfigOption(context.Context, sdk.SetSessionConfigOptionRequest) (sdk.SetSessionConfigOptionResponse, error) {
	return sdk.SetSessionConfigOptionResponse{}, sdk.NewMethodNotFound(sdk.AgentMethodSessionSetConfigOption)
}

func (f *front) ListSessions(context.Context, sdk.ListSessionsRequest) (sdk.ListSessionsResponse, error) {
	return sdk.ListSessionsResponse{}, sdk.NewMethodNotFound(sdk.AgentMethodSessionList)
}

func (f *front) ResumeSession(context.Context, sdk.ResumeSessionRequest) (sdk.ResumeSessionResponse, error) {
	return sdk.ResumeSessionResponse{}, sdk.NewMethodNotFound(sdk.AgentMethodSessionResume)
}

func (f *front) CloseSession(context.Context, sdk.CloseSessionRequest) (sdk.CloseSessionResponse, error) {
	return sdk.CloseSessionResponse{}, sdk.NewMethodNotFound(sdk.AgentMethodSessionClose)
}

func (f *front) Logout(context.Context, sdk.LogoutRequest) (sdk.LogoutResponse, error) {
	return sdk.LogoutResponse{}, sdk.NewMethodNotFound(sdk.AgentMethodLogout)
}