// Package gateway implements the agent-facing edge: the single MCP endpoint
// every AI agent connects to. It is an MCP *server* to agents (Streamable HTTP
// + JSON-RPC 2.0) and delegates tool aggregation/dispatch to the Router.
package gateway

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"sync/atomic"
	"time"

	"gitea.thetempleofdoom.com/drjones/mcp-gateway-nexus/internal/mcp"
	"gitea.thetempleofdoom.com/drjones/mcp-gateway-nexus/internal/router"
)

// Gateway serves the aggregated MCP endpoint plus health/metrics.
type Gateway struct {
	info    mcp.Implementation
	mcpPath string
	router  Router
	log     *slog.Logger

	metrics gwMetrics
}

// Router is the subset of *router.Router the gateway depends on (kept small so
// it is trivially mockable in tests).
type Router interface {
	ListTools() []mcp.ToolDefinition
	CallTool(ctx context.Context, qualified string, args json.RawMessage) (json.RawMessage, *mcp.RPCError, error)
	Statuses() []router.Status
}

type gwMetrics struct {
	requests  atomic.Int64
	toolCalls atomic.Int64
	toolErrs  atomic.Int64
}

// Options configures a Gateway.
type Options struct {
	Name    string
	Version string
	MCPPath string
	Router  Router
	Logger  *slog.Logger
}

// New creates a Gateway.
func New(opts Options) *Gateway {
	log := opts.Logger
	if log == nil {
		log = slog.Default()
	}
	path := opts.MCPPath
	if path == "" {
		path = "/mcp"
	}
	return &Gateway{
		info:    mcp.Implementation{Name: opts.Name, Version: opts.Version},
		mcpPath: path,
		router:  opts.Router,
		log:     log,
	}
}

// Handler returns the HTTP handler exposing the MCP endpoint and ops routes.
func (g *Gateway) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc(g.mcpPath, g.handleMCP)
	mux.HandleFunc("/healthz", g.handleHealth)
	mux.HandleFunc("/metrics", g.handleMetrics)
	mux.HandleFunc("/", g.handleRoot)
	return mux
}

func (g *Gateway) handleRoot(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.NotFound(w, r)
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{
		"name":     g.info.Name,
		"version":  g.info.Version,
		"mcp":      g.mcpPath,
		"protocol": mcp.ProtocolVersion,
	})
}

func (g *Gateway) handleHealth(w http.ResponseWriter, _ *http.Request) {
	statuses := g.router.Statuses()
	ready := 0
	for _, s := range statuses {
		if s.Ready {
			ready++
		}
	}
	writeJSON(w, http.StatusOK, map[string]any{
		"status":            "ok",
		"upstreams_total":   len(statuses),
		"upstreams_ready":   ready,
		"tools":             len(g.router.ListTools()),
		"upstream_statuses": statuses,
	})
}

// handleMetrics exposes a minimal Prometheus text exposition. A full
// client_golang integration lands with the Metrics module (Phase 4).
func (g *Gateway) handleMetrics(w http.ResponseWriter, _ *http.Request) {
	w.Header().Set("Content-Type", "text/plain; version=0.0.4")
	fmt.Fprintf(w, "# HELP nexus_gateway_requests_total Total MCP requests received.\n")
	fmt.Fprintf(w, "# TYPE nexus_gateway_requests_total counter\n")
	fmt.Fprintf(w, "nexus_gateway_requests_total %d\n", g.metrics.requests.Load())
	fmt.Fprintf(w, "# HELP nexus_gateway_tool_calls_total Total tools/call requests.\n")
	fmt.Fprintf(w, "# TYPE nexus_gateway_tool_calls_total counter\n")
	fmt.Fprintf(w, "nexus_gateway_tool_calls_total %d\n", g.metrics.toolCalls.Load())
	fmt.Fprintf(w, "# HELP nexus_gateway_tool_errors_total Total failed tools/call requests.\n")
	fmt.Fprintf(w, "# TYPE nexus_gateway_tool_errors_total counter\n")
	fmt.Fprintf(w, "nexus_gateway_tool_errors_total %d\n", g.metrics.toolErrs.Load())
	fmt.Fprintf(w, "# HELP nexus_gateway_tools Registered tools currently exposed.\n")
	fmt.Fprintf(w, "# TYPE nexus_gateway_tools gauge\n")
	fmt.Fprintf(w, "nexus_gateway_tools %d\n", len(g.router.ListTools()))
}

// handleMCP is the single MCP endpoint. It accepts a JSON-RPC request or
// notification and responds with application/json.
func (g *Gateway) handleMCP(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		w.Header().Set("Allow", http.MethodPost)
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	g.metrics.requests.Add(1)

	body, err := io.ReadAll(io.LimitReader(r.Body, 32<<20))
	if err != nil {
		writeRPC(w, mcp.NewError(nil, mcp.CodeParseError, "read body: "+err.Error()))
		return
	}
	var msg mcp.Message
	if err := json.Unmarshal(body, &msg); err != nil {
		writeRPC(w, mcp.NewError(nil, mcp.CodeParseError, "invalid JSON-RPC: "+err.Error()))
		return
	}

	// Notifications get 202 with no body.
	if msg.IsNotification() {
		w.WriteHeader(http.StatusAccepted)
		return
	}
	resp := g.dispatch(r.Context(), &msg)
	writeRPC(w, resp)
}

func (g *Gateway) dispatch(ctx context.Context, msg *mcp.Message) *mcp.Message {
	switch msg.Method {
	case mcp.MethodInitialize:
		res := mcp.InitializeResult{
			ProtocolVersion: mcp.ProtocolVersion,
			Capabilities:    mcp.Capabilities{Tools: &mcp.ToolsCapability{ListChanged: true}},
			ServerInfo:      g.info,
			Instructions:    "MCP Nexus aggregates every upstream MCP server behind one endpoint. Tools are namespaced as {namespace}.{tool}.",
		}
		out, _ := mcp.NewResult(msg.ID, res)
		return out

	case mcp.MethodPing:
		out, _ := mcp.NewResult(msg.ID, struct{}{})
		return out

	case mcp.MethodToolsList:
		res := mcp.ListToolsResult{Tools: g.router.ListTools()}
		out, _ := mcp.NewResult(msg.ID, res)
		return out

	case mcp.MethodToolsCall:
		return g.handleToolCall(ctx, msg)

	default:
		return mcp.NewError(msg.ID, mcp.CodeMethodNotFound, "unknown method: "+msg.Method)
	}
}

func (g *Gateway) handleToolCall(ctx context.Context, msg *mcp.Message) *mcp.Message {
	g.metrics.toolCalls.Add(1)
	var params mcp.CallToolParams
	if err := msg.UnmarshalParams(&params); err != nil {
		g.metrics.toolErrs.Add(1)
		return mcp.NewError(msg.ID, mcp.CodeInvalidParams, "invalid params: "+err.Error())
	}
	if params.Name == "" {
		g.metrics.toolErrs.Add(1)
		return mcp.NewError(msg.ID, mcp.CodeInvalidParams, "tool name is required")
	}

	start := time.Now()
	raw, rpcErr, err := g.router.CallTool(ctx, params.Name, params.Arguments)
	if err != nil {
		g.metrics.toolErrs.Add(1)
		g.log.Warn("tool call transport error", "tool", params.Name, "error", err)
		return mcp.NewError(msg.ID, mcp.CodeInternalError, err.Error())
	}
	if rpcErr != nil {
		g.metrics.toolErrs.Add(1)
		return &mcp.Message{JSONRPC: mcp.JSONRPCVersion, ID: msg.ID, Error: rpcErr}
	}
	g.log.Info("tool call", "tool", params.Name, "ms", time.Since(start).Milliseconds())
	// Pass the upstream MCP result through unmodified.
	return &mcp.Message{JSONRPC: mcp.JSONRPCVersion, ID: msg.ID, Result: raw}
}

func writeRPC(w http.ResponseWriter, msg *mcp.Message) {
	writeJSON(w, http.StatusOK, msg)
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}