package jsonrpc

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"sync"
)

// Conn is a JSON-RPC 2.0 connection over a framed stream.
//
// It is safe for concurrent use: several goroutines may call and notify at
// once, and replies are matched back to their callers by request id.
type Conn struct {
	stream io.ReadWriteCloser
	reader *bufio.Reader
	framer Framer

	onNotify  NotificationFunc
	onRequest RequestFunc

	writeMu sync.Mutex

	mu      sync.Mutex
	nextID  int64
	pending map[int64]chan *message
	closed  bool
}

// NewConn returns a connection over stream, using framer to mark one message
// off from the next. Call Run in a goroutine of its own to start reading.
//
// onNotify and onRequest may both be nil, in which case notifications are
// dropped and the peer's requests are answered with "method not found".
//
//	conn := jsonrpc.NewConn(stream, lsp.Framing{}, client.onNotify, client.onRequest)
//	go conn.Run()
func NewConn(stream io.ReadWriteCloser, framer Framer, onNotify NotificationFunc, onRequest RequestFunc) *Conn {
	return &Conn{
		stream:    stream,
		reader:    bufio.NewReader(stream),
		framer:    framer,
		onNotify:  onNotify,
		onRequest: onRequest,
		pending:   map[int64]chan *message{},
	}
}

// Run reads messages until the stream ends or the connection is closed, then
// fails every call still waiting for an answer.
//
// It returns nil for a clean end of stream and the read error otherwise.
func (c *Conn) Run() error {
	err := c.readLoop()
	c.failPending()
	return err
}

// readLoop dispatches every message that arrives.
func (c *Conn) readLoop() error {
	for {
		body, err := c.framer.ReadMessage(c.reader)
		if err != nil {
			if errors.Is(err, io.EOF) || c.isClosed() {
				return nil
			}
			return err
		}

		var msg message
		if err := json.Unmarshal(body, &msg); err != nil {
			return fmt.Errorf("jsonrpc: unreadable message: %w", err)
		}
		c.dispatch(&msg)
	}
}

// dispatch sends a message where it belongs: to the call that is waiting for
// it, to the notification handler, or to the request handler.
//
// A notification and a request differ only by the id: the peer expects an
// answer to exactly the messages that carry one.
func (c *Conn) dispatch(msg *message) {
	switch {
	case msg.Method != "" && len(msg.ID) == 0:
		if c.onNotify != nil {
			c.onNotify(msg.Method, msg.Params)
		}
	case msg.Method != "":
		c.answer(msg)
	default:
		c.deliver(msg)
	}
}

// deliver hands a response to the call that is waiting for it.
//
// The id is looked up among the ids *this* connection sent. A peer numbering
// its own requests from zero — which the Agent Client Protocol's agents do —
// therefore cannot collide with a call in flight here.
func (c *Conn) deliver(msg *message) {
	var id int64
	if err := json.Unmarshal(msg.ID, &id); err != nil {
		return // a response to an id we never sent; nothing to do with it
	}

	c.mu.Lock()
	waiting, ok := c.pending[id]
	delete(c.pending, id)
	c.mu.Unlock()

	if ok {
		waiting <- msg
	}
}

// answer hands a request the peer made to the handler, which replies now or
// later.
func (c *Conn) answer(msg *message) {
	request := &Request{Method: msg.Method, Params: msg.Params, conn: c, id: msg.ID}

	if c.onRequest == nil {
		request.Reply(nil, &ResponseError{Code: CodeMethodNotFound, Message: msg.Method})
		return
	}
	c.onRequest(request)
}

// Call sends a request and waits for its answer, decoding it into result.
//
// Passing a nil result discards the answer. The call ends early if ctx does,
// which is what stops a slow request from freezing the caller.
//
//	var items CompletionList
//	err := conn.Call(ctx, "textDocument/completion", params, &items)
func (c *Conn) Call(ctx context.Context, method string, params, result any) error {
	id, waiting, err := c.register()
	if err != nil {
		return err
	}
	defer c.forget(id)

	if err := c.sendRequest(id, method, params); err != nil {
		return err
	}

	select {
	case <-ctx.Done():
		return ctx.Err()
	case reply := <-waiting:
		return decodeReply(reply, result)
	}
}

// Notify sends a notification, which expects no answer.
//
//	err := conn.Notify("session/cancel", map[string]string{"sessionId": id})
func (c *Conn) Notify(method string, params any) error {
	encodedParams, err := encodeParams(params)
	if err != nil {
		return err
	}
	return c.write(&message{JSONRPC: Version, Method: method, Params: encodedParams})
}

// sendRequest encodes and writes one request.
func (c *Conn) sendRequest(id int64, method string, params any) error {
	encodedID, err := json.Marshal(id)
	if err != nil {
		return fmt.Errorf("jsonrpc: encoding the request id: %w", err)
	}
	encodedParams, err := encodeParams(params)
	if err != nil {
		return err
	}
	return c.write(&message{JSONRPC: Version, ID: encodedID, Method: method, Params: encodedParams})
}

// register allocates a request id and the channel its answer will arrive on.
func (c *Conn) register() (int64, chan *message, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.closed {
		return 0, nil, ErrClosed
	}
	c.nextID++
	// The channel is buffered so that the read loop never blocks handing over
	// an answer whose caller has already given up.
	waiting := make(chan *message, 1)
	c.pending[c.nextID] = waiting
	return c.nextID, waiting, nil
}

// forget drops a pending request, whether it was answered or abandoned.
func (c *Conn) forget(id int64) {
	c.mu.Lock()
	delete(c.pending, id)
	c.mu.Unlock()
}

// write frames and sends one message.
func (c *Conn) write(msg *message) error {
	body, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("jsonrpc: encoding a message: %w", err)
	}

	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	if c.isClosed() {
		return ErrClosed
	}
	return c.framer.WriteMessage(c.stream, body)
}

// failPending wakes every waiting call once the connection has ended.
//
// Requests the *peer* sent are not answered here. There is no answer to give,
// and the peer whose connection has just ended is not waiting for one.
func (c *Conn) failPending() {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.closed = true
	for id, waiting := range c.pending {
		close(waiting)
		delete(c.pending, id)
	}
}

// isClosed reports whether the connection has been closed.
func (c *Conn) isClosed() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.closed
}

// Close shuts the connection down. Calls still waiting fail with ErrClosed.
func (c *Conn) Close() error {
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return nil
	}
	c.closed = true
	c.mu.Unlock()

	return c.stream.Close()
}