package aifin

import (
	"testing"

	"github.com/stretchr/testify/require"
)

func TestEndpointPoliciesPinFiveUsedWindFamilies(t *testing.T) {
	t.Parallel()

	require.Equal(t, []EndpointPolicy{
		{
			Family:   FamilyStockData,
			Endpoint: "https://mcp.wind.com.cn/vserver_stock_data/mcp/",
			AllowedTools: []string{
				"get_stock_price_indicators",
				"get_stock_kline",
				"search_stocks",
				"get_stock_fundamentals",
			},
		},
		{
			Family:   FamilyFundData,
			Endpoint: "https://mcp.wind.com.cn/vserver_fund_data/mcp/",
			AllowedTools: []string{
				"get_fund_price_indicators",
				"get_fund_kline",
				"search_funds",
				"get_fund_performance",
			},
		},
		{
			Family:   FamilyIndexData,
			Endpoint: "https://mcp.wind.com.cn/vserver_index_data/mcp/",
			AllowedTools: []string{
				"get_index_price_indicators",
				"get_index_kline",
			},
		},
		{
			Family:   FamilyFinancialDocs,
			Endpoint: "https://mcp.wind.com.cn/vserver_financial_docs/mcp/",
			AllowedTools: []string{
				"get_company_announcements",
				"get_financial_news",
			},
		},
		{
			Family:   FamilyEconomicData,
			Endpoint: "https://mcp.wind.com.cn/vserver_economic_data/mcp/",
			AllowedTools: []string{
				"natural_language_get_edb_data",
			},
		},
	}, EndpointPolicies())

	policies := EndpointPolicies()
	policies[0].Endpoint = "https://attacker.invalid/"
	policies[0].AllowedTools[0] = "call_tool"
	require.Equal(
		t,
		"https://mcp.wind.com.cn/vserver_stock_data/mcp/",
		EndpointPolicies()[0].Endpoint,
	)
	require.Equal(
		t,
		"get_stock_price_indicators",
		EndpointPolicies()[0].AllowedTools[0],
	)
}

func TestActionRoutesPinExactToolAndFamilyMapping(t *testing.T) {
	t.Parallel()

	require.Equal(t, map[string]ActionRoute{
		ActionGetStockPriceIndicators: {
			Family: FamilyStockData,
			Tool:   "get_stock_price_indicators",
		},
		ActionGetStockKline: {
			Family: FamilyStockData,
			Tool:   "get_stock_kline",
		},
		ActionSearchStocks: {
			Family: FamilyStockData,
			Tool:   "search_stocks",
		},
		ActionGetStockFundamentals: {
			Family: FamilyStockData,
			Tool:   "get_stock_fundamentals",
		},
		ActionGetFundPriceIndicators: {
			Family: FamilyFundData,
			Tool:   "get_fund_price_indicators",
		},
		ActionGetFundKline: {
			Family: FamilyFundData,
			Tool:   "get_fund_kline",
		},
		ActionSearchFunds: {
			Family: FamilyFundData,
			Tool:   "search_funds",
		},
		ActionGetFundPerformance: {
			Family: FamilyFundData,
			Tool:   "get_fund_performance",
		},
		ActionGetIndexPriceIndicators: {
			Family: FamilyIndexData,
			Tool:   "get_index_price_indicators",
		},
		ActionGetIndexKline: {
			Family: FamilyIndexData,
			Tool:   "get_index_kline",
		},
		ActionGetCompanyAnnouncements: {
			Family: FamilyFinancialDocs,
			Tool:   "get_company_announcements",
		},
		ActionGetFinancialNews: {
			Family: FamilyFinancialDocs,
			Tool:   "get_financial_news",
		},
		ActionNaturalLanguageGetEDBData: {
			Family: FamilyEconomicData,
			Tool:   "natural_language_get_edb_data",
		},
	}, ActionRoutes())

	routes := ActionRoutes()
	routes[ActionSearchStocks] = ActionRoute{
		Family: FamilyEconomicData,
		Tool:   "call_tool",
	}
	require.Equal(t, "search_stocks", ActionRoutes()[ActionSearchStocks].Tool)
}