package acp

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"strconv"
	"sync"
	"sync/atomic"
)

// Handler is the meaning behind the methods: one implementation per agent.
//
// Every method may be called concurrently, and Prompt in particular is
// expected to run for minutes while Cancel arrives on another goroutine. A
// handler that serializes its own state is a handler that works; this package
// deliberately does no locking on its behalf, because only the handler knows
// what a session is.
type Handler interface {
	Initialize(context.Context, InitializeRequest) (InitializeResponse, error)
	// Authenticate must not fail for an unknown method id. Clients send one
	// unconditionally and do not catch the error, so a refusal here is a
	// session that never starts rather than a login prompt.
	Authenticate(context.Context, AuthenticateRequest) error
	NewSession(context.Context, NewSessionRequest) (NewSessionResponse, error)
	Prompt(context.Context, PromptRequest) (PromptResponse, error)
	// Cancel is a notification: there is nothing to answer and nothing to
	// fail. It must return promptly — the turn it cancels reports the outcome
	// itself, by ending with StopCancelled.
	Cancel(context.Context, CancelParams)
}

// Extender is the optional half of [Handler]: methods outside the core set.
// A handler that does not implement it answers CodeMethodNotFound to anything
// unrecognised, which is the correct answer and not an error.
type Extender interface {
	Extra(ctx context.Context, method string, params json.RawMessage) (any, error)
}

// ErrClosed is what a pending [Conn.Request] fails with when the connection
// ends before the answer arrives.
var ErrClosed = errors.New("acp: connection closed")

// Conn is one ACP connection.
type Conn struct {
	h Handler
	r io.Reader
	w io.Writer

	// wmu serializes writes. Every frame is marshalled whole and written in
	// one call, so interleaving is the only thing to guard against.
	wmu sync.Mutex

	nextID atomic.Int64

	pmu     sync.Mutex
	pending map[string]chan *Message

	closeOnce sync.Once
	done      chan struct{}

	// wg counts handler goroutines, so Serve can wait for them to finish
	// writing before it returns and the streams go away underneath them.
	wg sync.WaitGroup
}

// NewConn wires a handler to a pair of streams. Nothing is read or written
// until [Conn.Serve] runs.
func NewConn(r io.Reader, w io.Writer, h Handler) *Conn {
	return &Conn{
		h:       h,
		r:       r,
		w:       w,
		pending: make(map[string]chan *Message),
		done:    make(chan struct{}),
	}
}

// Serve reads until the stream ends, dispatching as it goes. A clean EOF is
// how a client says goodbye, so it returns nil; anything else is the error
// that stopped it.
//
// Cancelling ctx stops the handlers but not the read itself: a blocking read
// on an arbitrary io.Reader cannot be interrupted portably. Closing the reader
// is what ends Serve.
func (c *Conn) Serve(ctx context.Context) error {
	ctx, cancel := context.WithCancel(ctx)

	err := c.loop(ctx)

	// Order matters. Cancel first so a turn in flight stops, then release
	// anyone blocked waiting for an answer that is no longer coming, then wait
	// for the handlers to finish with the streams.
	cancel()
	c.close()
	c.wg.Wait()
	return err
}

func (c *Conn) loop(ctx context.Context) error {
	// ReadBytes rather than a Scanner: a prompt can carry a base64 image, and
	// a line-length limit chosen here would be a decode failure nobody could
	// explain.
	br := bufio.NewReaderSize(c.r, 64*1024)
	for {
		line, err := br.ReadBytes('\n')
		if len(bytes.TrimSpace(line)) > 0 {
			c.handleLine(ctx, line)
		}
		switch {
		case errors.Is(err, io.EOF):
			return nil
		case err != nil:
			return err
		}
	}
}

func (c *Conn) handleLine(ctx context.Context, line []byte) {
	var m Message
	if err := json.Unmarshal(line, &m); err != nil {
		// A frame we cannot parse might still have a recoverable id, in which
		// case the peer is waiting and deserves an answer. If not, there is
		// nobody to tell and the loop carries on: one bad line is not a reason
		// to end a session.
		if id := salvageID(line); id != nil {
			c.respondError(id, Errorf(CodeParse, "invalid JSON"))
		}
		return
	}

	switch {
	case m.ID != nil && m.Method == "":
		c.deliver(&m)
	case m.ID != nil:
		c.wg.Add(1)
		go func() {
			defer c.wg.Done()
			c.serveRequest(ctx, &m)
		}()
	case m.Method != "":
		c.wg.Add(1)
		go func() {
			defer c.wg.Done()
			c.serveNotification(ctx, &m)
		}()
	}
}

// salvageID pulls the id out of a line that did not parse as a whole, so a
// peer blocked on a request gets an error instead of a hang.
//
// It streams tokens rather than unmarshalling, because the interesting case is
// a frame that was cut off: the id is usually intact and everything after it
// is not, and a whole-value decode fails on exactly those lines.
func salvageID(line []byte) *json.RawMessage {
	dec := json.NewDecoder(bytes.NewReader(line))
	if tok, err := dec.Token(); err != nil || tok != json.Delim('{') {
		return nil
	}
	for {
		tok, err := dec.Token()
		if err != nil {
			return nil
		}
		key, ok := tok.(string)
		if !ok {
			return nil // the closing brace, or something that is not an object
		}
		var raw json.RawMessage
		if err := dec.Decode(&raw); err != nil {
			return nil
		}
		if key == "id" {
			return &raw
		}
	}
}

func (c *Conn) serveRequest(ctx context.Context, m *Message) {
	result, err := c.call(ctx, m.Method, m.Params)
	if err != nil {
		c.respondError(m.ID, asError(err))
		return
	}
	c.respond(m.ID, result)
}

func (c *Conn) serveNotification(ctx context.Context, m *Message) {
	if m.Method != MethodCancel {
		// Unknown notifications are ignored by design: that is what makes it
		// safe for a client to send one we have not heard of.
		return
	}
	var params CancelParams
	if err := decodeParams(m.Params, &params); err != nil {
		return
	}
	c.h.Cancel(ctx, params)
}

// call dispatches one request to the handler. The returned value is whatever
// goes in the result field; nil becomes an empty object, because a JSON-RPC
// response with no result at all is not a success.
func (c *Conn) call(ctx context.Context, method string, params json.RawMessage) (any, error) {
	switch method {
	case MethodInitialize:
		var req InitializeRequest
		if err := decodeParams(params, &req); err != nil {
			return nil, err
		}
		return c.h.Initialize(ctx, req)

	case MethodAuthenticate:
		var req AuthenticateRequest
		if err := decodeParams(params, &req); err != nil {
			return nil, err
		}
		return nil, c.h.Authenticate(ctx, req)

	case MethodNewSession:
		var req NewSessionRequest
		if err := decodeParams(params, &req); err != nil {
			return nil, err
		}
		return c.h.NewSession(ctx, req)

	case MethodPrompt:
		var req PromptRequest
		if err := decodeParams(params, &req); err != nil {
			return nil, err
		}
		return c.h.Prompt(ctx, req)

	default:
		if ex, ok := c.h.(Extender); ok {
			return ex.Extra(ctx, method, params)
		}
		return nil, Errorf(CodeMethodNotFound, "unknown method %q", method)
	}
}

// Notify sends a notification. It never has an id, which is what stops the
// peer from waiting for an answer to it.
func (c *Conn) Notify(ctx context.Context, method string, params any) error {
	raw, err := json.Marshal(params)
	if err != nil {
		return err
	}
	return c.write(&Message{JSONRPC: jsonrpcVersion, Method: method, Params: raw})
}

// Request calls the peer and waits for the answer, decoding it into out when
// out is non-nil. It returns ctx.Err() if the caller gives up first and
// [ErrClosed] if the connection ends before the peer replies — a permission
// request outliving its connection has to fail rather than hang the turn.
func (c *Conn) Request(ctx context.Context, method string, params, out any) error {
	raw, err := json.Marshal(params)
	if err != nil {
		return err
	}

	id := json.RawMessage(strconv.FormatInt(c.nextID.Add(1), 10))
	key := string(id)
	ch := make(chan *Message, 1)

	c.pmu.Lock()
	c.pending[key] = ch
	c.pmu.Unlock()
	defer func() {
		c.pmu.Lock()
		delete(c.pending, key)
		c.pmu.Unlock()
	}()

	if err := c.write(&Message{JSONRPC: jsonrpcVersion, ID: &id, Method: method, Params: raw}); err != nil {
		return err
	}

	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-c.done:
		return ErrClosed
	case reply := <-ch:
		if reply.Error != nil {
			return reply.Error
		}
		if out == nil || len(reply.Result) == 0 {
			return nil
		}
		return json.Unmarshal(reply.Result, out)
	}
}

func (c *Conn) deliver(m *Message) {
	c.pmu.Lock()
	ch := c.pending[string(*m.ID)]
	c.pmu.Unlock()
	if ch != nil {
		// Buffered, and the entry is removed by the waiter, so this cannot
		// block and cannot be delivered twice.
		select {
		case ch <- m:
		default:
		}
	}
}

func (c *Conn) respond(id *json.RawMessage, result any) {
	if result == nil {
		result = struct{}{}
	}
	raw, err := json.Marshal(result)
	if err != nil {
		c.respondError(id, asError(err))
		return
	}
	_ = c.write(&Message{JSONRPC: jsonrpcVersion, ID: id, Result: raw})
}

func (c *Conn) respondError(id *json.RawMessage, e *Error) {
	if id == nil {
		null := json.RawMessage("null")
		id = &null
	}
	_ = c.write(&Message{JSONRPC: jsonrpcVersion, ID: id, Error: e})
}

func (c *Conn) write(m *Message) error {
	b, err := json.Marshal(m)
	if err != nil {
		return err
	}
	b = append(b, '\n')

	c.wmu.Lock()
	defer c.wmu.Unlock()
	_, err = c.w.Write(b)
	return err
}

func (c *Conn) close() {
	c.closeOnce.Do(func() { close(c.done) })
}