package middleware

import (
	"context"
	"errors"
	"testing"

	mcpContext "gitea.com/gitea/gitea-mcp/pkg/context"
	"gitea.com/gitea/gitea-mcp/pkg/debug"
	"gitea.com/gitea/gitea-mcp/pkg/flag"

	"github.com/mark3labs/mcp-go/mcp"
)

func TestToolLogging_PropagatesContextValues(t *testing.T) {
	origDebug := flag.Debug
	defer func() { flag.Debug = origDebug }()
	flag.Debug = false // exercise the no-log code path; context wiring runs regardless

	var (
		gotToolName string
		gotReqID    string
	)
	handler := func(ctx context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		gotToolName = debug.ToolNameFromContext(ctx)
		gotReqID = debug.RequestIDFromContext(ctx)
		return mcp.NewToolResultText("ok"), nil
	}

	wrapped := ToolLogging(handler)

	req := mcp.CallToolRequest{}
	req.Params.Name = "get_me"
	req.Params.Arguments = map[string]any{"x": 1}

	res, err := wrapped(context.Background(), req)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if res == nil {
		t.Fatal("expected non-nil result")
	}
	if gotToolName != "get_me" {
		t.Fatalf("tool name not propagated, got %q", gotToolName)
	}
	if gotReqID == "" {
		t.Fatal("request id not propagated")
	}
}

func TestToolLogging_ReturnsHandlerError(t *testing.T) {
	want := errors.New("boom")
	wrapped := ToolLogging(func(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		return nil, want
	})

	req := mcp.CallToolRequest{}
	req.Params.Name = "broken_tool"

	_, err := wrapped(context.Background(), req)
	if !errors.Is(err, want) {
		t.Fatalf("error not surfaced, got %v", err)
	}
}

func TestTokenSourceFromContext(t *testing.T) {
	origFlag := flag.TokenSource
	defer func() { flag.TokenSource = origFlag }()

	t.Run("falls back to flag.TokenSource", func(t *testing.T) {
		flag.TokenSource = "env"
		if got := tokenSourceFromContext(context.Background()); got != "env" {
			t.Fatalf("got %q want env", got)
		}
	})

	t.Run("returns none when no source is known", func(t *testing.T) {
		flag.TokenSource = ""
		if got := tokenSourceFromContext(context.Background()); got != "none" {
			t.Fatalf("got %q want none", got)
		}
	})

	t.Run("context value beats flag", func(t *testing.T) {
		flag.TokenSource = "env"
		ctx := context.WithValue(context.Background(), mcpContext.TokenSourceContextKey, "header")
		if got := tokenSourceFromContext(ctx); got != "header" {
			t.Fatalf("got %q want header", got)
		}
	})
}