package registry

import (
	"fmt"
	"testing"

	"github.com/modelcontextprotocol/go-sdk/mcp"

	"github.com/txn2/mcp-data-platform/pkg/query"
	"github.com/txn2/mcp-data-platform/pkg/semantic"
)

const (
	regTestTrino     = "trino"
	regTestProd      = "prod"
	regTestTest      = "test"
	regTestToolCount = 3
)

// mockToolkit is a simple mock for testing.
type mockToolkit struct {
	kind       string
	name       string
	connection string
	tools      []string
	closeCalls int
}

func (m *mockToolkit) Kind() string                            { return m.kind }
func (m *mockToolkit) Name() string                            { return m.name }
func (m *mockToolkit) Connection() string                      { return m.connection }
func (m *mockToolkit) RegisterTools(_ *mcp.Server)             {} //nolint:revive // unused-receiver: mock
func (m *mockToolkit) Tools() []string                         { return m.tools }
func (m *mockToolkit) SetSemanticProvider(_ semantic.Provider) {} //nolint:revive // unused-receiver: mock
func (m *mockToolkit) SetQueryProvider(_ query.Provider)       {} //nolint:revive // unused-receiver: mock
func (m *mockToolkit) Close() error                            { m.closeCalls++; return nil }

// mockToolkitWithCloseError is a toolkit that returns an error on Close.
type mockToolkitWithCloseError struct {
	mockToolkit
}

func (m *mockToolkitWithCloseError) Close() error { //nolint:revive // unused-receiver: mock
	return fmt.Errorf("close error")
}

func TestRegistry_RegisterAndGet(t *testing.T) {
	reg := NewRegistry()
	toolkit := &mockToolkit{kind: regTestTrino, name: regTestProd}

	if err := reg.Register(toolkit); err != nil {
		t.Fatalf("Register() error = %v", err)
	}

	got, ok := reg.Get(regTestTrino, "prod")
	if !ok {
		t.Fatal("Get() returned false")
	}
	if got.Kind() != regTestTrino {
		t.Errorf("Kind() = %q, want %q", got.Kind(), regTestTrino)
	}
}

func TestRegistry_RegisterDuplicate(t *testing.T) {
	reg := NewRegistry()
	toolkit := &mockToolkit{kind: regTestTrino, name: regTestProd}

	_ = reg.Register(toolkit)
	err := reg.Register(toolkit)
	if err == nil {
		t.Error("Register() expected error for duplicate")
	}
}

func TestRegistry_GetNotFound(t *testing.T) {
	reg := NewRegistry()
	_, ok := reg.Get("nonexistent", "name")
	if ok {
		t.Error("Get() returned true for nonexistent toolkit")
	}
}

func TestRegistry_GetByKind(t *testing.T) {
	reg := NewRegistry()
	_ = reg.Register(&mockToolkit{kind: regTestTrino, name: regTestProd})
	_ = reg.Register(&mockToolkit{kind: regTestTrino, name: "staging"})
	_ = reg.Register(&mockToolkit{kind: "datahub", name: "main"})

	trinoToolkits := reg.GetByKind(regTestTrino)
	if len(trinoToolkits) != 2 {
		t.Errorf("GetByKind(trino) returned %d toolkits, want 2", len(trinoToolkits))
	}
}

func TestRegistry_AllAndAllTools(t *testing.T) {
	reg := NewRegistry()
	_ = reg.Register(&mockToolkit{kind: regTestTrino, name: regTestProd, tools: []string{"trino_query", "trino_describe"}})
	_ = reg.Register(&mockToolkit{kind: "datahub", name: "main", tools: []string{"datahub_search"}})

	all := reg.All()
	if len(all) != 2 {
		t.Errorf("All() returned %d toolkits, want 2", len(all))
	}

	tools := reg.AllTools()
	if len(tools) != regTestToolCount {
		t.Errorf("AllTools() returned %d tools, want 3", len(tools))
	}
}

func TestRegistry_Close(t *testing.T) {
	reg := NewRegistry()
	toolkit := &mockToolkit{kind: regTestTrino, name: regTestProd}
	_ = reg.Register(toolkit)

	if err := reg.Close(); err != nil {
		t.Errorf("Close() error = %v", err)
	}
	if toolkit.closeCalls != 1 {
		t.Errorf("closeCalls = %d, want 1", toolkit.closeCalls)
	}
}

func TestRegistry_CloseWithError(t *testing.T) {
	reg := NewRegistry()
	toolkit := &mockToolkitWithCloseError{mockToolkit: mockToolkit{kind: regTestTrino, name: regTestProd}}
	_ = reg.Register(toolkit)

	err := reg.Close()
	if err == nil {
		t.Error("Close() expected error when toolkit fails")
	}
}

func TestRegistry_Providers(_ *testing.T) {
	reg := NewRegistry()
	toolkit := &mockToolkit{kind: regTestTrino, name: regTestProd}
	_ = reg.Register(toolkit)

	reg.SetSemanticProvider(semantic.NewNoopProvider())
	reg.SetQueryProvider(query.NewNoopProvider())
	// Just verify it doesn't panic
}

func TestRegistry_RegisterAllTools(_ *testing.T) {
	reg := NewRegistry()
	_ = reg.Register(&mockToolkit{kind: regTestTrino, name: regTestProd, tools: []string{"trino_query"}})
	_ = reg.Register(&mockToolkit{kind: "datahub", name: "main", tools: []string{"datahub_search"}})

	server := mcp.NewServer(&mcp.Implementation{Name: "test", Version: "1.0.0"}, nil)
	reg.RegisterAllTools(server) // Should not panic
}

func TestRegistry_RegisterWithPresetProviders(t *testing.T) {
	reg := NewRegistry()
	reg.SetSemanticProvider(semantic.NewNoopProvider())
	reg.SetQueryProvider(query.NewNoopProvider())

	toolkit := &mockToolkit{kind: regTestTrino, name: regTestProd}
	if err := reg.Register(toolkit); err != nil {
		t.Fatalf("Register() error = %v", err)
	}

	got, ok := reg.Get(regTestTrino, "prod")
	if !ok {
		t.Fatal("Get() returned false")
	}
	if got.Kind() != regTestTrino {
		t.Errorf("Kind() = %q, want %q", got.Kind(), regTestTrino)
	}
}

func TestRegistry_Factory(t *testing.T) {
	reg := NewRegistry()
	factory := func(name string, _ map[string]any) (Toolkit, error) {
		return &mockToolkit{kind: "custom", name: name}, nil
	}
	reg.RegisterFactory("custom", factory)

	err := reg.CreateAndRegister(ToolkitConfig{
		Kind:   "custom",
		Name:   regTestTest,
		Config: map[string]any{},
	})
	if err != nil {
		t.Fatalf("CreateAndRegister() error = %v", err)
	}

	_, ok := reg.Get("custom", regTestTest)
	if !ok {
		t.Error("Get() returned false after CreateAndRegister")
	}
}

func TestRegistry_FactoryError(t *testing.T) {
	reg := NewRegistry()
	factory := func(_ string, _ map[string]any) (Toolkit, error) {
		return nil, fmt.Errorf("factory error")
	}
	reg.RegisterFactory("failing", factory)

	err := reg.CreateAndRegister(ToolkitConfig{
		Kind:   "failing",
		Name:   regTestTest,
		Config: map[string]any{},
	})
	if err == nil {
		t.Error("CreateAndRegister() expected error when factory fails")
	}
}

func TestRegistry_UnknownKind(t *testing.T) {
	reg := NewRegistry()

	err := reg.CreateAndRegister(ToolkitConfig{
		Kind:   "unknown",
		Name:   regTestTest,
		Config: map[string]any{},
	})
	if err == nil {
		t.Error("CreateAndRegister() expected error for unknown kind")
	}
}

func TestRegisterBuiltinFactories(t *testing.T) {
	reg := NewRegistry()
	RegisterBuiltinFactories(reg)

	t.Run("trino aggregate factory registered", func(t *testing.T) {
		_, ok := reg.GetAggregateFactory(regTestTrino)
		if !ok {
			t.Error("expected trino aggregate factory to be registered")
		}
	})

	t.Run("datahub factory registered", func(t *testing.T) {
		// Should fail with invalid config (missing url)
		err := reg.CreateAndRegister(ToolkitConfig{
			Kind:   "datahub",
			Name:   regTestTest,
			Config: map[string]any{},
		})
		if err == nil {
			t.Error("expected error for missing datahub config")
		}
	})

	t.Run("s3 factory registered", func(_ *testing.T) {
		// S3 factory is registered, try to create (may succeed with AWS defaults)
		_ = reg.CreateAndRegister(ToolkitConfig{
			Kind:   "s3",
			Name:   regTestTest,
			Config: map[string]any{},
		})
		// Just verify the factory is called - actual creation depends on AWS SDK defaults
	})

	t.Run("mcp (gateway) aggregate factory registered", func(t *testing.T) {
		_, ok := reg.GetAggregateFactory("mcp")
		if !ok {
			t.Error("expected mcp aggregate factory to be registered (gateway feature)")
		}
	})
}

func TestTrinoFactory(t *testing.T) {
	// Test with invalid config
	_, err := TrinoFactory(regTestTest, map[string]any{})
	if err == nil {
		t.Error("TrinoFactory() expected error for missing host")
	}
}

func TestRegisterAggregateFactory(t *testing.T) {
	reg := NewRegistry()

	called := false
	factory := func(defaultName string, instances map[string]map[string]any) (Toolkit, error) {
		called = true
		if defaultName != "primary" {
			t.Errorf("defaultName = %q, want 'primary'", defaultName)
		}
		if len(instances) != 2 {
			t.Errorf("expected 2 instances, got %d", len(instances))
		}
		return &mockToolkit{kind: "agg", name: defaultName, tools: []string{"agg_tool"}}, nil
	}

	reg.RegisterAggregateFactory("agg", factory)

	got, ok := reg.GetAggregateFactory("agg")
	if !ok {
		t.Fatal("expected aggregate factory to be registered")
	}

	tk, err := got("primary", map[string]map[string]any{
		"primary":   {"host": "a"},
		"secondary": {"host": "b"},
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if !called {
		t.Error("factory was not called")
	}
	if tk.Name() != "primary" {
		t.Errorf("Name() = %q, want 'primary'", tk.Name())
	}
}

func TestGetAggregateFactory_NotFound(t *testing.T) {
	reg := NewRegistry()
	_, ok := reg.GetAggregateFactory("nonexistent")
	if ok {
		t.Error("expected false for unregistered aggregate factory")
	}
}

func TestTrinoAggregateFactory(t *testing.T) {
	t.Run("valid multi-instance config", func(t *testing.T) {
		tk, err := TrinoAggregateFactory("warehouse", map[string]map[string]any{
			"warehouse": {
				"host": "warehouse.example.com",
				"user": "trino",
				"port": 8080,
			},
			"elasticsearch": {
				"host": "es.example.com",
				"user": "trino",
				"port": 8080,
			},
		})
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		if tk.Kind() != "trino" {
			t.Errorf("Kind() = %q, want 'trino'", tk.Kind())
		}
		if tk.Name() != "warehouse" {
			t.Errorf("Name() = %q, want 'warehouse'", tk.Name())
		}
		if tk.Connection() != "warehouse" {
			t.Errorf("Connection() = %q, want 'warehouse'", tk.Connection())
		}
	})

	t.Run("invalid config returns error", func(t *testing.T) {
		_, err := TrinoAggregateFactory("bad", map[string]map[string]any{
			"bad": {"timeout": "invalid-duration"},
		})
		if err == nil {
			t.Error("expected error for invalid config")
		}
	})

	t.Run("missing host returns error", func(t *testing.T) {
		_, err := TrinoAggregateFactory("empty", map[string]map[string]any{
			"empty": {},
		})
		if err == nil {
			t.Error("expected error for missing host")
		}
	})
}

func TestDataHubFactory(t *testing.T) {
	// Test with invalid config
	_, err := DataHubFactory(regTestTest, map[string]any{})
	if err == nil {
		t.Error("DataHubFactory() expected error for missing url")
	}
}

func TestGatewayAggregateFactory_NoInstancesReturnsEmptyToolkit(t *testing.T) {
	tk, err := GatewayAggregateFactory(regTestTest, map[string]map[string]any{})
	if err != nil {
		t.Fatalf("GatewayAggregateFactory: %v", err)
	}
	if tk.Kind() != "mcp" {
		t.Errorf("Kind: got %q, want %q", tk.Kind(), "mcp")
	}
	if got := tk.Tools(); len(got) != 0 {
		t.Errorf("expected empty Tools() for no instances, got %v", got)
	}
	_ = tk.Close()
}

func TestGatewayAggregateFactory_BadInstanceConfigReturnsError(t *testing.T) {
	_, err := GatewayAggregateFactory(regTestTest, map[string]map[string]any{
		"broken": {}, // missing endpoint
	})
	if err == nil {
		t.Fatal("expected parse error for missing endpoint")
	}
}

func TestAPIGatewayAggregateFactory_NoInstancesReturnsEmptyToolkit(t *testing.T) {
	tk, err := APIGatewayAggregateFactory(regTestTest, map[string]map[string]any{})
	if err != nil {
		t.Fatalf("APIGatewayAggregateFactory: %v", err)
	}
	if tk.Kind() != "api" {
		t.Errorf("Kind: got %q, want %q", tk.Kind(), "api")
	}
	if got := tk.Tools(); len(got) != 3 ||
		got[0] != "api_invoke_endpoint" ||
		got[1] != "api_list_endpoints" ||
		got[2] != "api_get_endpoint_schema" {
		t.Errorf("expected [api_invoke_endpoint api_list_endpoints api_get_endpoint_schema], got %v", got)
	}
	_ = tk.Close()
}

func TestAPIGatewayAggregateFactory_BadInstanceConfigReturnsError(t *testing.T) {
	_, err := APIGatewayAggregateFactory(regTestTest, map[string]map[string]any{
		"broken": {}, // missing base_url
	})
	if err == nil {
		t.Fatal("expected parse error for missing base_url")
	}
}

func TestAPIGatewayAggregateFactory_LoadsValidInstance(t *testing.T) {
	tk, err := APIGatewayAggregateFactory(regTestTest, map[string]map[string]any{
		"acme": {
			"base_url":   "https://api.example.com",
			"auth_mode":  "bearer",
			"credential": "tok",
		},
	})
	if err != nil {
		t.Fatalf("APIGatewayAggregateFactory: %v", err)
	}
	if tk.Kind() != "api" {
		t.Errorf("Kind: got %q, want %q", tk.Kind(), "api")
	}
	_ = tk.Close()
}

func TestGatewayAggregateFactory_UnreachableInstanceAbsorbed(t *testing.T) {
	// Unreachable endpoint must not fail the factory — the toolkit still
	// constructs and the platform startup continues. The failed initial
	// connection is not stored.
	tk, err := GatewayAggregateFactory(regTestTest, map[string]map[string]any{
		"broken": {
			"endpoint":        "http://127.0.0.1:1/mcp",
			"connect_timeout": "250ms",
			"call_timeout":    "1s",
		},
	})
	if err != nil {
		t.Fatalf("expected no error for unreachable instance, got %v", err)
	}
	if got := tk.Tools(); len(got) != 0 {
		t.Errorf("expected empty Tools(), got %v", got)
	}
	_ = tk.Close()
}

func TestS3Factory(_ *testing.T) {
	// S3Factory may succeed with AWS SDK defaults (env vars, IAM role, etc.)
	// Just verify it can be called
	_, _ = S3Factory(regTestTest, map[string]any{})
}

func TestGetToolkitForTool_Found(t *testing.T) {
	reg := NewRegistry()
	_ = reg.Register(&mockToolkit{
		kind:       regTestTrino,
		name:       "production",
		connection: "prod-trino",
		tools:      []string{"trino_query", "trino_describe"},
	})

	match := reg.GetToolkitForTool("trino_query")
	assertToolMatch(t, match, ToolkitMatch{Kind: regTestTrino, Name: "production", Connection: "prod-trino", Found: true})
}

func TestGetToolkitForTool_NotFound(t *testing.T) {
	reg := NewRegistry()
	_ = reg.Register(&mockToolkit{
		kind:       regTestTrino,
		name:       "production",
		connection: "prod-trino",
		tools:      []string{"trino_query"},
	})

	match := reg.GetToolkitForTool("unknown_tool")
	assertToolMatch(t, match, ToolkitMatch{})
}

func TestGetToolkitForTool_MultipleToolkits(t *testing.T) {
	reg := NewRegistry()
	_ = reg.Register(&mockToolkit{
		kind: regTestTrino, name: "production", connection: "prod-trino",
		tools: []string{"trino_query", "trino_describe"},
	})
	_ = reg.Register(&mockToolkit{
		kind: "datahub", name: "main", connection: "main-datahub",
		tools: []string{"datahub_search", "datahub_get_entity"},
	})
	_ = reg.Register(&mockToolkit{
		kind: "s3", name: "storage", connection: "s3-storage",
		tools: []string{"s3_list_buckets", "s3_get_object"},
	})

	tests := []struct {
		tool string
		want ToolkitMatch
	}{
		{"trino_query", ToolkitMatch{Kind: regTestTrino, Name: "production", Connection: "prod-trino", Found: true}},
		{"datahub_search", ToolkitMatch{Kind: "datahub", Name: "main", Connection: "main-datahub", Found: true}},
		{"s3_list_buckets", ToolkitMatch{Kind: "s3", Name: "storage", Connection: "s3-storage", Found: true}},
		{"unknown", ToolkitMatch{}},
	}

	for _, tt := range tests {
		t.Run(tt.tool, func(t *testing.T) {
			match := reg.GetToolkitForTool(tt.tool)
			assertToolMatch(t, match, tt.want)
		})
	}
}

// mockMultiConnToolkit implements ConnectionResolver, mapping each tool
// name to its source connection. Models the gateway toolkit pattern.
type mockMultiConnToolkit struct {
	mockToolkit
	toolToConn map[string]string
}

func (m *mockMultiConnToolkit) ConnectionForTool(toolName string) string {
	return m.toolToConn[toolName]
}

func TestGetToolkitForTool_MultiConnectionResolverWins(t *testing.T) {
	reg := NewRegistry()
	gw := &mockMultiConnToolkit{
		mockToolkit: mockToolkit{
			kind: "mcp", name: "primary", connection: "default-primary",
			tools: []string{"vendorA__list", "vendorB__list"},
		},
		toolToConn: map[string]string{
			"vendorA__list": "vendorA",
			"vendorB__list": "vendorB",
		},
	}
	if err := reg.Register(gw); err != nil {
		t.Fatalf("Register: %v", err)
	}

	matchA := reg.GetToolkitForTool("vendorA__list")
	if matchA.Connection != "vendorA" {
		t.Errorf("Connection for vendorA__list = %q, want vendorA", matchA.Connection)
	}
	matchB := reg.GetToolkitForTool("vendorB__list")
	if matchB.Connection != "vendorB" {
		t.Errorf("Connection for vendorB__list = %q, want vendorB", matchB.Connection)
	}
}

func TestGetToolkitForTool_FallbackWhenResolverReturnsEmpty(t *testing.T) {
	reg := NewRegistry()
	gw := &mockMultiConnToolkit{
		mockToolkit: mockToolkit{
			kind: "mcp", name: "primary", connection: "default-fallback",
			tools: []string{"orphan__tool"},
		},
		toolToConn: map[string]string{}, // resolver returns "" → fallback
	}
	if err := reg.Register(gw); err != nil {
		t.Fatalf("Register: %v", err)
	}
	match := reg.GetToolkitForTool("orphan__tool")
	if match.Connection != "default-fallback" {
		t.Errorf("Connection = %q, want default-fallback (resolver empty → fallback)", match.Connection)
	}
}

func assertToolMatch(t *testing.T, got, want ToolkitMatch) {
	t.Helper()
	if got.Found != want.Found {
		t.Errorf("found = %v, want %v", got.Found, want.Found)
	}
	if got.Kind != want.Kind {
		t.Errorf("kind = %q, want %q", got.Kind, want.Kind)
	}
	if got.Name != want.Name {
		t.Errorf("name = %q, want %q", got.Name, want.Name)
	}
	if got.Connection != want.Connection {
		t.Errorf("connection = %q, want %q", got.Connection, want.Connection)
	}
}
