package server

import (
	"context"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/modelcontextprotocol/go-sdk/mcp"
	log "github.com/nacorid/logger"
)

type LoggingMiddleware struct{}

func NewLoggingMiddleware() *LoggingMiddleware {
	return &LoggingMiddleware{}
}

// OnCall logs native MCP tool requests streamed over SSE/WebSockets
func (lm *LoggingMiddleware) OnCall(next mcp.MethodHandler) mcp.MethodHandler {
	return func(ctx context.Context, method string, request mcp.Request) (mcp.Result, error) {
		if method == "tools/call" {
			start := time.Now()

			result, err := next(ctx, method, request)

			duration := fmt.Sprintf("%.2fms", float64(time.Since(start))/float64(time.Millisecond))

			toolName := "unknown"
			if callReq, ok := request.(*mcp.CallToolRequest); ok && callReq.Params != nil {
				toolName = callReq.Params.Name
			}

			log.InfoWithContext(ctx, "MCP Tool Processed",
				"transport", "mcp_native",
				"tool", toolName,
				"duration", duration,
				"success", err == nil,
			)
			return result, err
		}
		return next(ctx, method, request)
	}
}

func (lm *LoggingMiddleware) OnCallHTTP(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()

		ipAddr := r.RemoteAddr
		if forwarded := r.Header.Get("X-Forwarded-For"); forwarded != "" {
			ipAddr = strings.TrimSpace(strings.Split(forwarded, ",")[0])
		}

		next.ServeHTTP(w, r)

		duration := fmt.Sprintf("%.2fms", float64(time.Since(start))/float64(time.Millisecond))
		log.InfoWithContext(r.Context(), "REST Tool Processed",
			"transport", "rest",
			"path", r.URL.Path,
			"duration", duration,
			"remote_ip", ipAddr,
		)
	})
}