// Package pluginsdk is the worker side of a Cletus plugin: the message
// protocol, chat tools, screen and service handlers, and calls to the host.
// See docs/specification.md and docs/host-api.md.
package pluginsdk

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

// APIVersion is the worker protocol version the host expects at handshake.
const APIVersion = 1

// MaxMessageBytes is the largest message (one line) either side may send.
const MaxMessageBytes = 64 << 20

// Func handles one screen (ui.*) request. params is the JSON array of the
// screen's arguments; decode it with Args.
type Func func(ctx context.Context, h *Host, params json.RawMessage) (any, error)

// ServiceFunc handles one service method another plugin calls.
type ServiceFunc func(ctx context.Context, h *Host, call ServiceCall) (any, error)

// ServiceCall is a service request the host routed from another plugin.
type ServiceCall struct {
	Caller   string          `json:"caller"`
	Mutating bool            `json:"mutating"`
	Args     json.RawMessage `json:"args"`
}

// Handlers is everything a worker answers.
type Handlers struct {
	// Tools are the chat tools. Each must also be listed in plugin.json.
	Tools []Tool
	// UI maps a screen method name ("check", answered as ui.check) to its handler.
	UI map[string]Func
	// Services maps "<service>.<method>" ("gauge.check") to its handler.
	// Each must be declared under "provides" in plugin.json.
	Services map[string]ServiceFunc
	// Busy reports whether work is running that must not be interrupted.
	// Nil means never busy.
	Busy func() bool
}

// ID returns the plugin id the host started the worker as.
func ID() string { return os.Getenv("CLETUS_PLUGIN_ID") }

// DataDir returns the worker's data folder for the signed-in account. It is
// also the worker's working directory. Empty when not started by a host.
func DataDir() string { return os.Getenv("CLETUS_PLUGIN_DATA") }

// Serve runs the worker on standard input and output until the host shuts
// it down or closes the pipe. Standard output belongs to the protocol, so
// Serve points os.Stdout at standard error: a stray fmt.Println lands in
// the worker log instead of breaking the connection.
func Serve(h Handlers) error {
	out := os.Stdout
	os.Stdout = os.Stderr
	return ServeIO(context.Background(), os.Stdin, out, h)
}

type message struct {
	ID       string          `json:"id,omitempty"`
	ParentID string          `json:"parent_id,omitempty"`
	Method   string          `json:"method,omitempty"`
	Params   json.RawMessage `json:"params,omitempty"`
	Result   json.RawMessage `json:"result,omitempty"`
	Error    *wireError      `json:"error,omitempty"`
}

type wireError struct {
	Message string `json:"message"`
}

type server struct {
	handlers Handlers
	tools    map[string]Tool
	out      io.Writer
	writeMu  sync.Mutex

	mu      sync.Mutex
	active  map[string]context.CancelFunc
	pending map[string]chan message
	next    int
	wg      sync.WaitGroup
}

// ServeIO runs the worker on the given streams. It returns when ctx ends,
// the input closes, or the host sends plugin.shutdown.
func ServeIO(ctx context.Context, in io.Reader, out io.Writer, h Handlers) error {
	s := &server{handlers: h, tools: map[string]Tool{}, out: out,
		active: map[string]context.CancelFunc{}, pending: map[string]chan message{}}
	for _, t := range h.Tools {
		if _, dup := s.tools[t.Name]; dup {
			return fmt.Errorf("tool %q is registered twice", t.Name)
		}
		s.tools[t.Name] = t
	}
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	defer s.wg.Wait()
	defer s.cancelAll()

	lines := make(chan []byte)
	readErr := make(chan error, 1)
	go func() {
		reader := bufio.NewReaderSize(in, 64<<10)
		for {
			line, err := readLine(reader)
			if len(line) > 0 {
				select {
				case lines <- line:
				case <-ctx.Done():
					return
				}
			}
			if err != nil {
				readErr <- err
				return
			}
		}
	}()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case err := <-readErr:
			if errors.Is(err, io.EOF) {
				return nil
			}
			return err
		case line := <-lines:
			var m message
			if err := json.Unmarshal(line, &m); err != nil {
				return fmt.Errorf("invalid message from host: %w", err)
			}
			if stop := s.dispatch(ctx, m); stop {
				return nil
			}
		}
	}
}

func readLine(r *bufio.Reader) ([]byte, error) {
	var line []byte
	for {
		chunk, err := r.ReadSlice('\n')
		line = append(line, chunk...)
		if len(line) > MaxMessageBytes+1 {
			return nil, errors.New("message from host is too large")
		}
		if err == bufio.ErrBufferFull {
			continue
		}
		return []byte(strings.TrimSpace(string(line))), err
	}
}

// dispatch handles one incoming message and reports whether to stop.
func (s *server) dispatch(ctx context.Context, m message) bool {
	switch {
	case m.Method == "$/cancel":
		var id string
		if json.Unmarshal(m.Params, &id) == nil {
			s.mu.Lock()
			if cancel := s.active[id]; cancel != nil {
				cancel()
			}
			s.mu.Unlock()
		}
	case m.Method == "plugin.shutdown":
		s.reply(m.ID, map[string]bool{"ok": true}, nil)
		return true
	case m.Method != "":
		if m.ID == "" {
			return false // a notification this worker doesn't use
		}
		reqCtx, cancel := context.WithCancel(ctx)
		s.mu.Lock()
		s.active[m.ID] = cancel
		s.mu.Unlock()
		s.wg.Add(1)
		go func() {
			defer s.wg.Done()
			defer func() {
				s.mu.Lock()
				delete(s.active, m.ID)
				s.mu.Unlock()
				cancel()
			}()
			result, err := s.handle(reqCtx, m)
			s.reply(m.ID, result, err)
		}()
	case m.ID != "":
		s.mu.Lock()
		ch := s.pending[m.ID]
		delete(s.pending, m.ID)
		s.mu.Unlock()
		if ch != nil {
			ch <- m
		}
	}
	return false
}

func (s *server) handle(ctx context.Context, m message) (result any, err error) {
	defer func() {
		if r := recover(); r != nil {
			result, err = nil, fmt.Errorf("the plugin failed while handling %s: %v", m.Method, r)
		}
	}()
	h := &Host{call: func(ctx context.Context, method string, params any) (json.RawMessage, error) {
		return s.callHost(ctx, m.ID, method, params)
	}}
	switch {
	case m.Method == "plugin.handshake":
		return map[string]int{"api_version": APIVersion}, nil
	case m.Method == "plugin.busy":
		return s.handlers.Busy != nil && s.handlers.Busy(), nil
	case m.Method == "tools.list":
		return s.listTools(), nil
	case m.Method == "tools.invoke":
		return s.invokeTool(ctx, h, m.Params)
	case strings.HasPrefix(m.Method, "ui."):
		fn := s.handlers.UI[strings.TrimPrefix(m.Method, "ui.")]
		if fn == nil {
			return nil, fmt.Errorf("unknown method %s", m.Method)
		}
		return fn(ctx, h, m.Params)
	case strings.HasPrefix(m.Method, "svc."):
		fn := s.handlers.Services[strings.TrimPrefix(m.Method, "svc.")]
		if fn == nil {
			return nil, fmt.Errorf("unknown method %s", m.Method)
		}
		var call ServiceCall
		if err := json.Unmarshal(m.Params, &call); err != nil {
			return nil, fmt.Errorf("invalid service request: %w", err)
		}
		return fn(ctx, h, call)
	}
	return nil, fmt.Errorf("unknown method %s", m.Method)
}

func (s *server) reply(id string, result any, err error) {
	if id == "" {
		return
	}
	m := message{ID: id}
	if err != nil {
		m.Error = &wireError{Message: err.Error()}
	} else {
		raw, marshalErr := json.Marshal(result)
		if marshalErr != nil {
			m.Error = &wireError{Message: "the plugin returned a result that is not JSON: " + marshalErr.Error()}
		} else {
			m.Result = raw
		}
	}
	s.send(m)
}

func (s *server) send(m message) {
	data, err := json.Marshal(m)
	if err != nil || len(data) > MaxMessageBytes {
		if m.ID != "" && m.Method == "" {
			data, _ = json.Marshal(message{ID: m.ID, Error: &wireError{Message: "the plugin's reply is too large"}})
		} else {
			return
		}
	}
	s.writeMu.Lock()
	defer s.writeMu.Unlock()
	_, _ = s.out.Write(append(data, '\n'))
}

func (s *server) callHost(ctx context.Context, parent, method string, params any) (json.RawMessage, error) {
	raw, err := json.Marshal(params)
	if err != nil {
		return nil, fmt.Errorf("%s: %w", method, err)
	}
	ch := make(chan message, 1)
	s.mu.Lock()
	s.next++
	id := fmt.Sprintf("w%d", s.next)
	s.pending[id] = ch
	s.mu.Unlock()
	s.send(message{ID: id, ParentID: parent, Method: method, Params: raw})
	select {
	case reply := <-ch:
		if reply.Error != nil {
			return nil, errors.New(reply.Error.Message)
		}
		return reply.Result, nil
	case <-ctx.Done():
		s.mu.Lock()
		delete(s.pending, id)
		s.mu.Unlock()
		cancel, _ := json.Marshal(id)
		s.send(message{Method: "$/cancel", Params: cancel})
		return nil, ctx.Err()
	}
}

func (s *server) cancelAll() {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, cancel := range s.active {
		cancel()
	}
}

// Args decodes a screen request's argument array into targets, in order.
// Missing trailing arguments leave their targets unchanged.
func Args(params json.RawMessage, targets ...any) error {
	var items []json.RawMessage
	if len(params) > 0 && string(params) != "null" {
		if err := json.Unmarshal(params, &items); err != nil {
			return fmt.Errorf("screen arguments must be a JSON array: %w", err)
		}
	}
	for i, target := range targets {
		if i >= len(items) {
			break
		}
		if err := json.Unmarshal(items[i], target); err != nil {
			return fmt.Errorf("argument %d: %w", i+1, err)
		}
	}
	return nil
}