package mcp

import (
	"encoding/json"
	"fmt"
)

// MCP protocol surface. The patterm server originally exposed each
// tool as its own JSON-RPC method (and the harness still drives it
// that way). Real MCP clients (claude, codex, opencode) speak the
// model-context-protocol RPC dialect: they send `initialize` first,
// then `tools/list`, then `tools/call` with `{name, arguments}`. This
// file wraps those four entry points around the existing tool dispatch
// without changing the underlying tool implementations.

// supportedProtocolVersion is the MCP protocol revision we advertise
// when a client doesn't pin a specific version. Claude Code accepts
// the dated-string scheme used by the MCP spec.
const supportedProtocolVersion = "2025-06-18"

// serverInfo identifies the server back to the client during the
// initialize handshake. The version is intentionally kept generic so
// it doesn't need bumping per release; clients only key behavior off
// name + protocol version.
var serverInfo = map[string]any{
	"name":    "patterm",
	"version": "0.1.0",
}

// serverInstructions is returned in the MCP `initialize` response. MCP
// clients show this to the underlying LLM as context for how to use
// the server. Failure modes we've seen and want to head off:
//   - The agent assumes patterm is something it has to launch (running
//     `patterm` or `patterm mcp-stdio` from its own shell). It's
//     already attached — it just calls the tools.
//   - The agent reaches for shell tools (perl / nc / socat / curl) to
//     poke patterm's Unix socket directly. That socket connection
//     carries no caller identity, so any sub-agent the agent spawns
//     that way ends up as a stray top-level tab instead of a child
//     under the spawning agent. Always go through the MCP tools.
//   - The agent shells out to `claude` / `codex` / `opencode` to start
//     a peer instead of calling `spawn_agent`. Those peers won't show
//     up as sub-agents and won't be tied into the patterm lifecycle.
//
// Keep this short — clients vary in how much they surface to the LLM.
const serverInstructions = "You are already running INSIDE patterm; the `patterm` MCP server is connected over the same stdio MCP transport you use for any other MCP server. Use the MCP tools you see in tools/list — do NOT (a) try to launch `patterm` or `patterm mcp-stdio` yourself, (b) poke the Unix socket through perl / nc / socat / curl, or (c) shell out to `claude` / `codex` / `opencode` to start a peer. Any of those bypasses caller-identity and the new agent will land as a stray top-level tab instead of a child under you. Start with `whoami` for your role and the full tool list, then `help('topics')` for orientation. `spawn_agent` is the only correct way to start a sub-agent; `spawn_process` is for non-LLM commands; `list_processes` / `get_process_output` inspect them; `send_input` / `send_message` drive them. Whatever you spawn is yours to `close_process` when done. When you `send_message` a sub-agent, its reply comes back into YOUR pane as `[sub-agent:<name>] …`, not into the sub-agent's output — to wait for it, use `timer_fire_when_idle_any([sub_agent])` and then read your own pane; do NOT `wait_for_pattern` on the sub-agent, that will deadlock until timeout."

// toolDescriptor is the shape returned by `tools/list`. inputSchema is
// a JSON Schema object — we provide a minimal `{type: "object"}` schema
// for each tool, which lets MCP clients accept arbitrary arguments and
// rely on patterm's own server-side validation for typing.
type toolDescriptor struct {
	Name        string         `json:"name"`
	Description string         `json:"description"`
	InputSchema map[string]any `json:"inputSchema"`
}

// objectSchema builds an inputSchema for a tool that takes an object
// with the listed properties. required lists property names that must
// be present; passing nil makes them all optional. We always emit a
// concrete `properties` object (never null) because some MCP clients
// reject schemas where `properties` is not an object.
func objectSchema(properties map[string]any, required []string) map[string]any {
	if properties == nil {
		properties = map[string]any{}
	}
	s := map[string]any{
		"type":                 "object",
		"properties":           properties,
		"additionalProperties": true,
	}
	if len(required) > 0 {
		s["required"] = required
	}
	return s
}

func stringProp(desc string) map[string]any {
	return map[string]any{"type": "string", "description": desc}
}

func numberProp(desc string) map[string]any {
	return map[string]any{"type": "number", "description": desc}
}

func integerProp(desc string) map[string]any {
	return map[string]any{"type": "integer", "description": desc}
}

func booleanProp(desc string) map[string]any {
	return map[string]any{"type": "boolean", "description": desc}
}

func arrayOfStringsProp(desc string) map[string]any {
	return map[string]any{
		"type":        "array",
		"description": desc,
		"items":       map[string]any{"type": "string"},
	}
}

// toolCatalog is the full list advertised via tools/list. Descriptions
// are intentionally short — clients are expected to fetch help() for
// detail. Schemas mirror the param structs in tools.go.
func toolCatalog() []toolDescriptor {
	return []toolDescriptor{
		{
			Name:        "spawn_agent",
			Description: "Spawn a sub-agent from an agent preset and optionally seed it with initial instructions. This is the ONLY correct way to start a sub-agent under you — do not shell out to `claude` / `codex` / `opencode` and do not poke patterm's Unix socket via perl / nc / socat. Either bypasses caller identity and the new agent lands as a stray top-level tab instead of your child. Caller owns lifecycle: when the sub-agent's work is done (it reports back via send_message, or you no longer need it), call close_process on its process_id to free the pane and tear down the PTY. See help('spawning') and help('lifecycle').",
			InputSchema: objectSchema(map[string]any{
				"agent":              stringProp("Preset name (e.g. \"claude\", \"codex\")."),
				"agent_instructions": stringProp("Initial prompt typed into the agent after it's ready."),
				"name":               stringProp("Display name for the new pane."),
			}, []string{"agent"}),
		},
		{
			Name:        "spawn_process",
			Description: "Spawn a process: a terminal, a process preset, or a freeform argv command. Caller owns lifecycle: when the process is no longer needed, call close_process to remove its entry (live children are SIGKILL'd first). See help('lifecycle').",
			InputSchema: objectSchema(map[string]any{
				"kind":        stringProp("\"terminal\" or \"command\"."),
				"preset":      stringProp("Process preset name (mutually exclusive with argv)."),
				"argv":        map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "Argv vector for freeform commands."},
				"name":        stringProp("Display name for the pane."),
				"working_dir": stringProp("Working directory for the spawned process."),
				"env":         map[string]any{"type": "object", "additionalProperties": map[string]any{"type": "string"}, "description": "Extra environment variables."},
				"shell":       booleanProp("Run argv through sh -lc."),
			}, nil),
		},
		{
			Name:        "start_process",
			Description: "(Re)attach a PTY to a session-persistent command process that has exited.",
			InputSchema: objectSchema(map[string]any{
				"process_id": stringProp("Target process id."),
			}, []string{"process_id"}),
		},
		{
			Name:        "restart_process",
			Description: "Signal the target process and restart it under a fresh PTY.",
			InputSchema: objectSchema(map[string]any{
				"process_id": stringProp("Target process id."),
				"signal":     integerProp("Signal to send before relaunch (default SIGTERM)."),
			}, []string{"process_id"}),
		},
		{
			Name:        "stop_process",
			Description: "Send a signal to a running process without removing its entry.",
			InputSchema: objectSchema(map[string]any{
				"process_id": stringProp("Target process id."),
				"signal":     integerProp("Signal to send (default SIGTERM)."),
			}, []string{"process_id"}),
		},
		{
			Name:        "close_process",
			Description: "Remove the process entry entirely; live children are SIGKILL'd first.",
			InputSchema: objectSchema(map[string]any{
				"process_id": stringProp("Target process id."),
			}, []string{"process_id"}),
		},
		{
			Name:        "rename_process",
			Description: "Rename the pane label for a process.",
			InputSchema: objectSchema(map[string]any{
				"process_id": stringProp("Target process id."),
				"name":       stringProp("New display name."),
			}, []string{"process_id", "name"}),
		},
		{
			Name:        "select_process",
			Description: "Focus the named process in the host TUI.",
			InputSchema: objectSchema(map[string]any{
				"process_id": stringProp("Target process id."),
			}, []string{"process_id"}),
		},
		{
			Name:        "list_processes",
			Description: "List visible processes, optionally filtered by kind (\"agent\", \"command\", \"terminal\").",
			InputSchema: objectSchema(map[string]any{
				"kind": stringProp("Optional kind filter."),
			}, nil),
		},
		{
			Name:        "get_process_status",
			Description: "Return rich status (status, geometry, cursor, screen version) for one process.",
			InputSchema: objectSchema(map[string]any{
				"process_id": stringProp("Target process id."),
			}, []string{"process_id"}),
		},
		{
			Name:        "get_project_status",
			Description: "One-shot orientation: project, caller, processes, scratchpads.",
			InputSchema: objectSchema(nil, nil),
		},
		{
			Name:        "get_process_output",
			Description: "Read rendered grid (\"grid\") or ANSI-stripped stream (\"stream\") output, with screen-version watermark.",
			InputSchema: objectSchema(map[string]any{
				"process_id":   stringProp("Target process id."),
				"mode":         stringProp("\"grid\" (default) or \"stream\"."),
				"since_offset": integerProp("Watermark offset from a previous call."),
			}, []string{"process_id"}),
		},
		{
			Name:        "get_process_raw_output",
			Description: "Read the raw ANSI byte stream since since_offset.",
			InputSchema: objectSchema(map[string]any{
				"process_id":   stringProp("Target process id."),
				"since_offset": integerProp("Byte offset from a previous call."),
			}, []string{"process_id"}),
		},
		{
			Name:        "search_output",
			Description: "Search a process's rendered or raw output and return matching lines.",
			InputSchema: objectSchema(map[string]any{
				"process_id": stringProp("Target process id."),
				"pattern":    stringProp("Regex pattern."),
				"kind":       stringProp("\"rendered\" (default) or \"raw\"."),
				"limit":      integerProp("Max matches (default 20)."),
			}, []string{"process_id", "pattern"}),
		},
		{
			Name:        "wait_for_pattern",
			Description: "Block until pattern appears in the TARGET process's own output, or timeout elapses. Use this for waiting on text the target itself will emit (a shell prompt, a build's \"tests passed\" line, etc.). Anti-pattern: do NOT use this to wait for a sub-agent's reply to send_message — replies are routed into the CALLER's pane tagged `[sub-agent:<name>]`, not into the sub-agent's output, so this call will spin to timeout. For sub-agent coordination use `timer_fire_when_idle_any` and then read your own pane.",
			InputSchema: objectSchema(map[string]any{
				"process_id":      stringProp("Target process id."),
				"pattern":         stringProp("Regex pattern."),
				"timeout_seconds": numberProp("Max time to wait (seconds)."),
				"scope":           stringProp("\"grid\" (default) or \"scrollback\"."),
			}, []string{"process_id", "pattern"}),
		},
		{
			Name:        "get_process_ports",
			Description: "Return URL-form port sightings observed in a process's output.",
			InputSchema: objectSchema(map[string]any{
				"process_id": stringProp("Target process id."),
			}, []string{"process_id"}),
		},
		{
			Name:        "send_input",
			Description: "Type text, paste a block, or fire a named key into a process. Optional tail-after-send.",
			InputSchema: objectSchema(map[string]any{
				"process_id": stringProp("Target process id."),
				"kind":       stringProp("\"text\", \"paste\", or \"key\"."),
				"text":       stringProp("Text payload for kind=text/paste."),
				"key":        stringProp("Named key for kind=key (e.g. \"enter\", \"escape\")."),
				"submit":     booleanProp("Whether to append a submit keystroke."),
				"wait_ms":    integerProp("After sending, wait this many ms before tailing."),
				"tail_mode":  stringProp("\"none\" (default), \"stream\", or \"grid\"."),
			}, []string{"process_id", "kind"}),
		},
		{
			Name:        "send_message",
			Description: "Deliver a text message to another process as orchestrator-owned input. Fire-and-forget: returns immediately, without waiting for the recipient to read or act. If the recipient replies via send_message, that reply arrives in YOUR pane tagged `[sub-agent:<name>]` (child→parent) or `[orchestrator]` (parent→child) — NOT in the recipient's output. To wait for a sub-agent's reply, schedule `timer_fire_when_idle_any([sub_agent_id], body=…)` and then read your own pane when the timer fires. Do not `wait_for_pattern` on the recipient for a reply; it will deadlock.",
			InputSchema: objectSchema(map[string]any{
				"target_process_id": stringProp("Recipient process id."),
				"message":           stringProp("Message body."),
			}, []string{"target_process_id", "message"}),
		},
		{
			Name:        "request_human_attention",
			Description: "Flag a process pane as needing human review.",
			InputSchema: objectSchema(map[string]any{
				"process_id": stringProp("Target process id."),
				"reason":     stringProp("Short description shown to the human."),
			}, []string{"process_id", "reason"}),
		},
		{
			Name:        "timer_wait",
			Description: "Schedule a delay timer that injects a fixed `[system]` line into your pane when it fires (legacy; prefer timer_set).",
			InputSchema: objectSchema(map[string]any{
				"seconds": numberProp("Delay duration."),
				"label":   stringProp("Optional label for diagnostics."),
			}, []string{"seconds"}),
		},
		{
			Name:        "timer_set",
			Description: "Schedule a one-shot delay timer that delivers `body` to the owning agent as a fresh user turn when it fires.",
			InputSchema: objectSchema(map[string]any{
				"seconds":          numberProp("Delay duration."),
				"body":             stringProp("Message delivered verbatim to the owning agent as a user turn when the timer fires."),
				"label":            stringProp("Optional label for diagnostics."),
				"owner_process_id": stringProp("Owner process id; defaults to the caller. Top-level callers must supply this explicitly."),
			}, []string{"seconds", "body"}),
		},
		{
			Name:        "timer_fire_when_idle_any",
			Description: "Canonical way to wait for a sub-agent to finish working: send_message the sub-agent, then schedule this with watched=[sub_agent_id]; when it fires, the reply is already sitting in your own pane tagged `[sub-agent:<name>]`. Schedules a timer that fires when any watched process enters idle (already-idle entries excluded), or when max_wait_seconds elapses.",
			InputSchema: objectSchema(map[string]any{
				"watched":          arrayOfStringsProp("Process ids to watch."),
				"body":             stringProp("Message delivered verbatim to the owning agent when the timer fires."),
				"label":            stringProp("Optional label for diagnostics."),
				"max_wait_seconds": numberProp("Optional cap; 0 means no fallback fire."),
				"owner_process_id": stringProp("Owner process id; defaults to the caller."),
			}, []string{"watched", "body"}),
		},
		{
			Name:        "timer_fire_when_idle_all",
			Description: "Canonical way to wait for several sub-agents to finish working in parallel: send_message each one, then schedule this with watched=[…ids]; when it fires, each reply is in your own pane tagged `[sub-agent:<name>]`. Schedules a timer that fires when all watched processes are idle (already-idle entries count as satisfied), or when max_wait_seconds elapses.",
			InputSchema: objectSchema(map[string]any{
				"watched":          arrayOfStringsProp("Process ids to watch."),
				"body":             stringProp("Message delivered verbatim to the owning agent when the timer fires."),
				"label":            stringProp("Optional label for diagnostics."),
				"max_wait_seconds": numberProp("Optional cap; 0 means no fallback fire."),
				"owner_process_id": stringProp("Owner process id; defaults to the caller."),
			}, []string{"watched", "body"}),
		},
		{
			Name:        "timer_cancel",
			Description: "Cancel one pending timer owned by the caller.",
			InputSchema: objectSchema(map[string]any{
				"timer_id": stringProp("Timer id returned by a previous timer_* call."),
			}, []string{"timer_id"}),
		},
		{
			Name:        "timer_pause",
			Description: "Pause one pending timer owned by the caller. Idle-aware timers stop listening to state changes; delay timers preserve their remaining wait.",
			InputSchema: objectSchema(map[string]any{
				"timer_id": stringProp("Timer id."),
			}, []string{"timer_id"}),
		},
		{
			Name:        "timer_resume",
			Description: "Resume one paused timer owned by the caller.",
			InputSchema: objectSchema(map[string]any{
				"timer_id": stringProp("Timer id."),
			}, []string{"timer_id"}),
		},
		{
			Name:        "timer_list",
			Description: "List pending and paused timers owned by the caller.",
			InputSchema: objectSchema(nil, nil),
		},
		{
			Name:        "scratchpad_list",
			Description: "List shared per-project scratchpad entries.",
			InputSchema: objectSchema(nil, nil),
		},
		{
			Name:        "scratchpad_read",
			Description: "Read a scratchpad entry, returning content and revision.",
			InputSchema: objectSchema(map[string]any{
				"name": stringProp("Scratchpad name."),
			}, []string{"name"}),
		},
		{
			Name:        "scratchpad_write",
			Description: "Write a scratchpad entry with optimistic concurrency on expected_revision.",
			InputSchema: objectSchema(map[string]any{
				"name":              stringProp("Scratchpad name."),
				"content":           stringProp("New content."),
				"expected_revision": stringProp("Last-seen revision token."),
			}, []string{"name", "content"}),
		},
		{
			Name:        "scratchpad_append",
			Description: "Append to a scratchpad entry without revision checking.",
			InputSchema: objectSchema(map[string]any{
				"name":    stringProp("Scratchpad name."),
				"content": stringProp("Text to append."),
			}, []string{"name", "content"}),
		},
		{
			Name:        "whoami",
			Description: "Return the caller's identity, role, parent, project metadata, and available tools.",
			InputSchema: objectSchema(nil, nil),
		},
		{
			Name:        "help",
			Description: "Return human-readable help for a topic (e.g. tool name).",
			InputSchema: objectSchema(map[string]any{
				"topic": stringProp("Topic or tool name (empty for index)."),
			}, nil),
		},
	}
}

// handleProtocolMethod handles MCP protocol-level methods. Returns
// (result, handled). When handled is false, the caller falls back to
// the legacy direct-tool dispatch. For notifications, result is nil
// and handled is true.
func (s *Server) handleProtocolMethod(callerID, method string, params json.RawMessage, isNotification bool) (any, bool, int, string, any) {
	switch method {
	case "initialize":
		var p struct {
			ProtocolVersion string         `json:"protocolVersion"`
			Capabilities    map[string]any `json:"capabilities"`
			ClientInfo      map[string]any `json:"clientInfo"`
		}
		_ = unmarshalParamsOptional(params, &p)
		protoVersion := p.ProtocolVersion
		if protoVersion == "" {
			protoVersion = supportedProtocolVersion
		}
		result := map[string]any{
			"protocolVersion": protoVersion,
			"capabilities": map[string]any{
				"tools": map[string]any{"listChanged": false},
			},
			"serverInfo":   serverInfo,
			"instructions": serverInstructions,
		}
		return result, true, 0, "", nil

	case "notifications/initialized", "notifications/cancelled", "notifications/roots/list_changed":
		// Notifications get no response — handled is true so the caller
		// doesn't fall through to legacy dispatch, but result is nil.
		return nil, true, 0, "", nil

	case "ping":
		return map[string]any{}, true, 0, "", nil

	case "tools/list":
		return map[string]any{"tools": toolCatalog()}, true, 0, "", nil

	case "tools/call":
		var p struct {
			Name      string          `json:"name"`
			Arguments json.RawMessage `json:"arguments"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, true, codeInvalidParams, err.Error(), nil
		}
		if p.Name == "" {
			return nil, true, codeInvalidParams, "tools/call: name required", nil
		}
		s.mu.Lock()
		host := s.host
		s.mu.Unlock()
		if host == nil {
			return nil, true, codeInternal, "patterm: tool host not initialized", nil
		}
		result, code, errMsg, data := callTool(host, callerID, p.Name, p.Arguments)
		if errMsg != "" {
			// MCP convention: errors during tool execution come back as
			// successful tools/call results with isError=true, so the
			// model sees the failure as content rather than a transport
			// error. Genuine transport errors (parse, etc.) stay as
			// JSON-RPC errors and are handled outside this branch.
			content := errMsg
			if data != nil {
				if kindMap, ok := data.(map[string]string); ok {
					if k, present := kindMap["kind"]; present && k != "" {
						content = fmt.Sprintf("%s (%s)", errMsg, k)
					}
				}
			}
			_ = code // code stays useful for legacy callers; tools/call surfaces text.
			return map[string]any{
				"content": []map[string]any{{"type": "text", "text": content}},
				"isError": true,
			}, true, 0, "", nil
		}
		return wrapToolResult(result), true, 0, "", nil

	case "resources/list":
		// We don't expose resources; respond with an empty list rather
		// than a method-not-found to keep clients happy.
		return map[string]any{"resources": []any{}}, true, 0, "", nil

	case "prompts/list":
		return map[string]any{"prompts": []any{}}, true, 0, "", nil

	case "logging/setLevel":
		return map[string]any{}, true, 0, "", nil
	}
	return nil, false, 0, "", nil
}

// wrapToolResult turns a structured tool result into an MCP tools/call
// response. Plain strings (e.g. "ok") become text content; structured
// values are JSON-encoded into a single text block and also exposed
// under structuredContent so capable clients can read the shape.
func wrapToolResult(result any) map[string]any {
	var text string
	switch v := result.(type) {
	case nil:
		text = "ok"
	case string:
		text = v
	default:
		b, err := json.Marshal(v)
		if err != nil {
			text = fmt.Sprintf("%v", v)
		} else {
			text = string(b)
		}
	}
	out := map[string]any{
		"content": []map[string]any{{"type": "text", "text": text}},
		"isError": false,
	}
	if result != nil {
		switch result.(type) {
		case string:
			// Skip — plain string already lives in content.
		default:
			out["structuredContent"] = result
		}
	}
	return out
}