lore-engine-poc/workers/mcp-server/main.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/http"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/neo4j/neo4j-go-driver/v5/neo4j"
)

var httpClient = &http.Client{Timeout: 30 * time.Second}

// ── Context Logger ───────────────────────────────────────────────────────────

type contextKey string
const loggerKey contextKey = "logger"

func contextWithLogger(ctx context.Context, l *slog.Logger) context.Context {
	return context.WithValue(ctx, loggerKey, l)
}

func loggerFromContext(ctx context.Context) *slog.Logger {
	if l, ok := ctx.Value(loggerKey).(*slog.Logger); ok {
		return l
	}
	return slog.Default()
}

// ── Config ────────────────────────────────────────────────────────────────────

type Config struct {
	Neo4jURL   string
	Neo4jUser  string
	Neo4jPass  string
	EmbedURL   string
	EmbedModel string
	Port       string
}

func configFromEnv() Config {
	return Config{
		Neo4jURL:   getEnv("NEO4J_URL", "bolt://neo4j:7687"),
		Neo4jUser:  getEnv("NEO4J_USER", "neo4j"),
		Neo4jPass:  getEnv("NEO4J_PASSWORD", "changeme"),
		EmbedURL:   getEnv("EMBED_URL", "http://ollama-gpu:11434"),
		EmbedModel: getEnv("EMBED_MODEL", "nomic-embed-text"),
		Port:       getEnv("MCP_PORT", "9000"),
	}
}

// ── OpenAI-compatible embed ───────────────────────────────────────────────────

type embedReq struct {
	Model string `json:"model"`
	Input string `json:"input"`
}
type embedRespItem struct {
	Embedding []float32 `json:"embedding"`
}
type embedResp struct {
	Data []embedRespItem `json:"data"`
}

func embed(ctx context.Context, cfg Config, text string) ([]float32, error) {
	body, err := json.Marshal(embedReq{Model: cfg.EmbedModel, Input: text})
	if err != nil {
		return nil, fmt.Errorf("marshal embed request: %w", err)
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		cfg.EmbedURL+"/v1/embeddings", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("create embed HTTP request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("execute embed HTTP request: %w", err)
	}
	defer resp.Body.Close()
	var er embedResp
	if err := json.NewDecoder(resp.Body).Decode(&er); err != nil {
		return nil, fmt.Errorf("decode embed response: %w", err)
	}
	if len(er.Data) == 0 {
		return nil, fmt.Errorf("empty embedding response")
	}
	return er.Data[0].Embedding, nil
}

// ── MCP JSON-RPC types ────────────────────────────────────────────────────────
// Implements the MCP 2024-11-05 specification over HTTP+SSE transport.
// Spec: https://spec.modelcontextprotocol.io/specification/

type JSONRPCRequest struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      any             `json:"id"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params,omitempty"`
}

type JSONRPCResponse struct {
	JSONRPC string `json:"jsonrpc"`
	ID      any    `json:"id"`
	Result  any    `json:"result,omitempty"`
	Error   *RPCError `json:"error,omitempty"`
}

type RPCError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

func okResp(id any, result any) JSONRPCResponse {
	return JSONRPCResponse{JSONRPC: "2.0", ID: id, Result: result}
}

func errResp(id any, code int, msg string) JSONRPCResponse {
	return JSONRPCResponse{JSONRPC: "2.0", ID: id, Error: &RPCError{Code: code, Message: msg}}
}

// ── MCP Tool definitions ──────────────────────────────────────────────────────

type MCPTool struct {
	Name        string         `json:"name"`
	Description string         `json:"description"`
	InputSchema map[string]any `json:"inputSchema"`
}

var mcpTools = []MCPTool{
	{
		Name:        "semantic_search",
		Description: "Find messages and chunks semantically similar to a query using vector similarity over the knowledge graph",
		InputSchema: map[string]any{
			"type": "object",
			"properties": map[string]any{
				"query": map[string]any{"type": "string", "description": "Natural language search query"},
				"limit": map[string]any{"type": "integer", "description": "Max results to return (default 5)"},
			},
			"required": []string{"query"},
		},
	},
	{
		Name:        "graph_traverse",
		Description: "Traverse the knowledge graph from a named entity to find related entities and messages",
		InputSchema: map[string]any{
			"type": "object",
			"properties": map[string]any{
				"entity": map[string]any{"type": "string", "description": "Entity name to start traversal from"},
				"depth":  map[string]any{"type": "integer", "description": "Traversal depth 1-3 (default 2)"},
			},
			"required": []string{"entity"},
		},
	},
	{
		Name:        "get_context",
		Description: "Get full context for a specific message including its chunks and all related entities",
		InputSchema: map[string]any{
			"type": "object",
			"properties": map[string]any{
				"message_id": map[string]any{"type": "string", "description": "Message ID to retrieve context for"},
			},
			"required": []string{"message_id"},
		},
	},
	{
		Name:        "get_person_profile",
		Description: "Get topics, interests, and message history associated with a named person",
		InputSchema: map[string]any{
			"type": "object",
			"properties": map[string]any{
				"name": map[string]any{"type": "string", "description": "Person's name"},
			},
			"required": []string{"name"},
		},
	},
	{
		Name:        "query_as_npc",
		Description: "Query the knowledge graph from a specific NPC's perspective, scoped to only what they have personally witnessed. Returns semantic search results and encounter graph context filtered to the NPC's knowledge horizon.",
		InputSchema: map[string]any{
			"type": "object",
			"properties": map[string]any{
				"npc_name": map[string]any{"type": "string", "description": "The NPC's name (must match a Person node in the graph)"},
				"question": map[string]any{"type": "string", "description": "The question the NPC is trying to answer"},
				"limit":    map[string]any{"type": "integer", "description": "Max chunk results to return (default 5)"},
			},
			"required": []string{"npc_name", "question"},
		},
	},
	{
		Name:        "log_encounter",
		Description: "Log a D&D encounter directly to the knowledge graph. Creates an Encounter node with WITNESSED edges for each participant. Call this after each NPC conversation so the NPC remembers it next time.",
		InputSchema: map[string]any{
			"type": "object",
			"properties": map[string]any{
				"title":        map[string]any{"type": "string", "description": "Short title for the encounter"},
				"participants": map[string]any{"type": "string", "description": "Comma-separated list of participant names"},
				"summary":      map[string]any{"type": "string", "description": "Brief summary of what happened or was discussed"},
				"location":     map[string]any{"type": "string", "description": "Location name where the encounter happened (optional)"},
				"type":         map[string]any{"type": "string", "description": "Encounter type: conversation, combat, discovery (default: conversation)"},
			},
			"required": []string{"title", "participants", "summary"},
		},
	},
	{
		Name:        "get_unresolved",
		Description: "List provisional entity nodes (lore_verified=false) — entities created from encounter data that have no matching lore document yet. Use this to identify gaps in the lore that need a document written, or aliases that should be added to an existing character's file.",
		InputSchema: map[string]any{
			"type": "object",
			"properties": map[string]any{
				"type":  map[string]any{"type": "string", "description": "Filter by entity type: Person, Location, Faction, etc. Omit for all."},
				"limit": map[string]any{"type": "integer", "description": "Max results (default 30)"},
			},
		},
	},
	{
		Name:        "get_contradictions",
		Description: "Return flagged contradictions — cases where two source documents make conflicting claims about the same entity (e.g., two lore files disagree on where a character was located). Use this to surface deception, unreliable narrators, or factual disputes hidden in the record.",
		InputSchema: map[string]any{
			"type": "object",
			"properties": map[string]any{
				"subject": map[string]any{"type": "string", "description": "Optional entity name to filter by (e.g., 'Silas Viper'). Omit to return all contradictions."},
				"limit":   map[string]any{"type": "integer", "description": "Max results to return (default 20)"},
			},
		},
	},
	{
		Name:        "list_encounters",
		Description: "List all past encounters stored in the campaign knowledge graph, ordered by recency",
		InputSchema: map[string]any{
			"type": "object",
			"properties": map[string]any{
				"limit": map[string]any{"type": "integer", "description": "Max encounters to return (default 10)"},
			},
		},
	},
	{
		Name:        "search_encounters",
		Description: "Search and filter past encounters by keyword, location, or participant name",
		InputSchema: map[string]any{
			"type": "object",
			"properties": map[string]any{
				"query":       map[string]any{"type": "string", "description": "Optional keyword search in titles/summaries"},
				"location":    map[string]any{"type": "string", "description": "Optional location name to filter by"},
				"participant": map[string]any{"type": "string", "description": "Optional participant name to filter by"},
				"limit":       map[string]any{"type": "integer", "description": "Max results to return (default 10)"},
			},
		},
	},
	{
		Name:        "get_encounter",
		Description: "Get complete details for a single campaign encounter by ID, including participants and featured entities",
		InputSchema: map[string]any{
			"type": "object",
			"properties": map[string]any{
				"id": map[string]any{"type": "string", "description": "The encounter ID"},
			},
			"required": []string{"id"},
		},
	},
}

// ── SSE session registry ──────────────────────────────────────────────────────
// Each GET /sse opens a persistent connection. Messages arrive via POST /message?sessionId=X
// and are routed to the right SSE stream via this registry.

type sseSession struct {
	ch     chan string
	ctx    context.Context
	cancel context.CancelFunc
}

type sessionRegistry struct {
	mu       sync.RWMutex
	sessions map[string]*sseSession
}

func newRegistry() *sessionRegistry {
	return &sessionRegistry{sessions: make(map[string]*sseSession)}
}

func (r *sessionRegistry) add(id string, s *sseSession) {
	r.mu.Lock()
	r.sessions[id] = s
	r.mu.Unlock()
}

func (r *sessionRegistry) get(id string) (*sseSession, bool) {
	r.mu.RLock()
	s, ok := r.sessions[id]
	r.mu.RUnlock()
	return s, ok
}

func (r *sessionRegistry) remove(id string) {
	r.mu.Lock()
	delete(r.sessions, id)
	r.mu.Unlock()
}

// ── Tool execution ────────────────────────────────────────────────────────────

func executeTool(ctx context.Context, cfg Config, driver neo4j.DriverWithContext,
	name string, params map[string]any) (any, error) {

	session := driver.NewSession(ctx, neo4j.SessionConfig{})
	defer session.Close(ctx)

	switch name {
	case "semantic_search":
		return handleSemanticSearch(ctx, cfg, session, params)
	case "graph_traverse":
		return handleGraphTraverse(ctx, session, params)
	case "get_context":
		return handleGetContext(ctx, session, params)
	case "get_person_profile":
		return handleGetPersonProfile(ctx, session, params)
	case "query_as_npc":
		return handleQueryAsNPC(ctx, cfg, session, params)
	case "log_encounter":
		return handleLogEncounter(ctx, session, params)
	case "get_unresolved":
		return handleGetUnresolved(ctx, session, params)
	case "get_contradictions":
		return handleGetContradictions(ctx, session, params)
	case "list_encounters":
		return handleListEncounters(ctx, session, params)
	case "search_encounters":
		return handleSearchEncounters(ctx, session, params)
	case "get_encounter":
		return handleGetEncounter(ctx, session, params)
	default:
		return nil, fmt.Errorf("unknown tool: %s", name)
	}
}

// ── MCP request dispatcher ────────────────────────────────────────────────────

func dispatch(ctx context.Context, cfg Config, driver neo4j.DriverWithContext,
	req JSONRPCRequest) JSONRPCResponse {

	switch req.Method {

	// ── Lifecycle ──────────────────────────────────────────────────────────────
	case "initialize":
		return okResp(req.ID, map[string]any{
			"protocolVersion": "2024-11-05",
			"capabilities": map[string]any{
				"tools": map[string]any{},
			},
			"serverInfo": map[string]any{
				"name":    "graphmcp",
				"version": "1.0.0",
			},
		})

	case "notifications/initialized":
		// Client confirms init — no response needed for notifications
		return JSONRPCResponse{}

	case "ping":
		return okResp(req.ID, map[string]any{})

	// ── Tool discovery ─────────────────────────────────────────────────────────
	case "tools/list":
		return okResp(req.ID, map[string]any{"tools": mcpTools})

	// ── Tool invocation ────────────────────────────────────────────────────────
	case "tools/call":
		var p struct {
			Name      string         `json:"name"`
			Arguments map[string]any `json:"arguments"`
		}
		if err := json.Unmarshal(req.Params, &p); err != nil {
			return errResp(req.ID, -32602, "invalid params: "+err.Error())
		}

		logger := loggerFromContext(ctx)
		logger.Info("Executing MCP tool", "tool", p.Name, "arguments", p.Arguments)

		result, err := executeTool(ctx, cfg, driver, p.Name, p.Arguments)
		if err != nil {
			logger.Error("MCP tool execution failed", "tool", p.Name, "err", err)
			return errResp(req.ID, -32603, err.Error())
		}

		logger.Info("MCP tool execution completed successfully", "tool", p.Name)

		// MCP tools/call result must be wrapped in content array
		resultJSON, err := json.Marshal(result)
		if err != nil {
			logger.Error("marshal MCP tool result failed", "tool", p.Name, "err", err)
			return errResp(req.ID, -32603, "marshal result: "+err.Error())
		}
		return okResp(req.ID, map[string]any{
			"content": []map[string]any{
				{
					"type": "text",
					"text": string(resultJSON),
				},
			},
		})

	default:
		return errResp(req.ID, -32601, "method not found: "+req.Method)
	}
}

// ── HTTP handlers ─────────────────────────────────────────────────────────────

func sseHandler(cfg Config, driver neo4j.DriverWithContext,
	reg *sessionRegistry) http.HandlerFunc {

	return func(w http.ResponseWriter, r *http.Request) {
		logger := loggerFromContext(r.Context())
		// SSE requires these headers
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
		w.Header().Set("Connection", "keep-alive")
		w.Header().Set("Access-Control-Allow-Origin", "*")

		flusher, ok := w.(http.Flusher)
		if !ok {
			http.Error(w, "streaming not supported", http.StatusInternalServerError)
			return
		}

		// Generate session ID and register
		sessionID := fmt.Sprintf("%d", time.Now().UnixNano())
		ctx, cancel := context.WithCancel(r.Context())
		sess := &sseSession{
			ch:     make(chan string, 32),
			ctx:    ctx,
			cancel: cancel,
		}
		reg.add(sessionID, sess)
		defer reg.remove(sessionID)
		defer cancel()

		// Send endpoint event — tells client where to POST messages
		// Format: event: endpoint\ndata: /message?sessionId=X\n\n
		fmt.Fprintf(w, "event: endpoint\ndata: /message?sessionId=%s\n\n", sessionID)
		flusher.Flush()

		logger.Info("SSE session opened", "sessionId", sessionID)

		// Keep-alive ticker — prevents proxy timeouts
		ticker := time.NewTicker(15 * time.Second)
		defer ticker.Stop()

		for {
			select {
			case msg := <-sess.ch:
				fmt.Fprintf(w, "data: %s\n\n", msg)
				flusher.Flush()
			case <-ticker.C:
				fmt.Fprintf(w, ": keepalive\n\n")
				flusher.Flush()
			case <-ctx.Done():
				logger.Info("SSE session closed", "sessionId", sessionID)
				return
			}
		}
	}
}

func messageHandler(cfg Config, driver neo4j.DriverWithContext,
	reg *sessionRegistry) http.HandlerFunc {

	return func(w http.ResponseWriter, r *http.Request) {
		logger := loggerFromContext(r.Context())
		if r.Method != http.MethodPost {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}

		sessionID := r.URL.Query().Get("sessionId")
		sess, ok := reg.get(sessionID)
		if !ok {
			http.Error(w, "session not found", http.StatusNotFound)
			return
		}

		var req JSONRPCRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			http.Error(w, "invalid JSON", http.StatusBadRequest)
			return
		}

		// Acknowledge immediately — response comes via SSE
		w.WriteHeader(http.StatusAccepted)

		// Dispatch in goroutine so POST returns fast
		go func() {
			resp := dispatch(sess.ctx, cfg, driver, req)

			// Notifications have no ID — don't send a response
			if resp.ID == nil && resp.Result == nil && resp.Error == nil {
				return
			}

			data, err := json.Marshal(resp)
			if err != nil {
				logger.Error("marshal response", "err", err)
				return
			}

			select {
			case sess.ch <- string(data):
			case <-sess.ctx.Done():
			}
		}()
	}
}

// ── Main ──────────────────────────────────────────────────────────────────────

func main() {
	cfg := configFromEnv()
	logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))
	slog.SetDefault(logger)

	ctx := context.Background()
	ctx = contextWithLogger(ctx, logger)

	driver, err := neo4j.NewDriverWithContext(cfg.Neo4jURL,
		neo4j.BasicAuth(cfg.Neo4jUser, cfg.Neo4jPass, ""))
	if err != nil {
		logger.Error("neo4j driver error", "err", err)
		os.Exit(1)
	}
	defer driver.Close(ctx)

	if err := driver.VerifyConnectivity(ctx); err != nil {
		logger.Error("neo4j connectivity check failed", "err", err)
		os.Exit(1)
	}

	reg := newRegistry()
	mux := http.NewServeMux()

	// ── Streamable HTTP — Open WebUI "MCP (Streamable HTTP)" type
	// Single POST /mcp endpoint, synchronous JSON-RPC response.
	// In Open WebUI Admin Settings → External Tools:
	//   Type: MCP (Streamable HTTP)
	//   URL:  http://mcp-server:9000/mcp
	mux.HandleFunc("/mcp", func(w http.ResponseWriter, r *http.Request) {
		reqCtx := contextWithLogger(r.Context(), logger)
		w.Header().Set("Content-Type", "application/json")
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")
		if r.Method == http.MethodOptions {
			w.WriteHeader(http.StatusOK)
			return
		}
		if r.Method != http.MethodPost {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		var req JSONRPCRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			w.WriteHeader(http.StatusBadRequest)
			if err := json.NewEncoder(w).Encode(errResp(nil, -32700, "parse error")); err != nil {
				logger.Error("failed to write parse error response", "err", err)
			}
			return
		}
		resp := dispatch(reqCtx, cfg, driver, req)
		if resp.ID == nil && resp.Result == nil && resp.Error == nil {
			w.WriteHeader(http.StatusAccepted)
			return
		}
		if err := json.NewEncoder(w).Encode(resp); err != nil {
			logger.Error("failed to encode response", "err", err)
		}
	})

	// ── SSE transport — kept for Claude Desktop / other SSE clients
	mux.HandleFunc("/sse", func(w http.ResponseWriter, r *http.Request) {
		reqCtx := contextWithLogger(r.Context(), logger)
		sseHandler(cfg, driver, reg)(w, r.WithContext(reqCtx))
	})
	mux.HandleFunc("/message", func(w http.ResponseWriter, r *http.Request) {
		reqCtx := contextWithLogger(r.Context(), logger)
		messageHandler(cfg, driver, reg)(w, r.WithContext(reqCtx))
	})

	// Health check
	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		if err := json.NewEncoder(w).Encode(map[string]string{"status": "ok"}); err != nil {
			logger.Error("failed to encode health status", "err", err)
		}
	})

	addr := ":" + cfg.Port
	logger.Info("mcp-server listening", "addr", addr,
		"streamable_http", addr+"/mcp",
		"sse", addr+"/sse")

	if err := http.ListenAndServe(addr, mux); err != nil {
		logger.Error("server error", "err", err)
		os.Exit(1)
	}
}

// ── Tool handlers (unchanged from original) ───────────────────────────────────

func handleSemanticSearch(ctx context.Context, cfg Config,
	session neo4j.SessionWithContext, params map[string]any) (any, error) {

	query, _ := params["query"].(string)
	limit := 5
	if l, ok := params["limit"].(float64); ok {
		limit = int(l)
	}

	vec, err := embed(ctx, cfg, query)
	if err != nil {
		return nil, fmt.Errorf("embed query: %w", err)
	}

	return session.ExecuteRead(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
		// Query message chunks
		res, err := tx.Run(ctx, `
			CALL db.index.vector.queryNodes('chunk_embeddings', $limit, $vec)
			YIELD node AS chunk, score
			MATCH (m:Message)-[:HAS_CHUNK]->(chunk)
			RETURN chunk.text AS text, score, m.id AS msgID,
			       m.author AS author, m.timestamp AS timestamp
			ORDER BY score DESC
		`, map[string]any{"limit": limit, "vec": vec})
		if err != nil {
			return nil, fmt.Errorf("query message vector index: %w", err)
		}
		var rows []map[string]any
		for res.Next(ctx) {
			record := res.Record()
			row := map[string]any{"source": "message"}
			for _, key := range record.Keys {
				if v, ok := record.Get(key); ok {
					row[key] = v
				}
			}
			rows = append(rows, row)
		}
		if err := res.Err(); err != nil {
			return nil, fmt.Errorf("message vector query iterate: %w", err)
		}

		// Query lore chunks — d.title maps to "author", d.id to "msgID" for schema compat
		loreRes, err := tx.Run(ctx, `
			CALL db.index.vector.queryNodes('lore_chunk_embeddings', $limit, $vec)
			YIELD node AS chunk, score
			MATCH (d:LoreDocument)-[:HAS_CHUNK]->(chunk)
			RETURN chunk.text AS text, score, d.id AS msgID,
			       d.title AS author, d.uploaded_at AS timestamp
			ORDER BY score DESC
		`, map[string]any{"limit": limit, "vec": vec})
		if err != nil {
			return nil, fmt.Errorf("query lore vector index: %w", err)
		}
		for loreRes.Next(ctx) {
			record := loreRes.Record()
			row := map[string]any{"source": "lore"}
			for _, key := range record.Keys {
				if v, ok := record.Get(key); ok {
					row[key] = v
				}
			}
			rows = append(rows, row)
		}
		if err := loreRes.Err(); err != nil {
			return nil, fmt.Errorf("lore vector query iterate: %w", err)
		}

		// Merge and re-rank by score descending, truncate to caller's limit
		sort.Slice(rows, func(i, j int) bool {
			si, _ := rows[i]["score"].(float64)
			sj, _ := rows[j]["score"].(float64)
			return si > sj
		})
		if len(rows) > limit {
			rows = rows[:limit]
		}

		return rows, nil
	})
}

func handleGraphTraverse(ctx context.Context,
	session neo4j.SessionWithContext, params map[string]any) (any, error) {

	entity, _ := params["entity"].(string)
	depth := 2
	if d, ok := params["depth"].(float64); ok {
		depth = int(d)
	}

	return session.ExecuteRead(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
		res, err := tx.Run(ctx, fmt.Sprintf(`
			MATCH path = (start {name: $entity})-[*1..%d]-(related)
			RETURN [node IN nodes(path) | coalesce(node.name, node.id, '')] AS nodePath,
			       [rel  IN relationships(path) | type(rel)] AS relPath
			LIMIT 20
		`, depth), map[string]any{"entity": entity})
		if err != nil {
			return nil, fmt.Errorf("run traverse query: %w", err)
		}
		var rows []map[string]any
		for res.Next(ctx) {
			record := res.Record()
			row := map[string]any{}
			for _, key := range record.Keys {
				if v, ok := record.Get(key); ok {
					row[key] = v
				}
			}
			rows = append(rows, row)
		}
		if err := res.Err(); err != nil {
			return nil, fmt.Errorf("traverse iterate: %w", err)
		}
		return rows, nil
	})
}

func handleGetContext(ctx context.Context,
	session neo4j.SessionWithContext, params map[string]any) (any, error) {

	msgID, _ := params["message_id"].(string)

	return session.ExecuteRead(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
		res, err := tx.Run(ctx, `
			MATCH (m:Message {id: $msgID})
			OPTIONAL MATCH (m)-[:HAS_CHUNK]->(c:Chunk)
			OPTIONAL MATCH (m)-[:MENTIONS]->(e)
			RETURN m.content AS content, m.author AS author,
			       m.timestamp AS timestamp,
			       collect(DISTINCT c.text) AS chunks,
			       collect(DISTINCT {name: e.name, type: labels(e)[0]}) AS entities
		`, map[string]any{"msgID": msgID})
		if err != nil {
			return nil, fmt.Errorf("run get context query: %w", err)
		}
		var rows []map[string]any
		for res.Next(ctx) {
			record := res.Record()
			row := map[string]any{}
			for _, key := range record.Keys {
				if v, ok := record.Get(key); ok {
					row[key] = v
				}
			}
			rows = append(rows, row)
		}
		if err := res.Err(); err != nil {
			return nil, fmt.Errorf("get context iterate: %w", err)
		}
		return rows, nil
	})
}

func handleGetPersonProfile(ctx context.Context,
	session neo4j.SessionWithContext, params map[string]any) (any, error) {

	name, _ := params["name"].(string)

	return session.ExecuteRead(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
		res, err := tx.Run(ctx, `
			MATCH (p:Person {name: $name})<-[:MENTIONS]-(m:Message)
			OPTIONAL MATCH (m)-[:MENTIONS]->(t:Topic)
			RETURN p.name AS person,
			       count(DISTINCT m) AS messageCount,
			       collect(DISTINCT t.name) AS topics,
			       collect(DISTINCT {id: m.id, content: m.content, ts: m.timestamp})[..10] AS recentMessages
		`, map[string]any{"name": name})
		if err != nil {
			return nil, fmt.Errorf("run person profile query: %w", err)
		}
		var rows []map[string]any
		for res.Next(ctx) {
			record := res.Record()
			row := map[string]any{}
			for _, key := range record.Keys {
				if v, ok := record.Get(key); ok {
					row[key] = v
				}
			}
			rows = append(rows, row)
		}
		if err := res.Err(); err != nil {
			return nil, fmt.Errorf("person profile iterate: %w", err)
		}
		return rows, nil
	})
}

// ── NPC tools ─────────────────────────────────────────────────────────────────

// handleLogEncounter writes an Encounter node + WITNESSED edges synchronously so
// query_as_npc can see the result on the very next call (no stream hop).
// Partial edge failures are collected in the "warnings" response field rather than
// being silently dropped, so callers know when the graph is incomplete.
func handleLogEncounter(ctx context.Context,
	session neo4j.SessionWithContext, params map[string]any) (any, error) {

	logger := loggerFromContext(ctx)
	title, _ := params["title"].(string)
	participants, _ := params["participants"].(string)
	summary, _ := params["summary"].(string)
	location, _ := params["location"].(string)
	encType, _ := params["type"].(string)
	if strings.TrimSpace(title) == "" {
		return nil, fmt.Errorf("title is required")
	}
	if encType == "" {
		encType = "conversation"
	}
	var writeWarnings []string

	timestamp := time.Now().UTC().Format(time.RFC3339)
	encID := fmt.Sprintf("mcp-%d", time.Now().UnixMilli())

	_, err := session.ExecuteWrite(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
		_, err := tx.Run(ctx, `
			MERGE (enc:Encounter {id: $id})
			  ON CREATE SET
			    enc.title         = $title,
			    enc.type          = $type,
			    enc.location_name = $location,
			    enc.timestamp     = $timestamp,
			    enc.summary       = $summary
		`, map[string]any{
			"id": encID, "title": title, "type": encType,
			"location": location, "timestamp": timestamp, "summary": summary,
		})
		return nil, err
	})
	if err != nil {
		return nil, fmt.Errorf("create encounter: %w", err)
	}

	if strings.TrimSpace(location) != "" {
		canonical := resolveCanonical(ctx, session, location)
		_, err := session.ExecuteWrite(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
			if canonical != "" {
				_, err := tx.Run(ctx, `
					MATCH (loc {name: $canonical})
					MATCH (enc:Encounter {id: $encID})
					MERGE (enc)-[:OCCURRED_AT]->(loc)
				`, map[string]any{"canonical": canonical, "encID": encID})
				return nil, err
			}
			_, err := tx.Run(ctx, `
				MERGE (loc:Location {name: $location})
				  ON CREATE SET loc.lore_verified = false, loc.source = "encounter"
				MATCH (enc:Encounter {id: $encID})
				MERGE (enc)-[:OCCURRED_AT]->(loc)
			`, map[string]any{"location": location, "encID": encID})
			return nil, err
		})
		if err != nil {
			logger.Warn("failed to link encounter location", "encID", encID, "location", location, "err", err)
			writeWarnings = append(writeWarnings, fmt.Sprintf("location %q: %v", location, err))
		}
	}

	for _, name := range strings.Split(participants, ",") {
		name = strings.TrimSpace(name)
		if name == "" {
			continue
		}
		canonical := resolveCanonical(ctx, session, name)
		_, err := session.ExecuteWrite(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
			if canonical != "" {
				_, err := tx.Run(ctx, `
					MATCH (e {name: $canonical})
					MERGE (enc:Encounter {id: $encID})
					MERGE (e)-[w:WITNESSED]->(enc)
					  ON CREATE SET w.at = $timestamp
				`, map[string]any{"canonical": canonical, "encID": encID, "timestamp": timestamp})
				return nil, err
			}
			_, err := tx.Run(ctx, `
				MERGE (p:Person {name: $name})
				  ON CREATE SET p.lore_verified = false, p.source = "encounter"
				MERGE (enc:Encounter {id: $encID})
				MERGE (p)-[w:WITNESSED]->(enc)
				  ON CREATE SET w.at = $timestamp
			`, map[string]any{"name": name, "encID": encID, "timestamp": timestamp})
			return nil, err
		})
		if err != nil {
			logger.Warn("failed to link encounter participant", "encID", encID, "participant", name, "err", err)
			writeWarnings = append(writeWarnings, fmt.Sprintf("participant %q: %v", name, err))
		}
	}

	result := map[string]any{
		"enc_id":       encID,
		"title":        title,
		"participants": participants,
		"location":     location,
		"timestamp":    timestamp,
	}
	if len(writeWarnings) > 0 {
		result["warnings"] = writeWarnings
	}
	return result, nil
}

// handleQueryAsNPC returns knowledge-graph results scoped to what a named NPC
// has personally witnessed. Two-phase: resolve horizon, then filtered search.
func handleQueryAsNPC(ctx context.Context, cfg Config,
	session neo4j.SessionWithContext, params map[string]any) (any, error) {

	npcName, _ := params["npc_name"].(string)
	textQuery, _ := params["question"].(string)
	limit := 5
	if l, ok := params["limit"].(float64); ok {
		limit = int(l)
	}
	if npcName == "" || textQuery == "" {
		return nil, fmt.Errorf("npc_name and question are required")
	}

	// Phase 1: resolve NPC knowledge horizon
	horizon, tier, err := resolveNPCHorizon(ctx, session, npcName)
	if err != nil {
		return nil, fmt.Errorf("resolve horizon: %w", err)
	}
	horizonSet := make(map[string]bool, len(horizon))
	for _, name := range horizon {
		horizonSet[name] = true
	}

	// Phase 2: embed question + scoped semantic search
	vec, err := embed(ctx, cfg, textQuery)
	if err != nil {
		return nil, fmt.Errorf("embed question: %w", err)
	}

	fetchLimit := limit * 4 // over-fetch so horizon filtering has candidates
	chunks, err := session.ExecuteRead(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
		var rows []map[string]any

		// Message chunks
		res, err := tx.Run(ctx, `
			CALL db.index.vector.queryNodes('chunk_embeddings', $limit, $vec)
			YIELD node AS chunk, score
			MATCH (m:Message)-[:HAS_CHUNK]->(chunk)
			OPTIONAL MATCH (m)-[:MENTIONS]->(e)
			RETURN chunk.text AS text, score, m.id AS msgID,
			       m.author AS author, m.timestamp AS timestamp,
			       collect(DISTINCT coalesce(e.name,"")) AS entities,
			       "message" AS source
			ORDER BY score DESC
		`, map[string]any{"limit": fetchLimit, "vec": vec})
		if err != nil {
			return nil, fmt.Errorf("query NPC message vector index: %w", err)
		}
		for res.Next(ctx) {
			record := res.Record()
			row := map[string]any{}
			for _, key := range record.Keys {
				if v, ok := record.Get(key); ok {
					row[key] = v
				}
			}
			if chunkOverlapsHorizon(row, horizonSet) {
				rows = append(rows, row)
			}
		}
		if err := res.Err(); err != nil {
			return nil, fmt.Errorf("NPC message vector iterate: %w", err)
		}

		// Lore chunks — scholars bypass horizon filter (broad world knowledge)
		loreRes, err := tx.Run(ctx, `
			CALL db.index.vector.queryNodes('lore_chunk_embeddings', $limit, $vec)
			YIELD node AS chunk, score
			MATCH (d:LoreDocument)-[:HAS_CHUNK]->(chunk)
			OPTIONAL MATCH (d)-[:FEATURES]->(e)
			RETURN chunk.text AS text, score, d.id AS msgID,
			       d.title AS author, d.uploaded_at AS timestamp,
			       collect(DISTINCT coalesce(e.name,"")) AS entities,
			       "lore" AS source
			ORDER BY score DESC
		`, map[string]any{"limit": fetchLimit, "vec": vec})
		if err != nil {
			return nil, fmt.Errorf("query NPC lore vector index: %w", err)
		}
		for loreRes.Next(ctx) {
			record := loreRes.Record()
			row := map[string]any{}
			for _, key := range record.Keys {
				if v, ok := record.Get(key); ok {
					row[key] = v
				}
			}
			if tier == "scholar" || chunkOverlapsHorizon(row, horizonSet) {
				rows = append(rows, row)
			}
		}
		if err := loreRes.Err(); err != nil {
			return nil, err
		}

		return rows, nil
	})
	if err != nil {
		return nil, err
	}

	allChunks := chunks.([]map[string]any)
	sort.Slice(allChunks, func(i, j int) bool {
		si, _ := allChunks[i]["score"].(float64)
		sj, _ := allChunks[j]["score"].(float64)
		return si > sj
	})
	if len(allChunks) > limit {
		allChunks = allChunks[:limit]
	}

	// Phase 3: graph context — encounters the NPC witnessed
	graphCtx, err := session.ExecuteRead(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
		res, err := tx.Run(ctx, `
			MATCH (npc:Person {name: $npc_name})-[:WITNESSED]->(enc:Encounter)
			OPTIONAL MATCH (enc)-[:FEATURED]->(e)
			OPTIONAL MATCH (enc)-[:OCCURRED_AT]->(loc:Location)
			RETURN enc.id AS enc_id, enc.title AS enc_title,
			       enc.type AS enc_type, enc.timestamp AS enc_timestamp,
			       enc.summary AS enc_summary,
			       collect(DISTINCT coalesce(e.name,""))   AS featured_entities,
			       collect(DISTINCT coalesce(loc.name,"")) AS locations
			ORDER BY enc.timestamp DESC
			LIMIT 20
		`, map[string]any{"npc_name": npcName})
		if err != nil {
			return nil, err
		}
		var gcRows []map[string]any
		for res.Next(ctx) {
			record := res.Record()
			row := map[string]any{}
			for _, key := range record.Keys {
				row[key], _ = record.Get(key)
			}
			gcRows = append(gcRows, row)
		}
		return gcRows, res.Err()
	})
	if err != nil {
		return nil, err
	}

	return map[string]any{
		"npc":           npcName,
		"tier":          tier,
		"horizon_count": len(horizon),
		"chunks":        allChunks,
		"graph_context": graphCtx,
	}, nil
}

func resolveNPCHorizon(ctx context.Context,
	session neo4j.SessionWithContext, npcName string) ([]string, string, error) {

	result, err := session.ExecuteRead(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
		res, err := tx.Run(ctx, `
			MATCH (npc:Person {name: $npc_name})
			OPTIONAL MATCH (npc)-[:WITNESSED]->(enc:Encounter)-[:FEATURED]->(featured)
			OPTIONAL MATCH (npc)-[:WITNESSED]->(enc2:Encounter)-[:OCCURRED_AT]->(loc)
			OPTIONAL MATCH (npc)-[:MEMBER_OF]->(fac)
			OPTIONAL MATCH (npc)-[:LOCATED_AT]->(home)
			WITH npc,
			     collect(DISTINCT coalesce(featured.name,"")) +
			     collect(DISTINCT coalesce(loc.name,""))      +
			     collect(DISTINCT coalesce(fac.name,""))      +
			     collect(DISTINCT coalesce(home.name,""))     AS rawHorizon
			RETURN [n IN rawHorizon WHERE n <> ""] AS horizon,
			       coalesce(npc.tier, "") AS tier
		`, map[string]any{"npc_name": npcName})
		if err != nil {
			return nil, err
		}
		if res.Next(ctx) {
			record := res.Record()
			horizonRaw, okHorizon := record.Get("horizon")
			tierRaw, okTier := record.Get("tier")
			var tier string
			if okTier {
				tier, _ = tierRaw.(string)
			}
			var horizonList []any
			if okHorizon {
				horizonList, _ = horizonRaw.([]any)
			}
			names := make([]string, 0, len(horizonList))
			for _, v := range horizonList {
				if s, ok := v.(string); ok && s != "" {
					names = append(names, s)
				}
			}
			return map[string]any{"horizon": names, "tier": tier}, nil
		}
		if err := res.Err(); err != nil {
			return nil, err
		}
		return map[string]any{"horizon": []string{}, "tier": ""}, nil
	})
	if err != nil {
		return nil, "", fmt.Errorf("read NPC horizon from db: %w", err)
	}
	m, ok := result.(map[string]any)
	if !ok {
		return nil, "", fmt.Errorf("unexpected horizon result type")
	}
	return m["horizon"].([]string), m["tier"].(string), nil
}

func resolveCanonical(ctx context.Context, session neo4j.SessionWithContext, name string) string {
	result, err := session.ExecuteRead(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
		res, err := tx.Run(ctx, `
			MATCH (e)
			WHERE (e.name = $name OR $name IN coalesce(e.aliases, []))
			  AND e.lore_verified = true
			  AND any(lbl IN labels(e) WHERE lbl IN ['Person','Location','Faction','Item','Creature','Event'])
			RETURN e.name AS canonical LIMIT 1
		`, map[string]any{"name": name})
		if err != nil {
			return "", err
		}
		if res.Next(ctx) {
			v, ok := res.Record().Get("canonical")
			if ok {
				if s, ok := v.(string); ok {
					return s, nil
				}
			}
		}
		return "", res.Err()
	})
	if err != nil || result == nil {
		return ""
	}
	s, ok := result.(string)
	if !ok {
		return ""
	}
	return s
}

func handleGetUnresolved(ctx context.Context,
	session neo4j.SessionWithContext, params map[string]any) (any, error) {

	entityType, _ := params["type"].(string)
	limit := 30
	if l, ok := params["limit"].(float64); ok {
		limit = int(l)
	}

	return session.ExecuteRead(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
		var (
			query string
			args  map[string]any
		)
		if entityType != "" {
			query = `
				MATCH (e)
				WHERE e.lore_verified = false AND $type IN labels(e)
				OPTIONAL MATCH (enc:Encounter)--(e)
				RETURN e.name AS name,
				       [l IN labels(e) WHERE l IN ['Person','Location','Faction','Item','Creature','Event']][0] AS type,
				       e.source AS source,
				       collect(DISTINCT enc.title) AS seen_in
				ORDER BY size(collect(DISTINCT enc.title)) DESC LIMIT $limit
			`
			args = map[string]any{"type": entityType, "limit": limit}
		} else {
			query = `
				MATCH (e)
				WHERE e.lore_verified = false
				  AND any(lbl IN labels(e) WHERE lbl IN ['Person','Location','Faction','Item','Creature','Event'])
				OPTIONAL MATCH (enc:Encounter)--(e)
				RETURN e.name AS name,
				       [l IN labels(e) WHERE l IN ['Person','Location','Faction','Item','Creature','Event']][0] AS type,
				       e.source AS source,
				       collect(DISTINCT enc.title) AS seen_in
				ORDER BY size(collect(DISTINCT enc.title)) DESC LIMIT $limit
			`
			args = map[string]any{"limit": limit}
		}

		res, err := tx.Run(ctx, query, args)
		if err != nil {
			return nil, fmt.Errorf("run unresolved query: %w", err)
		}
		var rows []map[string]any
		for res.Next(ctx) {
			record := res.Record()
			row := map[string]any{}
			for _, key := range record.Keys {
				if v, ok := record.Get(key); ok {
					row[key] = v
				}
			}
			rows = append(rows, row)
		}
		if err := res.Err(); err != nil {
			return nil, fmt.Errorf("unresolved iterate: %w", err)
		}
		return rows, nil
	})
}

func handleGetContradictions(ctx context.Context,
	session neo4j.SessionWithContext, params map[string]any) (any, error) {

	subject, _ := params["subject"].(string)
	limit := 20
	if l, ok := params["limit"].(float64); ok {
		limit = int(l)
	}

	return session.ExecuteRead(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
		var (
			query string
			args  map[string]any
		)
		if subject != "" {
			query = `
				MATCH (a {name: $subject})-[:HAS_CONTRADICTION]->(c:Contradiction)
				RETURN c.subject AS subject, c.predicate AS predicate,
				       c.claim_a AS claim_a, c.doc_a AS doc_a,
				       c.claim_b AS claim_b, c.doc_b AS doc_b,
				       c.detected_at AS detected_at
				ORDER BY c.detected_at DESC LIMIT $limit
			`
			args = map[string]any{"subject": subject, "limit": limit}
		} else {
			query = `
				MATCH (c:Contradiction)
				RETURN c.subject AS subject, c.predicate AS predicate,
				       c.claim_a AS claim_a, c.doc_a AS doc_a,
				       c.claim_b AS claim_b, c.doc_b AS doc_b,
				       c.detected_at AS detected_at
				ORDER BY c.detected_at DESC LIMIT $limit
			`
			args = map[string]any{"limit": limit}
		}

		res, err := tx.Run(ctx, query, args)
		if err != nil {
			return nil, fmt.Errorf("run contradictions query: %w", err)
		}
		var rows []map[string]any
		for res.Next(ctx) {
			record := res.Record()
			row := map[string]any{}
			for _, key := range record.Keys {
				if v, ok := record.Get(key); ok {
					row[key] = v
				}
			}
			rows = append(rows, row)
		}
		if err := res.Err(); err != nil {
			return nil, fmt.Errorf("contradictions iterate: %w", err)
		}
		return rows, nil
	})
}

func chunkOverlapsHorizon(row map[string]any, horizonSet map[string]bool) bool {
	entities, ok := row["entities"].([]any)
	if !ok {
		return false
	}
	for _, v := range entities {
		if s, ok := v.(string); ok && horizonSet[s] {
			return true
		}
	}
	return false
}

// ── Helpers ───────────────────────────────────────────────────────────────────

func getEnv(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}

func handleListEncounters(ctx context.Context, session neo4j.SessionWithContext, params map[string]any) (any, error) {
	limit := 10
	if l, ok := params["limit"].(float64); ok {
		limit = int(l)
	}
	return session.ExecuteRead(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
		res, err := tx.Run(ctx, `
			MATCH (enc:Encounter)
			RETURN enc.id AS id, enc.title AS title, enc.location_name AS location, enc.timestamp AS timestamp, enc.summary AS summary
			ORDER BY enc.timestamp DESC LIMIT $limit
		`, map[string]any{"limit": limit})
		if err != nil {
			return nil, fmt.Errorf("list encounters query: %w", err)
		}
		var rows []map[string]any
		for res.Next(ctx) {
			record := res.Record()
			row := map[string]any{}
			for _, key := range record.Keys {
				if v, ok := record.Get(key); ok {
					row[key] = v
				}
			}
			rows = append(rows, row)
		}
		if err := res.Err(); err != nil {
			return nil, fmt.Errorf("list encounters iterate: %w", err)
		}
		return rows, nil
	})
}

func handleSearchEncounters(ctx context.Context, session neo4j.SessionWithContext, params map[string]any) (any, error) {
	query, _ := params["query"].(string)
	location, _ := params["location"].(string)
	participant, _ := params["participant"].(string)
	limit := 10
	if l, ok := params["limit"].(float64); ok {
		limit = int(l)
	}

	return session.ExecuteRead(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
		var (
			matchCls = []string{"(enc:Encounter)"}
			whereCls = []string{}
			args     = map[string]any{"limit": limit}
		)

		if query != "" {
			whereCls = append(whereCls, "(toLower(enc.title) CONTAINS toLower($query) OR toLower(enc.summary) CONTAINS toLower($query))")
			args["query"] = query
		}
		if location != "" {
			whereCls = append(whereCls, "toLower(enc.location_name) CONTAINS toLower($location)")
			args["location"] = location
		}
		if participant != "" {
			matchCls = append(matchCls, "(p:Person)-[:WITNESSED]->(enc)")
			whereCls = append(whereCls, "toLower(p.name) CONTAINS toLower($participant)")
			args["participant"] = participant
		}

		matchStr := "MATCH " + strings.Join(matchCls, ", ")
		whereStr := ""
		if len(whereCls) > 0 {
			whereStr = "WHERE " + strings.Join(whereCls, " AND ")
		}

		cypher := fmt.Sprintf(`
			%s
			%s
			RETURN DISTINCT enc.id AS id, enc.title AS title, enc.location_name AS location, enc.timestamp AS timestamp, enc.summary AS summary
			ORDER BY enc.timestamp DESC LIMIT $limit
		`, matchStr, whereStr)

		res, err := tx.Run(ctx, cypher, args)
		if err != nil {
			return nil, fmt.Errorf("search encounters query: %w", err)
		}
		var rows []map[string]any
		for res.Next(ctx) {
			record := res.Record()
			row := map[string]any{}
			for _, key := range record.Keys {
				if v, ok := record.Get(key); ok {
					row[key] = v
				}
			}
			rows = append(rows, row)
		}
		if err := res.Err(); err != nil {
			return nil, fmt.Errorf("search encounters iterate: %w", err)
		}
		return rows, nil
	})
}

func handleGetEncounter(ctx context.Context, session neo4j.SessionWithContext, params map[string]any) (any, error) {
	id, _ := params["id"].(string)
	if id == "" {
		return nil, fmt.Errorf("id parameter is required")
	}

	return session.ExecuteRead(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
		res, err := tx.Run(ctx, `
			MATCH (enc:Encounter {id: $id})
			OPTIONAL MATCH (p:Person)-[:WITNESSED]->(enc)
			OPTIONAL MATCH (enc)-[:FEATURED]->(e)
			RETURN enc.id AS id, enc.title AS title, enc.location_name AS location,
			       enc.timestamp AS timestamp, enc.summary AS summary, enc.type AS type,
			       collect(DISTINCT p.name) AS participants,
			       collect(DISTINCT e.name) AS featured_entities
		`, map[string]any{"id": id})
		if err != nil {
			return nil, fmt.Errorf("get encounter query: %w", err)
		}
		if res.Next(ctx) {
			record := res.Record()
			row := map[string]any{}
			for _, key := range record.Keys {
				if v, ok := record.Get(key); ok {
					row[key] = v
				}
			}
			return row, nil
		}
		if err := res.Err(); err != nil {
			return nil, fmt.Errorf("get encounter iterate: %w", err)
		}
		return nil, fmt.Errorf("encounter not found: %s", id)
	})
}