package main

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

// An MCP server (Model Context Protocol, Streamable HTTP transport) on
// POST /mcp, so an assistant can read the router's state and edit its
// settings. It is a thin layer over what the web UI does:
//
//   - tokens (tokens.go) instead of the browser session; scope read or change;
//   - changes only edit settings.json, like the Advanced page: they stay
//     unapplied until a person applies them in the web UI (commit-confirm);
//   - secrets are never returned (the settings hold only paths to them), and
//     file paths in the settings must stay under the state directory;
//   - SSH keys, the release schema and the state version are out of reach;
//   - every call is written to the audit log.
//
// Only requests and responses of JSON-RPC 2.0 over POST are implemented: no
// server-initiated messages, no streaming.

var mcpVersions = []string{"2025-06-18", "2025-03-26", "2024-11-05"}

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

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

type rpcResponse struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id"`
	Result  any             `json:"result,omitempty"`
	Error   *rpcError       `json:"error,omitempty"`
}

const mcpInstructions = `This is the management interface of a nixos-router.
Read the state with get_status, list_interfaces, list_dhcp_leases, get_settings and get_update_status;
get_logs reads the system journal (a unit, a time range, a priority, a text to search for).
Settings are one JSON document (the resource router://settings; router://options describes every option).
Change them with update_settings (a JSON merge patch; pass the etag from get_settings), then run check_settings,
which evaluates the configuration and reports the router's own assertions. Changes are saved but not applied:
a person reviews them (get_pending_changes shows the difference) and applies them in the web UI.
Never put secrets into the settings: they refer to files (passwordFile, secretFile, ...) that the person provides.`

// mcp is the HTTP handler of the endpoint.
func (a *app) mcp(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Cache-Control", "no-store")
	if r.Method != http.MethodPost {
		w.Header().Set("Allow", "POST")
		http.Error(w, "this endpoint takes POST (JSON-RPC)", http.StatusMethodNotAllowed)
		return
	}
	tok, ok := a.bearer(w, r)
	if !ok {
		return
	}
	body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, 1<<20))
	if err != nil {
		http.Error(w, "request too large", http.StatusRequestEntityTooLarge)
		return
	}
	var req rpcRequest
	if err := json.Unmarshal(body, &req); err != nil || req.JSONRPC != "2.0" || req.Method == "" {
		a.rpcReply(w, rpcResponse{ID: json.RawMessage("null"), Error: &rpcError{-32700, "not a JSON-RPC 2.0 request"}})
		return
	}
	if len(req.ID) == 0 { // a notification
		w.WriteHeader(http.StatusAccepted)
		return
	}
	resp := rpcResponse{ID: req.ID}
	result, rerr := a.mcpDispatch(tok, req)
	if rerr != nil {
		resp.Error = rerr
	} else {
		resp.Result = result
	}
	a.rpcReply(w, resp)
}

// bearer authenticates an API request (MCP or REST) by its token; it has
// answered the request itself when it returns false.
func (a *app) bearer(w http.ResponseWriter, r *http.Request) (Token, bool) {
	// A browser page must not be able to talk to the router with a token it
	// does not have, but DNS rebinding could present a foreign Origin.
	if o := r.Header.Get("Origin"); o != "" {
		u, err := url.Parse(o)
		if err != nil || u.Host != r.Host {
			http.Error(w, "forbidden origin", http.StatusForbidden)
			return Token{}, false
		}
	}
	secret, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ")
	tok, valid := a.tokens.Lookup(strings.TrimSpace(secret))
	if !ok || !valid {
		time.Sleep(time.Second) // a guess costs a second
		w.Header().Set("WWW-Authenticate", `Bearer realm="router"`)
		http.Error(w, "a valid token is required", http.StatusUnauthorized)
		return Token{}, false
	}
	return tok, true
}

func (a *app) rpcReply(w http.ResponseWriter, resp rpcResponse) {
	resp.JSONRPC = "2.0"
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(resp)
}

func (a *app) mcpDispatch(tok Token, req rpcRequest) (any, *rpcError) {
	switch req.Method {
	case "initialize":
		var p struct {
			ProtocolVersion string `json:"protocolVersion"`
		}
		_ = json.Unmarshal(req.Params, &p)
		version := mcpVersions[0]
		for _, v := range mcpVersions {
			if v == p.ProtocolVersion {
				version = v
			}
		}
		return map[string]any{
			"protocolVersion": version,
			"capabilities":    map[string]any{"tools": map[string]any{}, "resources": map[string]any{}},
			"serverInfo":      map[string]any{"name": "nixos-router", "version": a.sys.installedRelease().Version},
			"instructions":    mcpInstructions,
		}, nil
	case "ping":
		return map[string]any{}, nil
	case "tools/list":
		var out []map[string]any
		for _, t := range mcpTools {
			if tok.allows(t.scope) {
				out = append(out, map[string]any{"name": t.name, "description": t.description, "inputSchema": t.schema})
			}
		}
		return map[string]any{"tools": out}, nil
	case "tools/call":
		var p struct {
			Name      string          `json:"name"`
			Arguments json.RawMessage `json:"arguments"`
		}
		if err := json.Unmarshal(req.Params, &p); err != nil {
			return nil, &rpcError{-32602, "bad parameters"}
		}
		return a.mcpCall(tok, p.Name, p.Arguments)
	case "resources/list":
		res := []map[string]any{
			{"uri": "router://settings", "name": "settings", "description": "The saved settings (settings.json).", "mimeType": "application/json"},
		}
		if a.optionsFile != "" {
			res = append(res, map[string]any{"uri": "router://options", "name": "options",
				"description": "Every router option: type, default, description (what settings.json may contain).", "mimeType": "application/json"})
		}
		return map[string]any{"resources": res}, nil
	case "resources/read":
		var p struct {
			URI string `json:"uri"`
		}
		if err := json.Unmarshal(req.Params, &p); err != nil {
			return nil, &rpcError{-32602, "bad parameters"}
		}
		var text string
		switch p.URI {
		case "router://settings":
			data, err := os.ReadFile(a.sys.settingsPath())
			if err != nil {
				return nil, &rpcError{-32603, err.Error()}
			}
			text = string(data)
		case "router://options":
			data, err := os.ReadFile(a.optionsFile)
			if a.optionsFile == "" || err != nil {
				return nil, &rpcError{-32002, "no such resource"}
			}
			text = string(data)
		default:
			return nil, &rpcError{-32002, "no such resource"}
		}
		return map[string]any{"contents": []map[string]any{{"uri": p.URI, "mimeType": "application/json", "text": text}}}, nil
	}
	return nil, &rpcError{-32601, "method not found: " + req.Method}
}

type mcpTool struct {
	name, description, scope string
	schema                   map[string]any
	run                      func(a *app, args json.RawMessage) (any, error)
}

// mcpCall runs a tool. A failing tool is not a protocol error: the assistant
// is told what went wrong and can correct itself.
func (a *app) mcpCall(tok Token, name string, args json.RawMessage) (any, *rpcError) {
	var tool *mcpTool
	for i := range mcpTools {
		if mcpTools[i].name == name {
			tool = &mcpTools[i]
		}
	}
	if tool == nil || !tok.allows(tool.scope) {
		return nil, &rpcError{-32602, "unknown tool: " + name}
	}
	res, err := a.runTool(tok, tool, args)
	if err != nil {
		return map[string]any{"isError": true, "content": []map[string]any{{"type": "text", "text": err.Error()}}}, nil
	}
	text, ok := res.(string)
	if !ok {
		data, merr := json.MarshalIndent(res, "", "  ")
		if merr != nil {
			return nil, &rpcError{-32603, merr.Error()}
		}
		text = string(data)
	}
	return map[string]any{"content": []map[string]any{{"type": "text", "text": text}}}, nil
}

// runTool runs a tool for a token (the scope is checked by the caller) and
// writes the call to the audit log.
func (a *app) runTool(tok Token, tool *mcpTool, args json.RawMessage) (any, error) {
	if len(args) == 0 {
		args = json.RawMessage("{}")
	}
	res, err := tool.run(a, args)
	entry := AuditEntry{Time: time.Now().UTC(), Token: tok.Name, Tool: tool.name, Args: string(args)}
	if err != nil {
		entry.Error = err.Error()
	}
	a.sys.audit(entry)
	if err == nil && tool.name == "update_settings" {
		if m, ok := res.(map[string]any); ok {
			if e, ok := m["etag"].(string); ok {
				a.sys.noteTokenEdit(tok.Name, e)
			}
		}
	}
	return res, err
}

// errStale: the settings changed since the caller read them.
var errStale = errors.New("the settings changed since you read them (someone else edited, or the etag is wrong): call get_settings again")

// argError is a tool's complaint about its arguments (HTTP 400 in the REST API).
type argError string

func (e argError) Error() string { return string(e) }

// etag identifies a version of the settings file.
func etag(data []byte) string {
	h := sha256.Sum256(data)
	return hex.EncodeToString(h[:8])
}

func finite(f float64) any {
	if math.IsNaN(f) || math.IsInf(f, 0) {
		return nil
	}
	return math.Round(f*10) / 10
}

var emptySchema = map[string]any{"type": "object", "properties": map[string]any{}}

var mcpTools = []mcpTool{
	{
		name: "get_status", scope: scopeRead, schema: emptySchema,
		description: "The router's state now: uplink (connected, addresses, gateway, pings), system (uptime, load, CPU, memory, disk, temperature, kernel), whether settings are unapplied, and whether an apply is running or waits for confirmation.",
		run: func(a *app, _ json.RawMessage) (any, error) {
			snap := a.metrics.Snapshot()
			st := a.sys.Status()
			u, s := snap.Uplink, snap.System
			var pings []map[string]any
			for _, p := range snap.Pings {
				pings = append(pings, map[string]any{"target": p.Target, "address": p.Address, "ok": p.OK, "rttMs": p.RTT.Milliseconds()})
			}
			return map[string]any{
				"time": snap.Time,
				"uplink": map[string]any{"device": u.Device, "type": u.Type, "connected": u.Connected, "reachable": u.Reachable,
					"ipv4": u.IPv4, "ipv6": u.IPv6, "gateway": u.Gateway, "since": u.Since},
				"pings": pings,
				"system": map[string]any{"uptimeSeconds": int(s.Uptime.Seconds()), "load": s.Load, "kernel": s.Kernel,
					"rebootForKernel": s.KernelNext, "rebootForConfiguration": s.BootPending, "cpuPercent": finite(s.CPU),
					"memoryUsedBytes": s.MemUsed, "memoryTotalBytes": s.MemTotal, "temperatureC": finite(s.Temperature),
					"diskUsedBytes": s.DiskUsed, "diskTotalBytes": s.DiskTotal},
				"settings": map[string]any{"unapplied": st.Unapplied, "applying": st.Applying,
					"lastApplyResult": st.ApplyResult, "confirmationDueInSeconds": int(st.RollbackIn.Seconds())},
			}, nil
		},
	},
	{
		name: "get_settings", scope: scopeRead, schema: emptySchema,
		description: "The saved settings (settings.json: `system` and `router`) with the etag to pass to update_settings, and whether they differ from the confirmed (running) system. Secrets are not in it, only paths of files that hold them.",
		run: func(a *app, _ json.RawMessage) (any, error) {
			data, err := os.ReadFile(a.sys.settingsPath())
			if err != nil {
				return nil, err
			}
			var v any
			if err := json.Unmarshal(data, &v); err != nil {
				return nil, err
			}
			return map[string]any{"etag": etag(data), "unapplied": a.sys.Status().Unapplied, "settings": v}, nil
		},
	},
	{
		name: "get_pending_changes", scope: scopeRead, schema: emptySchema,
		description: "The difference between the confirmed settings and the saved ones (lines starting with - are removed, + added); empty when there is none.",
		run: func(a *app, _ json.RawMessage) (any, error) {
			if d := a.sys.pendingChanges(); d != "" {
				return d, nil
			}
			return "No changes: the saved settings are the confirmed ones.", nil
		},
	},
	{
		name: "list_interfaces", scope: scopeRead, schema: emptySchema,
		description: "Network interfaces: physical cards (name now, name from the settings, MAC address) and VLAN devices, and how the settings use each (WAN, a LAN).",
		run: func(a *app, _ json.RawMessage) (any, error) {
			s, err := loadSettings(a.sys.settingsPath())
			if err != nil {
				return nil, err
			}
			var out []map[string]any
			for _, d := range a.devices(s, s.form()) {
				e := map[string]any{"name": d.Name, "kind": d.Kind, "usedAs": d.Role}
				if d.Kernel != "" {
					e["kernelName"] = d.Kernel
				}
				if d.MAC != "" {
					e["mac"] = d.MAC
				}
				if d.Custom != "" {
					e["deviceName"] = d.Custom
				}
				out = append(out, e)
			}
			return out, nil
		},
	},
	{
		name: "list_dhcp_leases", scope: scopeRead, schema: emptySchema,
		description: "The DHCP clients currently holding a lease: address, MAC, hostname, LAN, when the lease ends.",
		run: func(a *app, _ json.RawMessage) (any, error) {
			s, err := loadSettings(a.sys.settingsPath())
			if err != nil {
				return nil, err
			}
			var out []map[string]any
			for _, l := range readLeases(a.metrics.leasesPath, s.form().LANs, time.Now()) {
				out = append(out, map[string]any{"address": l.Address, "mac": l.MAC, "hostname": l.Hostname, "lan": l.LAN, "expires": l.Expires})
			}
			return out, nil
		},
	},
	{
		name: "get_update_status", scope: scopeRead, schema: emptySchema,
		description: "The installed release (version, schema) and the state of the last update run.",
		run: func(a *app, _ json.RawMessage) (any, error) {
			st := a.sys.UpdateStatus()
			return map[string]any{"installed": a.sys.installedRelease(), "running": st.Running, "lastResult": st.Result, "log": st.Log}, nil
		},
	},
	{
		name: "get_apply_log", scope: scopeRead, schema: emptySchema,
		description: "The log of the last apply (build and activation) and its result.",
		run: func(a *app, _ json.RawMessage) (any, error) {
			st := a.sys.Status()
			return map[string]any{"applying": st.Applying, "result": st.ApplyResult, "log": st.Log}, nil
		},
	},
	logsTool,
	{
		name: "check_settings", scope: scopeRead,
		schema:      emptySchema,
		description: "Evaluates the saved settings the way an apply would (the router's assertions and the option types) without building or activating anything. Takes about a minute; if it is still running, call again to get the result. Reports whether the result is for the current saved settings.",
		run:         func(a *app, _ json.RawMessage) (any, error) { return a.checkSettings() },
	},
	{
		name: "update_settings", scope: scopeChange,
		schema: map[string]any{
			"type": "object",
			"properties": map[string]any{
				"etag":  map[string]any{"type": "string", "description": "The etag from get_settings: the change is refused if the settings changed since."},
				"patch": map[string]any{"type": "object", "description": "A JSON merge patch (RFC 7386) on the settings: objects merge, a value replaces, null removes a key, arrays are replaced whole."},
			},
			"required": []string{"etag", "patch"},
		},
		description: "Edit the saved settings with a JSON merge patch. The settings are saved but not applied: a person applies them in the web UI. Not allowed: ssh keys, the schema version, the state version, and file paths outside the router's state directory. Run check_settings afterwards.",
		run:         func(a *app, args json.RawMessage) (any, error) { return a.updateSettings(args) },
	},
	{
		name: "discard_changes", scope: scopeChange, schema: emptySchema,
		description: "Throw away the unapplied changes: the settings go back to the confirmed ones.",
		run: func(a *app, _ json.RawMessage) (any, error) {
			if err := a.sys.Discard(); err != nil {
				return nil, err
			}
			return "The settings are the confirmed ones again.", nil
		},
	},
}

// mergePatch applies a JSON merge patch (RFC 7386).
func mergePatch(target, patch any) any {
	p, ok := patch.(map[string]any)
	if !ok {
		return patch
	}
	t, ok := target.(map[string]any)
	if !ok {
		t = map[string]any{}
	}
	for k, v := range p {
		if v == nil {
			delete(t, k)
		} else {
			t[k] = mergePatch(t[k], v)
		}
	}
	return t
}

// checkMCPSettings refuses what an assistant may not change.
func (s System) checkMCPSettings(old, next Settings) error {
	for k := range next {
		if k != "version" && k != "system" && k != "router" {
			return fmt.Errorf("unknown top-level key %q (only system and router)", k)
		}
	}
	same := func(a, b any) bool { x, _ := json.Marshal(a); y, _ := json.Marshal(b); return string(x) == string(y) }
	if !same(old["version"], next["version"]) {
		return errors.New("`version` (the settings schema) cannot be changed")
	}
	oldSys, newSys := obj(Settings{"system": old["system"]}, "system"), obj(Settings{"system": next["system"]}, "system")
	for _, k := range []string{"sshKeys", "stateVersion"} {
		if !same(oldSys[k], newSys[k]) {
			return fmt.Errorf("`system.%s` cannot be changed through this interface", k)
		}
	}
	for k := range newSys {
		if k != "hostName" && k != "timeZone" && !same(oldSys[k], newSys[k]) {
			return fmt.Errorf("`system.%s` cannot be changed through this interface (only hostName and timeZone)", k)
		}
	}
	var walk func(path string, v any) error
	walk = func(path string, v any) error {
		switch x := v.(type) {
		case map[string]any:
			for k, e := range x {
				if str, ok := e.(string); ok && strings.HasSuffix(k, "File") {
					if !strings.HasPrefix(str, s.StateDir+"/") || strings.Contains(str, "..") {
						return fmt.Errorf("%s.%s: files must be under %s (the person who runs the router puts secrets there)", path, k, s.StateDir)
					}
				}
				if err := walk(path+"."+k, e); err != nil {
					return err
				}
			}
		case []any:
			for i, e := range x {
				if err := walk(fmt.Sprintf("%s[%d]", path, i), e); err != nil {
					return err
				}
			}
		}
		return nil
	}
	return walk("router", next["router"])
}

func (a *app) updateSettings(args json.RawMessage) (any, error) {
	var in struct {
		Etag  string `json:"etag"`
		Patch any    `json:"patch"`
	}
	if err := json.Unmarshal(args, &in); err != nil || in.Patch == nil {
		return nil, errors.New("give `etag` and `patch` (an object)")
	}
	if _, ok := in.Patch.(map[string]any); !ok {
		return nil, errors.New("`patch` must be an object")
	}
	return a.editSettings(in.Etag, func(next Settings) (Settings, error) {
		patched, _ := mergePatch(map[string]any(next), in.Patch).(map[string]any)
		return Settings(patched), nil
	})
}

// editSettings saves what edit makes of the settings, if etag is the current
// version of them (so nobody's edit is lost) and the result stays within what
// the API may change. The settings stay unapplied.
func (a *app) editSettings(wantEtag string, edit func(next Settings) (Settings, error)) (any, error) {
	a.settingsMu.Lock()
	defer a.settingsMu.Unlock()
	data, err := os.ReadFile(a.sys.settingsPath())
	if err != nil {
		return nil, err
	}
	if wantEtag != etag(data) {
		return nil, errStale
	}
	old, err := parseSettings(data)
	if err != nil {
		return nil, err
	}
	next, err := parseSettings(data) // a copy to edit
	if err != nil {
		return nil, err
	}
	next, err = edit(next)
	if err != nil {
		return nil, err
	}
	if err := a.sys.checkMCPSettings(old, next); err != nil {
		return nil, err
	}
	if err := next.save(a.sys.settingsPath()); err != nil {
		return nil, err
	}
	after, _ := os.ReadFile(a.sys.settingsPath())
	return map[string]any{"saved": true, "applied": false, "etag": etag(after),
		"changes": diffLines(pretty(data), pretty(after)),
		"next":    "Run check_settings (GET /api/v1/settings/check); a person applies the changes in the web UI."}, nil
}