package pkulaw

import (
	"fmt"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/extension/runner"
)

const (
	Gateway = "https://apim-gateway.pkulaw.com"

	TopologyAggregate = "aggregate"
	TopologyFallback  = "fallback"
)

type EndpointID string

const (
	EndpointAggregate    EndpointID = "aggregate"
	EndpointLawSemantic  EndpointID = "law_semantic"
	EndpointLawKeyword   EndpointID = "law_keyword"
	EndpointCaseSemantic EndpointID = "case_semantic"
	EndpointCaseKeyword  EndpointID = "case_keyword"
)

type EndpointPolicy struct {
	ID           EndpointID
	Endpoint     string
	AllowedTools []string
}

type ActionRoute struct {
	Endpoint EndpointID
	Tool     string
}

var endpointPolicies = []EndpointPolicy{
	{
		ID:       EndpointAggregate,
		Endpoint: Gateway + "/mcp-law-agg/1.0.0/mcp",
		AllowedTools: []string{
			"search_article",
			"get_article",
			"get_law_list",
			"search_case",
			"get_case_list",
		},
	},
	{
		ID:       EndpointLawSemantic,
		Endpoint: Gateway + "/mcp-law-search-service/mcp",
		AllowedTools: []string{
			"search_article",
			"get_article",
		},
	},
	{
		ID:       EndpointLawKeyword,
		Endpoint: Gateway + "/mcp-law/mcp",
		AllowedTools: []string{
			"get_law_list",
		},
	},
	{
		ID:       EndpointCaseSemantic,
		Endpoint: Gateway + "/mcp-case-search-service/mcp",
		AllowedTools: []string{
			"search_case",
		},
	},
	{
		ID:       EndpointCaseKeyword,
		Endpoint: Gateway + "/mcp-case/mcp",
		AllowedTools: []string{
			"get_case_list",
		},
	},
}

var aggregateRoutes = map[string]ActionRoute{
	ActionSearchLawArticles: {
		Endpoint: EndpointAggregate,
		Tool:     "search_article",
	},
	ActionGetLawArticle: {
		Endpoint: EndpointAggregate,
		Tool:     "get_article",
	},
	ActionSearchLaws: {
		Endpoint: EndpointAggregate,
		Tool:     "get_law_list",
	},
	ActionSearchCases: {
		Endpoint: EndpointAggregate,
		Tool:     "search_case",
	},
	ActionSearchCaseDocuments: {
		Endpoint: EndpointAggregate,
		Tool:     "get_case_list",
	},
}

var fallbackRoutes = map[string]ActionRoute{
	ActionSearchLawArticles: {
		Endpoint: EndpointLawSemantic,
		Tool:     "search_article",
	},
	ActionGetLawArticle: {
		Endpoint: EndpointLawSemantic,
		Tool:     "get_article",
	},
	ActionSearchLaws: {
		Endpoint: EndpointLawKeyword,
		Tool:     "get_law_list",
	},
	ActionSearchCases: {
		Endpoint: EndpointCaseSemantic,
		Tool:     "search_case",
	},
	ActionSearchCaseDocuments: {
		Endpoint: EndpointCaseKeyword,
		Tool:     "get_case_list",
	},
}

func EndpointPolicies() []EndpointPolicy {
	result := make([]EndpointPolicy, 0, len(endpointPolicies))
	for _, policy := range endpointPolicies {
		cloned := policy
		cloned.AllowedTools = append(
			[]string(nil),
			policy.AllowedTools...,
		)
		result = append(result, cloned)
	}
	return result
}

func cloneRoutes(routes map[string]ActionRoute) map[string]ActionRoute {
	cloned := make(map[string]ActionRoute, len(routes))
	for action, route := range routes {
		cloned[action] = route
	}
	return cloned
}

func ConnectionConfigContract() runner.ConnectionConfigContract {
	return runner.ConnectionConfigContract{
		Fields: []runner.ConnectionConfigField{{
			Name:     "topology",
			Required: true,
			Values: []string{
				TopologyAggregate,
				TopologyFallback,
			},
		}},
	}
}

func cloneConnectionConfigContract(
	contract runner.ConnectionConfigContract,
) runner.ConnectionConfigContract {
	cloned := runner.ConnectionConfigContract{
		Fields: make(
			[]runner.ConnectionConfigField,
			0,
			len(contract.Fields),
		),
	}
	for _, field := range contract.Fields {
		copied := field
		copied.Values = append([]string(nil), field.Values...)
		cloned.Fields = append(cloned.Fields, copied)
	}
	return cloned
}

func aggregateConnectionConfig() runner.ConnectionConfig {
	return runner.ConnectionConfig{"topology": TopologyAggregate}
}

func fallbackConnectionConfig() runner.ConnectionConfig {
	return runner.ConnectionConfig{"topology": TopologyFallback}
}

func validateRouting() error {
	if len(endpointPolicies) != 5 ||
		len(aggregateRoutes) != len(Definition().Actions) ||
		len(fallbackRoutes) != len(Definition().Actions) {
		return fmt.Errorf("PKULaw endpoint routing is incomplete")
	}
	toolsByEndpoint := make(
		map[EndpointID]map[string]struct{},
		len(endpointPolicies),
	)
	for _, policy := range endpointPolicies {
		if policy.ID == "" ||
			!strings.HasPrefix(policy.Endpoint, Gateway+"/") ||
			len(policy.AllowedTools) == 0 {
			return fmt.Errorf("PKULaw endpoint policy is invalid")
		}
		if _, duplicate := toolsByEndpoint[policy.ID]; duplicate {
			return fmt.Errorf("PKULaw endpoint policy is duplicated")
		}
		tools := make(map[string]struct{}, len(policy.AllowedTools))
		for _, tool := range policy.AllowedTools {
			if tool == "" {
				return fmt.Errorf("PKULaw tool policy is invalid")
			}
			if _, duplicate := tools[tool]; duplicate {
				return fmt.Errorf("PKULaw tool policy is duplicated")
			}
			tools[tool] = struct{}{}
		}
		toolsByEndpoint[policy.ID] = tools
	}
	for _, action := range Definition().Actions {
		for _, routes := range []map[string]ActionRoute{
			aggregateRoutes,
			fallbackRoutes,
		} {
			route, ok := routes[action.Name]
			if !ok {
				return fmt.Errorf("PKULaw action route is missing")
			}
			tools, endpointExists := toolsByEndpoint[route.Endpoint]
			if !endpointExists {
				return fmt.Errorf("PKULaw action endpoint is unavailable")
			}
			if _, allowed := tools[route.Tool]; !allowed {
				return fmt.Errorf("PKULaw action tool is not allowlisted")
			}
		}
	}
	return ConnectionConfigContract().ValidateSpec()
}