// Package stdio is the Novox Mesh SDK for Go (novox/hq ADR 0188, ADR 0193).
//
// A tools bundle is a process the node's runtime launches and speaks MCP over stdio to:
// `initialize`, `tools/list` once, `tools/call` per call, newline-framed JSON-RPC on stdin and
// stdout. This package is the loop, so a bundle writes its tools and nothing else. A tool named
// `<seat>.<verb>` is the module's implementation of that seat's verb; any other name is the
// module's own.
//
// No transport here: the bus, the subjects and the memberships are the runtime's (ADR 0039, 0175).
// What a bundle asks of the bus — emitting an event — it asks of the runtime over the same channel,
// and the runtime does it as the module (ADR 0193). stdout is the protocol; log to stderr.
package stdio

import (
	"bufio"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Protocol is the MCP version this speaks; the runtime says the same.
const Protocol = "2025-03-26"

// ServedAs is the word the runtime sets for a bundle it launches: the module it serves it as
// (ADR 0193). A bundle served under that name lists its own tools by their names.
const ServedAs = "MESH_SERVED_MODULE"

// Tool is one tool a bundle serves.
type Tool struct {
	Name        string
	Description string
	// Input is the JSON schema of the arguments; nil is a tool that takes nothing.
	Input map[string]any
	Run   func(args map[string]any) (any, error)
}

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

// channel is the one conversation with the runtime: what this bundle writes, serialised, and what
// it asked and waits for.
type channel struct {
	mu        sync.Mutex
	enc       *json.Encoder
	next      int
	waiting   map[string]chan message
	listeners []listener
	// watches are this bundle's state watches by the id each was asked with (ADR 0201).
	watches map[string]func(StateChange) error
}

type listener struct {
	pattern string
	handle  func(Envelope) error
}

// Envelope is one event as the mesh carries it: its type, the machine it came from, its body as JSON,
// and the ADR 0042 headers (x-event-id, x-source, x-node, x-time …).
type Envelope struct {
	Key     string            `json:"key"`
	Node    string            `json:"node"`
	Body    json.RawMessage   `json:"body"`
	Headers map[string]string `json:"headers,omitempty"`
}

func (c *channel) say(m map[string]any) {
	m["jsonrpc"] = "2.0"
	c.mu.Lock()
	defer c.mu.Unlock()
	_ = c.enc.Encode(m)
}

// ask sends the runtime a request and waits for its answer.
func (c *channel) ask(method string, params any, timeout time.Duration) (json.RawMessage, error) {
	return c.askAs(c.newID(), method, params, timeout)
}

func (c *channel) newID() string {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.next++
	return "bundle-" + strconv.Itoa(c.next)
}

// askAs is ask under an id the caller chose, for a request whose id something is registered under
// before it is sent.
func (c *channel) askAs(id, method string, params any, timeout time.Duration) (json.RawMessage, error) {
	c.mu.Lock()
	answer := make(chan message, 1)
	c.waiting[id] = answer
	c.mu.Unlock()
	c.say(map[string]any{"id": id, "method": method, "params": params})
	select {
	case m := <-answer:
		if m.Error != nil {
			return nil, errors.New(m.Error.Message)
		}
		return m.Result, nil
	case <-time.After(timeout):
		c.mu.Lock()
		delete(c.waiting, id)
		c.mu.Unlock()
		return nil, fmt.Errorf("the runtime did not answer %s in %s", method, timeout)
	}
}

var (
	current   *channel
	currentMu sync.Mutex
)

// Serve answers the runtime on stdin and stdout until stdin closes. An empty name serves as the
// module the runtime named in MESH_SERVED_MODULE.
func Serve(name string, tools []Tool) error {
	return ServeOn(name, tools, os.Stdin, os.Stdout)
}

// ServeOn is Serve over any pair of streams, for a test.
func ServeOn(name string, tools []Tool, in io.Reader, out io.Writer) error {
	if name == "" {
		name = os.Getenv(ServedAs)
	}
	if name == "" {
		name = "bundle"
	}
	byName := map[string]Tool{}
	for _, t := range tools {
		byName[t.Name] = t
	}
	c := &channel{enc: json.NewEncoder(out), waiting: map[string]chan message{}, watches: map[string]func(StateChange) error{}}
	currentMu.Lock()
	current = c
	currentMu.Unlock()
	scanner := bufio.NewScanner(in)
	scanner.Buffer(make([]byte, 1<<20), 1<<24)
	for scanner.Scan() {
		var m message
		if err := json.Unmarshal(scanner.Bytes(), &m); err != nil {
			continue
		}
		// The runtime's answer to something this bundle asked.
		if m.Method == "" && m.ID != nil {
			var key string
			if json.Unmarshal(*m.ID, &key) == nil {
				c.mu.Lock()
				answer, ok := c.waiting[key]
				delete(c.waiting, key)
				c.mu.Unlock()
				if ok {
					answer <- m
				}
			}
			continue
		}
		id := any(nil)
		if m.ID != nil {
			id = *m.ID
		}
		switch m.Method {
		case "initialize":
			c.say(map[string]any{"id": id, "result": map[string]any{"protocolVersion": Protocol,
				"capabilities": map[string]any{"tools": map[string]any{}},
				"serverInfo":   map[string]any{"name": name, "version": "1"}}})
		case "notifications/initialized":
		case "ping":
			if id != nil {
				c.say(map[string]any{"id": id, "result": map[string]any{}})
			}
		case "tools/list":
			listed := make([]map[string]any, 0, len(tools))
			for _, t := range tools {
				schema := t.Input
				if schema == nil {
					schema = map[string]any{"type": "object", "properties": map[string]any{}}
				} else if _, isSchema := schema["type"]; !isSchema {
					schema = map[string]any{"type": "object", "properties": schema}
				}
				listed = append(listed, map[string]any{"name": t.Name, "description": t.Description, "inputSchema": schema})
			}
			c.say(map[string]any{"id": id, "result": map[string]any{"tools": listed}})
		case "tools/call":
			var p struct {
				Name      string         `json:"name"`
				Arguments map[string]any `json:"arguments"`
			}
			_ = json.Unmarshal(m.Params, &p)
			t, ok := byName[p.Name]
			if !ok {
				c.say(map[string]any{"id": id, "error": map[string]any{"code": -32602, "message": fmt.Sprintf("%s has no tool %s", name, p.Name)}})
				continue
			}
			// Beside the loop, never in it: a tool that emits waits for the runtime's answer, which
			// arrives on the line this loop reads next — and calls do not queue behind a slow one.
			go func(id any, t Tool, args map[string]any) {
				if args == nil {
					args = map[string]any{}
				}
				result, err := t.Run(args)
				if err != nil {
					c.say(map[string]any{"id": id, "result": map[string]any{"content": []map[string]any{{"type": "text", "text": err.Error()}}, "isError": true}})
					return
				}
				text, _ := json.Marshal(result)
				c.say(map[string]any{"id": id, "result": map[string]any{"content": []map[string]any{{"type": "text", "text": string(text)}}}})
			}(id, t, p.Arguments)
		case "mesh/event":
			// One event the runtime read from this module's consumer (ADR 0198), answered once every
			// handler it is for has taken it; an error is answered as one, so the runtime does not
			// acknowledge it and the bus delivers it again. Beside the loop, like a tool call.
			var p struct {
				Envelope *Envelope `json:"envelope"`
			}
			if json.Unmarshal(m.Params, &p) != nil || p.Envelope == nil || p.Envelope.Key == "" {
				c.say(map[string]any{"id": id, "error": map[string]any{"code": -32602, "message": "mesh/event carries no envelope"}})
				continue
			}
			c.mu.Lock()
			matching := []listener{}
			for _, l := range c.listeners {
				if TopicMatches(l.pattern, p.Envelope.Key) {
					matching = append(matching, l)
				}
			}
			c.mu.Unlock()
			go func(id any, env Envelope, ls []listener) {
				for _, l := range ls {
					if err := l.handle(env); err != nil {
						c.say(map[string]any{"id": id, "error": map[string]any{"code": -32000, "message": err.Error()}})
						return
					}
				}
				c.say(map[string]any{"id": id, "result": map[string]any{}})
			}(id, *p.Envelope, matching)
		case "mesh/state":
			// One change to a state this bundle watches (ADR 0201), named by the watch's id, answered
			// once its handler has taken it. Beside the loop, like a tool call.
			var p struct {
				Watch  string       `json:"watch"`
				Change *StateChange `json:"change"`
			}
			_ = json.Unmarshal(m.Params, &p)
			c.mu.Lock()
			handle := c.watches[p.Watch]
			c.mu.Unlock()
			if handle == nil || p.Change == nil {
				c.say(map[string]any{"id": id, "error": map[string]any{"code": -32602, "message": "mesh/state names no watch of this bundle"}})
				continue
			}
			go func(id any, change StateChange) {
				if err := handle(change); err != nil {
					c.say(map[string]any{"id": id, "error": map[string]any{"code": -32000, "message": err.Error()}})
					return
				}
				c.say(map[string]any{"id": id, "result": map[string]any{}})
			}(id, *p.Change)
		default:
			if id != nil {
				c.say(map[string]any{"id": id, "error": map[string]any{"code": -32601, "message": "no " + m.Method}})
			}
		}
	}
	return scanner.Err()
}

func live() (*channel, error) {
	currentMu.Lock()
	defer currentMu.Unlock()
	if current == nil {
		return nil, errors.New("before Serve: a bundle reaches the bus through the runtime that launched it")
	}
	return current, nil
}

// Subscribe asks the runtime for this module's events whose type matches the pattern (`*` one
// segment, `#` any), and hands each to handle (ADR 0198). The runtime acknowledges an event on the bus
// only after every matching handler returned nil; an error has it delivered again. Call it from a
// goroutine started before Serve, or from a tool, once Serve runs.
func Subscribe(pattern string, handle func(Envelope) error) error {
	c, err := live()
	if err != nil {
		return err
	}
	c.mu.Lock()
	c.listeners = append(c.listeners, listener{pattern: pattern, handle: handle})
	c.mu.Unlock()
	_, err = c.ask("mesh/subscribe", map[string]any{"pattern": pattern}, 30*time.Second)
	return err
}

// Ask calls a tool on the mesh as this bundle's module, through the runtime (ADR 0198), and returns
// its result as JSON. The key is `<module>.<tool>` as the runtime calls it.
func Ask(key string, body any) (json.RawMessage, error) {
	c, err := live()
	if err != nil {
		return nil, err
	}
	return c.ask("mesh/ask", map[string]any{"key": key, "body": body}, 60*time.Second)
}

// TopicMatches says whether an event's type matches a pattern: `*` is one dot-separated segment, `#`
// (or `**`) any number, as the TypeScript SDK matches.
func TopicMatches(pattern, key string) bool {
	return matchFrom(strings.Split(pattern, "."), 0, strings.Split(key, "."), 0)
}

func matchFrom(p []string, pi int, k []string, ki int) bool {
	if pi == len(p) {
		return ki == len(k)
	}
	if p[pi] == "#" || p[pi] == "**" {
		for skip := ki; skip <= len(k); skip++ {
			if matchFrom(p, pi+1, k, skip) {
				return true
			}
		}
		return false
	}
	if ki == len(k) || (p[pi] != "*" && p[pi] != k[ki]) {
		return false
	}
	return matchFrom(p, pi+1, k, ki+1)
}

// EmitOptions is the provenance an emitter may attach (ADR 0042).
type EmitOptions struct {
	CausationID string
	Schema      string
}

// Emit asks the runtime to publish an event as this bundle's module, and returns once the bus has
// accepted it (ADR 0193). The envelope and headers are the TypeScript SDK's exactly: the source and
// node come from MESH_MODULE and MESH_NODE, which the runtime sets.
func Emit(eventType string, body any, opts ...EmitOptions) error {
	currentMu.Lock()
	c := current
	currentMu.Unlock()
	if c == nil {
		return errors.New("emit before Serve: a bundle emits through the runtime that launched it")
	}
	source, node := os.Getenv("MESH_MODULE"), os.Getenv("MESH_NODE")
	if source == "" {
		source = "unknown"
	}
	if node == "" {
		node = "unknown"
	}
	headers := map[string]string{
		"x-event-id": eventID(), "x-source": source, "x-node": node,
		"x-time": time.Now().UTC().Format(time.RFC3339Nano), "content-type": "application/json",
	}
	for _, o := range opts {
		if o.CausationID != "" {
			headers["x-causation-id"] = o.CausationID
		}
		if o.Schema != "" {
			headers["x-schema"] = o.Schema
		}
	}
	_, err := c.ask("mesh/publish", map[string]any{"key": eventType, "node": node, "body": body, "headers": headers}, 30*time.Second)
	return err
}

// eventID is a random UUID, version 4, as the TypeScript SDK's randomUUID.
func eventID() string {
	var b [16]byte
	_, _ = rand.Read(b[:])
	b[6] = (b[6] & 0x0f) | 0x40
	b[8] = (b[8] & 0x3f) | 0x80
	h := hex.EncodeToString(b[:])
	return h[0:8] + "-" + h[8:12] + "-" + h[12:16] + "-" + h[16:20] + "-" + h[20:32]
}