// 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"

	"codey/internal/acp"
	"codey/internal/agent"
	"codey/internal/approval"
	"codey/internal/llm"
)

// 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 {
	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
	}
}

// NewAgentFunc builds the agent for one session.
//
// 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. model is empty for the default. state is the
// conversation, passed in rather than created here so that selecting another
// model can rebuild the agent around a history that survives the switch.
type NewAgentFunc func(root, model string, state *agent.State) (*agent.Agent, 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)
	return s.conn.Serve(ctx)
}

// 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,
	}
	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.
func (s *Server) build(sess *session) error {
	a, err := s.opts.NewAgent(sess.root, sess.model, sess.state)
	if err != nil {
		return acp.Errorf(acp.CodeInternal, "cannot start an agent in %s: %v", sess.root, err)
	}
	a.Approve = sess.policy
	a.Observe = sess.emit.handle
	sess.agent = a
	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.agent.RunBlocks(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
}

// 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 ""
	}
}