package router

import (
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"strconv"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// agentHeader names the agent whose notes a request reads and writes.
//
// A namespace inside a project, not a credential: with auth on, the key decides
// the project and this only separates agents sharing it; with auth off it is
// just a name.
const agentHeader = "x-llmgw-agent"

const (
	defaultAgent   = "default"
	maxAgentID     = 100
	maxMemoryBody  = 64 << 10
	mcpProtocolVer = "2025-03-26"
)

// errMemoryOff is what a memory call meets when nothing is configured to keep
// notes. Refused rather than pretended: an agent that believes it saved
// something and finds nothing later has been lied to.
var errMemoryOff = errors.New(
	"this gateway keeps no memory notes: enable auth with a project that keeps memory, " +
		"or set threads.store and threads.retain_hours")

// memoryScope says whose notes this request is about.
func (r *Router) memoryScope(req *http.Request) (project, agent string, err error) {
	var caller string
	if p, ok := auth.FromContext(req.Context()); ok {
		caller = p.ProjectID
	}
	project, err = r.threadProject(caller)
	if err != nil {
		if errors.Is(err, errThreadsOff) {
			err = errMemoryOff
		}
		return "", "", err
	}
	agent = req.Header.Get(agentHeader)
	if agent == "" {
		agent = defaultAgent
	}
	if len(agent) > maxAgentID {
		return "", "", errors.New(agentHeader + " is too long")
	}
	return project, agent, nil
}

type noteJSON struct {
	Key       string   `json:"key"`
	Text      string   `json:"text"`
	Tags      []string `json:"tags"`
	CreatedAt string   `json:"created_at"`
	UpdatedAt string   `json:"updated_at"`
}

func toNoteJSON(n store.Note) noteJSON {
	tags := n.Tags
	if tags == nil {
		tags = []string{}
	}
	return noteJSON{
		Key: n.Key, Text: n.Text, Tags: tags,
		CreatedAt: n.CreatedAt.UTC().Format(time.RFC3339),
		UpdatedAt: n.UpdatedAt.UTC().Format(time.RFC3339),
	}
}

// memoryStatus maps a store error to what a client should be told.
func memoryStatus(err error) int {
	switch {
	case errors.Is(err, store.ErrNoNote):
		return http.StatusNotFound
	case errors.Is(err, errThreadsNeedAuth):
		return http.StatusUnauthorized
	case errors.Is(err, errMemoryOff), errors.Is(err, store.ErrMemoryNotEnabled),
		errors.Is(err, store.ErrMemoryFull):
		return http.StatusBadRequest
	default:
		return 0
	}
}

func writeMemoryError(w http.ResponseWriter, err error) {
	if code := memoryStatus(err); code != 0 {
		msg := err.Error()
		if code == http.StatusNotFound {
			// One answer for missing, another agent's and another project's, so
			// a caller cannot map which keys exist where they cannot read.
			msg = "no such note"
		}
		writeError(w, code, "invalid_request_error", msg)
		return
	}
	writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
}

// handlePutNote writes a note: PUT /v1/memory/{key}.
func (r *Router) handlePutNote(w http.ResponseWriter, req *http.Request) {
	project, agent, err := r.memoryScope(req)
	if err != nil {
		writeMemoryError(w, err)
		return
	}
	if r.noStore(callerSubject(req)) {
		writeErrorCode(w, http.StatusForbidden, "invalid_request_error", "no_store", errNoStoreMsg)
		return
	}
	var body struct {
		Text string   `json:"text"`
		Tags []string `json:"tags"`
	}
	raw, err := io.ReadAll(http.MaxBytesReader(w, req.Body, maxMemoryBody))
	if err != nil || json.Unmarshal(raw, &body) != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error",
			`body must be JSON like {"text": "...", "tags": ["..."]}`)
		return
	}
	note := store.Note{Key: req.PathValue("key"), Text: body.Text, Tags: body.Tags}
	if err := r.threads.PutNote(req.Context(), project, agent, note); err != nil {
		writeMemoryError(w, err)
		return
	}
	saved, err := r.threads.GetNote(req.Context(), project, agent, note.Key)
	if err != nil {
		writeMemoryError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, toNoteJSON(saved))
}

// handleGetNote reads one note: GET /v1/memory/{key}.
func (r *Router) handleGetNote(w http.ResponseWriter, req *http.Request) {
	project, agent, err := r.memoryScope(req)
	if err != nil {
		writeMemoryError(w, err)
		return
	}
	n, err := r.threads.GetNote(req.Context(), project, agent, req.PathValue("key"))
	if err != nil {
		writeMemoryError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, toNoteJSON(n))
}

// handleSearchNotes lists or searches: GET /v1/memory?q=&tag=&limit=.
func (r *Router) handleSearchNotes(w http.ResponseWriter, req *http.Request) {
	project, agent, err := r.memoryScope(req)
	if err != nil {
		writeMemoryError(w, err)
		return
	}
	q := req.URL.Query()
	limit, _ := strconv.Atoi(q.Get("limit"))
	notes, err := r.threads.SearchNotes(req.Context(), project, agent, q.Get("q"), q.Get("tag"), limit)
	if err != nil {
		writeMemoryError(w, err)
		return
	}
	out := make([]noteJSON, 0, len(notes))
	for _, n := range notes {
		out = append(out, toNoteJSON(n))
	}
	writeJSON(w, http.StatusOK, map[string]any{"object": "list", "data": out})
}

// handleDeleteNote forgets one: DELETE /v1/memory/{key}.
func (r *Router) handleDeleteNote(w http.ResponseWriter, req *http.Request) {
	project, agent, err := r.memoryScope(req)
	if err != nil {
		writeMemoryError(w, err)
		return
	}
	if err := r.threads.DeleteNote(req.Context(), project, agent, req.PathValue("key")); err != nil {
		writeMemoryError(w, err)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

// --- MCP ---
//
// The same notes, offered as tools an agent can be given. Only the small part
// of the protocol a tool server needs: initialize, tools/list and tools/call,
// answered as plain JSON over POST. It never streams and never calls the
// client, so it needs no session.

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

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"`
}

var mcpTools = []map[string]any{
	{
		"name":        "remember",
		"description": "Save a note that will still be there in later conversations. Writing to an existing key replaces it.",
		"inputSchema": map[string]any{
			"type":     "object",
			"required": []string{"key", "text"},
			"properties": map[string]any{
				"key":  map[string]any{"type": "string", "description": "Short unique name for the note"},
				"text": map[string]any{"type": "string", "description": "What to remember"},
				"tags": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "Optional lowercase labels"},
			},
		},
	},
	{
		"name":        "recall",
		"description": "Find saved notes by words, by tag, or by exact key. With nothing given, returns the most recent notes.",
		"inputSchema": map[string]any{
			"type": "object",
			"properties": map[string]any{
				"query": map[string]any{"type": "string", "description": "Words to search for"},
				"tag":   map[string]any{"type": "string", "description": "Only notes with this tag"},
				"key":   map[string]any{"type": "string", "description": "Exact key to read"},
				"limit": map[string]any{"type": "integer", "description": "Most notes to return (up to 100)"},
			},
		},
	},
	{
		"name":        "forget",
		"description": "Delete a saved note.",
		"inputSchema": map[string]any{
			"type":       "object",
			"required":   []string{"key"},
			"properties": map[string]any{"key": map[string]any{"type": "string"}},
		},
	},
}

// handleMCP serves POST /mcp.
func (r *Router) handleMCP(w http.ResponseWriter, req *http.Request) {
	raw, err := io.ReadAll(http.MaxBytesReader(w, req.Body, maxMemoryBody))
	var rpc rpcRequest
	if err != nil || json.Unmarshal(raw, &rpc) != nil || rpc.Method == "" {
		writeJSON(w, http.StatusBadRequest, rpcResponse{
			JSONRPC: "2.0", ID: json.RawMessage("null"),
			Error: &rpcError{Code: -32600, Message: "expected one JSON-RPC request object"},
		})
		return
	}
	// A notification has no id and gets no answer.
	if len(rpc.ID) == 0 {
		w.WriteHeader(http.StatusAccepted)
		return
	}

	resp := rpcResponse{JSONRPC: "2.0", ID: rpc.ID}
	switch rpc.Method {
	case "initialize":
		resp.Result = map[string]any{
			"protocolVersion": mcpProtocolVer,
			"capabilities":    map[string]any{"tools": map[string]any{}},
			"serverInfo":      map[string]any{"name": "llmgw-memory", "version": "1"},
		}
	case "ping":
		resp.Result = map[string]any{}
	case "tools/list":
		resp.Result = map[string]any{"tools": mcpTools}
	case "tools/call":
		resp.Result = r.callMemoryTool(req, rpc.Params)
	default:
		resp.Error = &rpcError{Code: -32601, Message: "method not found: " + rpc.Method}
	}
	writeJSON(w, http.StatusOK, resp)
}

func toolText(text string, isError bool) map[string]any {
	return map[string]any{
		"content": []map[string]any{{"type": "text", "text": text}},
		"isError": isError,
	}
}

func toolJSON(v any) map[string]any {
	b, _ := json.MarshalIndent(v, "", "  ")
	return toolText(string(b), false)
}

// callMemoryTool runs one tool. Failures are tool results with isError set, not
// protocol errors, because that is how a model gets to see what went wrong.
func (r *Router) callMemoryTool(req *http.Request, params json.RawMessage) map[string]any {
	var call struct {
		Name      string `json:"name"`
		Arguments struct {
			Key   string   `json:"key"`
			Text  string   `json:"text"`
			Tags  []string `json:"tags"`
			Query string   `json:"query"`
			Tag   string   `json:"tag"`
			Limit int      `json:"limit"`
		} `json:"arguments"`
	}
	if err := json.Unmarshal(params, &call); err != nil {
		return toolText("invalid tool call", true)
	}
	project, agent, err := r.memoryScope(req)
	if err != nil {
		return toolText(err.Error(), true)
	}
	a, ctx := call.Arguments, req.Context()

	switch call.Name {
	case "remember":
		if r.noStore(callerSubject(req)) {
			return toolText(errNoStoreMsg, true)
		}
		if err := r.threads.PutNote(ctx, project, agent, store.Note{Key: a.Key, Text: a.Text, Tags: a.Tags}); err != nil {
			return toolText(err.Error(), true)
		}
		return toolText("saved "+a.Key, false)
	case "recall":
		if a.Key != "" {
			n, err := r.threads.GetNote(ctx, project, agent, a.Key)
			if errors.Is(err, store.ErrNoNote) {
				return toolText("no note with key "+a.Key, false)
			}
			if err != nil {
				return toolText(err.Error(), true)
			}
			return toolJSON([]noteJSON{toNoteJSON(n)})
		}
		notes, err := r.threads.SearchNotes(ctx, project, agent, a.Query, a.Tag, a.Limit)
		if err != nil {
			return toolText(err.Error(), true)
		}
		out := make([]noteJSON, 0, len(notes))
		for _, n := range notes {
			out = append(out, toNoteJSON(n))
		}
		return toolJSON(out)
	case "forget":
		if err := r.threads.DeleteNote(ctx, project, agent, a.Key); err != nil {
			if errors.Is(err, store.ErrNoNote) {
				return toolText("no note with key "+a.Key, false)
			}
			return toolText(err.Error(), true)
		}
		return toolText("forgot "+a.Key, false)
	default:
		return toolText("unknown tool "+call.Name, true)
	}
}