// Package serve runs Codey as an agent another program drives, speaking the
// Agent Client Protocol over stdin and stdout.
//
// It is a front end, the same way the TUI and the one-shot `-p` mode are: it
// owns no agent behaviour and makes no decisions the other two do not. What it
// adds is a translation. The agent reports what it is doing through
// [agent.Agent.Observe] and asks permission through [agent.Agent.Approve];
// this package turns the first into `session/update` notifications and the
// second into a `session/request_permission` round trip, and turns the
// client's prompts back into turns.
//
// One process serves many sessions. A client that opens three threads gets
// three conversations here, each with its own history, its own tool set rooted
// where it asked, and its own record of what the user has permanently allowed
// — because "always allow this" is a statement about one conversation, not
// about the process.
package serve

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

	"github.com/johannSo/codey/internal/acp"
	"github.com/johannSo/codey/internal/agent"
	"github.com/johannSo/codey/internal/approval"
	"github.com/johannSo/codey/internal/llm"
	"github.com/johannSo/codey/internal/subagent"
)

// PermissionMode is how tool calls are gated, chosen on the command line by
// whoever launched us. It is a floor, not a policy: the client can still be
// asked about anything the mode does not settle.
type PermissionMode string

const (
	// ModeAsk asks about every mutating call. The default, and the only mode
	// that is safe against a tree you care about.
	ModeAsk PermissionMode = "default"
	// ModeAcceptEdits lets file edits through and still asks before a command
	// runs. An edit is reviewable after the fact; a command is not.
	ModeAcceptEdits PermissionMode = "acceptEdits"
	// ModeAuto asks only about calls that cannot be pinned down — an opaque
	// shell command, a search whose query is different every time. Everything
	// whose exact shape could have been remembered is allowed.
	ModeAuto PermissionMode = "auto"
	// ModeFullAccess never asks. Only against a throwaway clone or inside a
	// container.
	ModeFullAccess PermissionMode = "full-access"
)

// allows reports whether the mode settles this call without asking.
func (m PermissionMode) allows(req approval.Request) bool {
	// A prompt a standing ask rule insisted on is never settled by the floor:
	// the rule was written to stop exactly that.
	if req.Forced != "" {
		return false
	}
	switch m {
	case ModeFullAccess:
		return true
	case ModeAuto:
		return req.Rememberable
	case ModeAcceptEdits:
		return req.Tool == "edit_file" || req.Tool == "write_file"
	default:
		return false
	}
}

// AgentRequest is one session asking to be built.
type AgentRequest struct {
	// Root is the directory the client asked to work in, which is not
	// necessarily the one this process was started in: a client that manages
	// worktrees hands out a different one per thread.
	Root string
	// Model is empty for whatever the process was configured with.
	Model string
	// State is the conversation. It is passed in rather than created by the
	// builder so that selecting another model can rebuild the agent around a
	// history that survives the switch.
	State *agent.State
	// Approve gates tool calls — the session's own and its subagents'. What
	// you allow, you allow for the conversation, not for one agent in it.
	Approve *approval.Policy
}

// AgentSet is what comes back: the agent, and the subagents it may start.
type AgentSet struct {
	Agent *agent.Agent
	// Agents is nil when this session may not spawn.
	Agents Agents
}

// Agents is a session's subagent supervisor, as this package needs it.
// *subagent.Supervisor is the implementation; the interface is here so a test
// can stand in for one without starting real children.
type Agents interface {
	// Subscribe registers the two callbacks, which are called from the
	// children's own goroutines.
	Subscribe(observe func(subagent.Event), changed func())
	// TakeReports drains the reports nobody has been handed yet, and
	// ClaimReport says whether this front end is still the one to hand a given
	// report over — agent_wait may have collected it in the meantime.
	TakeReports() []subagent.Info
	ClaimReport(id string) bool
	Snapshot() []subagent.Info
	Shutdown()
}

// NewAgentFunc builds one session.
type NewAgentFunc func(AgentRequest) (AgentSet, error)

// Options configures a [Server].
type Options struct {
	// NewAgent is required.
	NewAgent NewAgentFunc
	// Version is reported in the initialize handshake.
	Version string
	// Permission is the approval floor for every session.
	Permission PermissionMode
	// Models is the picker the client shows, and DefaultModel is what is
	// selected in it. Both may be empty, in which case the client is told
	// nothing and uses whatever the agent was configured with.
	Models       []acp.ModelInfo
	DefaultModel string
}

// Server is the ACP end of the process.
type Server struct {
	opts Options
	conn *acp.Conn

	mu       sync.Mutex
	sessions map[string]*session
}

// New builds a server. Nothing runs until [Server.Run].
func New(opts Options) *Server {
	return &Server{opts: opts, sessions: make(map[string]*session)}
}

// Run serves one connection until the stream ends. A clean end of input is how
// a client says goodbye, and it returns nil.
func (s *Server) Run(ctx context.Context, r io.Reader, w io.Writer) error {
	s.conn = acp.NewConn(r, w, s)
	defer s.shutdown()
	return s.conn.Serve(ctx)
}

// shutdown stops every session's subagents. A child outliving the connection
// would keep editing a tree nobody is watching, and nobody would be left to
// read what it found.
func (s *Server) shutdown() {
	s.mu.Lock()
	sessions := make([]*session, 0, len(s.sessions))
	for _, sess := range s.sessions {
		sessions = append(sessions, sess)
	}
	s.mu.Unlock()

	for _, sess := range sessions {
		if sess.agents != nil {
			sess.agents.Shutdown()
		}
	}
}

// Initialize answers the handshake.
//
// The protocol version we answer with is ours, not the client's: agreeing to a
// revision we do not implement would be a handshake that succeeds and a
// session that misbehaves later, which is harder to diagnose than a refusal.
func (s *Server) Initialize(_ context.Context, _ acp.InitializeRequest) (acp.InitializeResponse, error) {
	resp := acp.InitializeResponse{
		ProtocolVersion: acp.ProtocolVersion,
		AgentInfo:       &acp.Implementation{Name: "codey", Title: "Codey", Version: s.opts.Version},
		AgentCapabilities: &acp.AgentCapabilities{
			// Images go through: a screenshot is the fastest way to say what
			// is wrong with a layout. Session loading does not, yet — and not
			// advertising it is what stops a client from asking.
			PromptCapabilities: &acp.PromptCapabilities{Image: true},
		},
		AuthMethods: []acp.AuthMethod{},
	}

	// The model list, for a client that is only probing. A health check
	// spawns the agent, shakes hands and stops; answering here is what puts
	// the models on its provider card without it having to open a session it
	// will not use.
	if len(s.opts.Models) > 0 {
		resp.Meta = map[string]any{acp.ModelStateMeta: s.modelState()}
	}
	return resp, nil
}

func (s *Server) modelState() *acp.SessionModelState {
	if len(s.opts.Models) == 0 {
		return nil
	}
	return &acp.SessionModelState{
		CurrentModelID:  s.opts.DefaultModel,
		AvailableModels: s.opts.Models,
	}
}

// Authenticate always succeeds.
//
// Codey takes its credential from the environment before this process starts,
// so there is nothing to log in to and no method id worth checking. Clients
// send this unconditionally and do not catch a failure, so an error here would
// be a session that never opens, over a question we were never asking.
func (s *Server) Authenticate(context.Context, acp.AuthenticateRequest) error { return nil }

// NewSession opens a conversation rooted where the client asked.
func (s *Server) NewSession(_ context.Context, req acp.NewSessionRequest) (acp.NewSessionResponse, error) {
	if req.Cwd == "" {
		return acp.NewSessionResponse{}, acp.Errorf(acp.CodeInvalidParams, "session/new needs a cwd")
	}

	sess := &session{
		id:          newSessionID(),
		root:        req.Cwd,
		state:       agent.NewState(),
		model:       s.opts.DefaultModel,
		agentsMoved: make(chan struct{}, 1),
	}
	sess.emit = newEmitter(notifier(s.conn, sess.id))
	sess.policy = approval.New(sess.promptFor(s.conn, s.opts.Permission))
	sess.policy.YOLO = s.opts.Permission == ModeFullAccess

	if err := s.build(sess); err != nil {
		return acp.NewSessionResponse{}, err
	}

	s.mu.Lock()
	s.sessions[sess.id] = sess
	s.mu.Unlock()

	resp := acp.NewSessionResponse{SessionID: sess.id}
	if state := s.modelState(); state != nil {
		state.CurrentModelID = sess.model
		resp.Models = state
	}
	return resp, nil
}

// build attaches an agent to a session, or replaces the one it has. The caller
// holds whatever lock the session needs.
//
// A rebuild — which is what selecting another model is — reuses the subagents
// the session already has. They belong to the conversation, not to the model
// the parent happens to be running on, and killing four working children
// because somebody changed a dropdown would be a surprising thing to do.
func (s *Server) build(sess *session) error {
	built, err := s.opts.NewAgent(AgentRequest{
		Root:    sess.root,
		Model:   sess.model,
		State:   sess.state,
		Approve: sess.policy,
	})
	if err != nil {
		return acp.Errorf(acp.CodeInternal, "cannot start an agent in %s: %v", sess.root, err)
	}

	built.Agent.Approve = sess.policy
	built.Agent.Observe = sess.emit.handle
	sess.agent = built.Agent

	if sess.agents == nil && built.Agents != nil {
		sess.agents = built.Agents
		sess.watchAgents()
	}
	return nil
}

// Prompt runs one turn.
func (s *Server) Prompt(ctx context.Context, req acp.PromptRequest) (acp.PromptResponse, error) {
	sess, err := s.lookup(req.SessionID)
	if err != nil {
		return acp.PromptResponse{}, err
	}

	blocks := promptBlocks(req.Prompt)
	if len(blocks) == 0 {
		return acp.PromptResponse{StopReason: acp.StopEndTurn}, nil
	}

	sess.mu.Lock()
	defer sess.mu.Unlock()

	turnCtx := sess.begin(ctx)
	defer sess.end()

	runErr := sess.runTurn(turnCtx, blocks)

	// Whatever happened, say the last of what was written before reporting how
	// the turn ended. A buffer dropped on the way out is a sentence the user
	// never sees.
	sess.emit.Flush()

	switch {
	case sess.wasCancelled(), errors.Is(runErr, context.Canceled):
		// Not an error: the session is intact and the next message can correct
		// course. That is the whole point of a stop button.
		return acp.PromptResponse{StopReason: acp.StopCancelled}, nil

	case errors.Is(runErr, agent.ErrMaxTurns):
		sess.emit.chunk(acp.UpdateAgentMessageChunk, runErr.Error())
		return acp.PromptResponse{StopReason: acp.StopMaxTurnRequests}, nil

	case runErr != nil:
		// The failure goes in the transcript rather than into a protocol error
		// alone, because a protocol error is a line in the client's log and
		// this is something the user has to read to know what to do next.
		sess.emit.chunk(acp.UpdateAgentMessageChunk, "codey: "+runErr.Error())
		return acp.PromptResponse{StopReason: acp.StopRefusal}, nil
	}

	return acp.PromptResponse{StopReason: acp.StopEndTurn}, nil
}

// maxReportRounds caps how many times one turn may be restarted to hand a
// subagent's report to the parent. A parent that answers every report by
// spawning another agent would otherwise never finish the turn it is in, and
// a cap is a cheaper guard than reasoning about when that is reasonable.
//
// Hitting it ends the turn with the outstanding reports still queued and
// unclaimed, so the next turn hands them over rather than losing them.
const maxReportRounds = 32

// runTurn runs the agent, and then keeps the turn open for as long as this
// session's subagents are still working.
//
// A turn is the protocol's unit of work, and a child's output belongs inside
// the one that started it: end the turn while four agents are still running
// and their rows have nowhere to appear, their reports have nobody to be
// handed to, and the thread looks finished while the machine is busy. So the
// turn ends when the work does.
//
// Each report the parent is entitled to comes back as a message, which is what
// the parent's own prompt promised would happen, and that runs the agent
// again. Cancelling breaks out of all of it: children keep working, and the
// next turn will pick their reports up.
func (s *session) runTurn(ctx context.Context, blocks []llm.ContentBlock) error {
	if _, err := s.agent.RunBlocks(ctx, blocks); err != nil {
		return err
	}

	for rounds := 0; ; {
		if ctx.Err() != nil {
			return ctx.Err()
		}
		// Counted against report deliveries only. Waking up because a child
		// changed status and finding nothing to hand over costs nothing and
		// must not eat the budget, or a long job would end its own turn.
		if rounds >= maxReportRounds {
			return nil
		}

		if text := s.reportBlocks(); text != "" {
			rounds++
			// What the agent said before the report arrived is finished. Left
			// buffered, the end of one message and the start of the next would
			// be glued into one chunk, and the client would show two answers
			// as a single run-on sentence.
			s.emit.Flush()
			if _, err := s.agent.RunBlocks(ctx, []llm.ContentBlock{llm.TextBlock(text)}); err != nil {
				return err
			}
			continue
		}
		if !s.agentsLive() {
			return nil
		}

		// Nothing to hand over and something still running: wait to be woken.
		select {
		case <-s.agentsMoved:
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

// Cancel stops the turn in flight. An unknown session is not an error worth
// reporting: a stop arriving after the session closed asks for something that
// already happened.
func (s *Server) Cancel(_ context.Context, params acp.CancelParams) {
	if sess, err := s.lookup(params.SessionID); err == nil {
		sess.stop()
	}
}

// Extra answers the methods outside the core set.
func (s *Server) Extra(_ context.Context, method string, params json.RawMessage) (any, error) {
	switch method {
	case "session/set_model":
		var req acp.SetModelRequest
		if err := json.Unmarshal(params, &req); err != nil {
			return nil, acp.Errorf(acp.CodeInvalidParams, "cannot decode params: %v", err)
		}
		sess, err := s.lookup(req.SessionID)
		if err != nil {
			return nil, err
		}
		sess.mu.Lock()
		defer sess.mu.Unlock()
		if req.ModelID == sess.model {
			return struct{}{}, nil
		}
		sess.model = req.ModelID
		// The conversation comes along: only the agent around it is replaced.
		if err := s.build(sess); err != nil {
			return nil, err
		}
		return struct{}{}, nil

	case "session/set_config_option":
		// Accepted and ignored. Some clients send one unconditionally on every
		// model selection, and refusing an option we never advertised would
		// fail a turn over a setting we do not have.
		return struct{}{}, nil

	default:
		return nil, acp.Errorf(acp.CodeMethodNotFound, "unknown method %q", method)
	}
}

func (s *Server) lookup(id string) (*session, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	sess, ok := s.sessions[id]
	if !ok {
		return nil, acp.Errorf(acp.CodeInvalidParams, "no such session %q", id)
	}
	return sess, nil
}

// promptBlocks turns the client's content blocks into the agent's.
//
// Text and images pass through. Anything else is named rather than dropped: a
// client that attached a file we cannot decode has told the user it sent
// something, and a model that never hears about it will answer a question
// nobody asked.
func promptBlocks(blocks []acp.ContentBlock) []llm.ContentBlock {
	var out []llm.ContentBlock
	var text strings.Builder

	for _, b := range blocks {
		switch b.Type {
		case "text":
			if b.Text != "" {
				if text.Len() > 0 {
					text.WriteString("\n")
				}
				text.WriteString(b.Text)
			}
		case "image":
			if b.Data != "" && b.MimeType != "" {
				// Text so far goes in first, so the image lands where it was
				// written rather than at the end of the message.
				if t := strings.TrimSpace(text.String()); t != "" {
					out = append(out, llm.TextBlock(t))
					text.Reset()
				}
				out = append(out, llm.ImageBlock(b.MimeType, b.Data))
			}
		default:
			if name := attachmentName(b); name != "" {
				if text.Len() > 0 {
					text.WriteString("\n")
				}
				fmt.Fprintf(&text, "[attached %s: %s]", b.Type, name)
			}
		}
	}

	if t := strings.TrimSpace(text.String()); t != "" {
		out = append(out, llm.TextBlock(t))
	}
	return out
}

func attachmentName(b acp.ContentBlock) string {
	switch {
	case b.Name != "":
		return b.Name
	case b.URI != "":
		return b.URI
	default:
		return ""
	}
}