package tianyancha

import (
	"slices"

	"github.com/tidbcloud/AgentTeam/internal/extension/registry"
	"github.com/tidbcloud/AgentTeam/internal/extension/spec"
)

// Action name constants for the Tianyancha catalog.
const (
	ActionSearchCompanies                   = "search-companies"
	ActionSearchCompaniesAdvanced           = "search-companies-advanced"
	ActionGetCompanyBasicInfo               = "get-company-basic-info"
	ActionVerifyCompanyIdentity             = "verify-company-identity"
	ActionGetCompanyContactInfo             = "get-company-contact-info"
	ActionListCompanyKeyPersonnel           = "list-company-key-personnel"
	ActionListCompanyShareholders           = "list-company-shareholders"
	ActionListCompanyInvestments            = "list-company-investments"
	ActionListCompanyEquityChanges          = "list-company-equity-changes"
	ActionListCompanyHistoricalShareholders = "list-company-historical-shareholders"
	ActionListCompanyHistoricalInvestments  = "list-company-historical-investments"
	ActionListCompanyBranches               = "list-company-branches"
	ActionListCompanyChanges                = "list-company-changes"
	ActionGetCompanyAnnualReports           = "get-company-annual-reports"
	ActionGetCompanyRisk                    = "get-company-risk"
	ActionGetCompanyRiskDetail              = "get-company-risk-detail"
	ActionGetCompanyJudicialRisk            = "get-company-judicial-risk"
	ActionGetCompanyDevelopmentInfo         = "get-company-development-info"
	ActionGetCompanyIntellectualProperty    = "get-company-intellectual-property"
	ActionListCompanyUltimateBeneficiaries  = "list-company-ultimate-beneficiaries"
	ActionGetCompanyActualControl           = "get-company-actual-control"
	ActionFindCompanyRelationshipPaths      = "find-company-relationship-paths"
	ActionSearchTenders                     = "search-tenders"
	ActionListCompanyNews                   = "list-company-news"
)

// Definition returns the Agent-visible Tianyancha catalog, assembled from the
// per-resource action groups below.
func Definition() registry.Definition {
	return registry.Definition{
		Domain:      Domain,
		Version:     1,
		Name:        "Tianyancha",
		Description: "Query Tianyancha company registration, risk, judicial, intellectual-property, tender, and news data by company identifier.",
		Actions: slices.Concat(
			searchActions(),
			companyDetailActions(),
			registryActions(),
			riskActions(),
			relationshipActions(),
		),
	}
}

func readSafety() spec.Safety {
	return spec.Safety{Class: "read", Idempotent: true, Destructive: false}
}

func searchActions() []registry.DefinitionAction {
	return []registry.DefinitionAction{
		{
			Name:        ActionSearchCompanies,
			Summary:     "Search Tianyancha companies by keyword",
			Description: "Search Tianyancha for companies by keyword and return matching company records.",
			Safety:      readSafety(),
			InputSchema: objectSchema(map[string]any{
				"keyword":  nonWhitespaceSchema("The company search keyword."),
				"pageNum":  pageNumSchema,
				"pageSize": pageSizeSchema,
			}, []string{"keyword"}),
			OutputSchema: pagedOutputSchema("companies", "The company records matching the search keyword."),
		},
		{
			Name:        ActionSearchCompaniesAdvanced,
			Summary:     "Search Tianyancha companies by keyword, industry, and region",
			Description: "Search Tianyancha for companies by keyword, industry, and region.",
			Safety:      readSafety(),
			InputSchema: objectSchema(map[string]any{
				"keyword":      nonWhitespaceSchema("The company search keyword."),
				"industryCode": nonWhitespaceSchema("The GB/T 4754-2017 industry code used to filter companies."),
				"areaCode":     nonWhitespaceSchema("The Tianyancha administrative area code used to filter companies."),
				"pageNum":      pageNumSchema,
				"pageSize":     pageSizeSchema,
			}, []string{"keyword"}),
			OutputSchema: pagedOutputSchema("companies", "The companies matching the search filters."),
		},
	}
}

func companyDetailActions() []registry.DefinitionAction {
	return []registry.DefinitionAction{
		{
			Name:        ActionGetCompanyBasicInfo,
			Summary:     "Get Tianyancha company registration details",
			Description: "Get Tianyancha company registration details by company identifier.",
			Safety:      readSafety(),
			InputSchema: objectSchema(map[string]any{
				"keyword": keywordSchema,
			}, []string{"keyword"}),
			OutputSchema: objectSchema(map[string]any{
				"company": nullableRawObjectSchema("The company registration record."),
			}, nil),
		},
		{
			Name:        ActionVerifyCompanyIdentity,
			Summary:     "Verify a company name, legal representative, and registration code",
			Description: "Verify that a company name, legal representative, and registration code identify the same company.",
			Safety:      readSafety(),
			InputSchema: objectSchema(map[string]any{
				"companyName":      nonWhitespaceSchema("The exact company name to verify."),
				"legalPersonName":  nonWhitespaceSchema("The legal representative name to verify."),
				"registrationCode": nonWhitespaceSchema("The unified social credit code, registration number, or organization code to verify."),
			}, []string{"companyName", "legalPersonName", "registrationCode"}),
			OutputSchema: objectSchema(map[string]any{
				"matchStatus": nullableIntegerSchema(
					"The match status: 0 for no match, 1 for a full match, or 2 when the company name may be a former name and the other fields match.",
					0, 2),
				"remark": nullableStringSchema("The verification explanation returned by Tianyancha."),
			}, nil),
		},
		{
			Name:        ActionGetCompanyContactInfo,
			Summary:     "Get a company's contact details",
			Description: "Get a company's phone numbers, email addresses, websites, and registered address.",
			Safety:      readSafety(),
			InputSchema: objectSchema(map[string]any{
				"keyword": keywordSchema,
			}, []string{"keyword"}),
			OutputSchema: objectSchema(map[string]any{
				"contact": nullableRawObjectSchema("The company contact payload."),
			}, nil),
		},
		{
			Name:        ActionGetCompanyAnnualReports,
			Summary:     "Get a company's annual reports",
			Description: "Get a company's annual reports, including disclosed financial, shareholder, investment, and social security information.",
			Safety:      readSafety(),
			InputSchema: objectSchema(map[string]any{
				"keyword": keywordSchema,
				"year":    nonWhitespaceSchema("The report year to retrieve. Omit it to retrieve every available year."),
			}, []string{"keyword"}),
			OutputSchema: objectSchema(map[string]any{
				"annualReports": rawRecordArraySchema("The annual report records returned by Tianyancha."),
				"total":         nonNegativeIntegerSchema("The total number of annual reports returned by Tianyancha."),
			}, nil),
		},
		{
			Name:        ActionGetCompanyDevelopmentInfo,
			Summary:     "Get a company's development information",
			Description: "Get a company's core team, products, financing history, investment events, competitors, and investment institutions.",
			Safety:      readSafety(),
			InputSchema: objectSchema(map[string]any{
				"keyword": keywordSchema,
			}, []string{"keyword"}),
			OutputSchema: objectSchema(map[string]any{
				"development": nullableRawObjectSchema("The combined company development payload."),
			}, nil),
		},
		{
			Name:        ActionGetCompanyIntellectualProperty,
			Summary:     "Get a company's intellectual property",
			Description: "Get a company's trademarks, patents, software copyrights, works copyrights, and ICP registrations.",
			Safety:      readSafety(),
			InputSchema: objectSchema(map[string]any{
				"keyword": keywordSchema,
			}, []string{"keyword"}),
			OutputSchema: objectSchema(map[string]any{
				"intellectualProperty": nullableRawObjectSchema("The combined company intellectual-property payload."),
			}, nil),
		},
	}
}

func registryActions() []registry.DefinitionAction {
	return []registry.DefinitionAction{
		{
			Name:         ActionListCompanyKeyPersonnel,
			Summary:      "List a company's directors, supervisors, and senior managers",
			Description:  "List a company's directors, supervisors, and senior managers from Tianyancha.",
			Safety:       readSafety(),
			InputSchema:  keywordPageInputSchema(nil),
			OutputSchema: pagedOutputSchema("personnel", "The directors, supervisors, and senior managers returned by Tianyancha."),
		},
		{
			Name:        ActionListCompanyShareholders,
			Summary:     "List a company's shareholders",
			Description: "List a company's shareholders and contribution details from Tianyancha.",
			Safety:      readSafety(),
			InputSchema: keywordPageInputSchema(map[string]any{
				"source": shareholderSourceSchema,
			}),
			OutputSchema: pagedOutputSchema("shareholders", "The shareholder and contribution records returned by Tianyancha."),
		},
		{
			Name:         ActionListCompanyInvestments,
			Summary:      "List a company's outbound investments",
			Description:  "List companies directly invested in by a company using Tianyancha.",
			Safety:       readSafety(),
			InputSchema:  keywordPageInputSchema(nil),
			OutputSchema: pagedOutputSchema("investments", "The directly invested company records returned by Tianyancha."),
		},
		{
			Name:         ActionListCompanyEquityChanges,
			Summary:      "List a company's equity changes",
			Description:  "List a company's disclosed equity changes, including shareholders before and after each change.",
			Safety:       readSafety(),
			InputSchema:  keywordPageInputSchema(nil),
			OutputSchema: pagedOutputSchema("equityChanges", "The disclosed equity change records returned by Tianyancha."),
		},
		{
			Name:         ActionListCompanyHistoricalShareholders,
			Summary:      "List a company's historical shareholders",
			Description:  "List a company's historical shareholders, ownership ratios, and subscribed capital from Tianyancha.",
			Safety:       readSafety(),
			InputSchema:  keywordPageInputSchema(nil),
			OutputSchema: pagedOutputSchema("historicalShareholders", "The historical shareholder records returned by Tianyancha."),
		},
		{
			Name:         ActionListCompanyHistoricalInvestments,
			Summary:      "List a company's historical outbound investments",
			Description:  "List companies previously invested in by a company using Tianyancha.",
			Safety:       readSafety(),
			InputSchema:  keywordPageInputSchema(nil),
			OutputSchema: pagedOutputSchema("historicalInvestments", "The historical outbound investment records returned by Tianyancha."),
		},
		{
			Name:         ActionListCompanyBranches,
			Summary:      "List a company's branches",
			Description:  "List a company's branch organizations and their registration details.",
			Safety:       readSafety(),
			InputSchema:  keywordPageInputSchema(nil),
			OutputSchema: pagedOutputSchema("branches", "The company branch records returned by Tianyancha."),
		},
		{
			Name:         ActionListCompanyChanges,
			Summary:      "List a company's registration changes",
			Description:  "List a company's registration changes, including before and after values.",
			Safety:       readSafety(),
			InputSchema:  keywordPageInputSchema(nil),
			OutputSchema: pagedOutputSchema("changes", "The company registration change records returned by Tianyancha."),
		},
		{
			Name:         ActionListCompanyUltimateBeneficiaries,
			Summary:      "List a company's ultimate beneficiaries",
			Description:  "List a company's ultimate beneficiaries, ownership percentages, and ownership chains from Tianyancha.",
			Safety:       readSafety(),
			InputSchema:  keywordPageInputSchema(nil),
			OutputSchema: pagedOutputSchema("beneficiaries", "The ultimate beneficiary records and ownership chains returned by Tianyancha."),
		},
		{
			Name:        ActionGetCompanyActualControl,
			Summary:     "Get a company's suspected actual controllers",
			Description: "Get a company's suspected actual controllers, ownership ratios, and control paths from Tianyancha.",
			Safety:      readSafety(),
			InputSchema: objectSchema(map[string]any{
				"keyword": keywordSchema,
			}, []string{"keyword"}),
			OutputSchema: objectSchema(map[string]any{
				"controllers": rawRecordArraySchema("The suspected actual controllers returned by Tianyancha."),
				"paths": map[string]any{
					"$schema":              schemaDialect,
					"type":                 "object",
					"description":          "The provider-defined ownership paths keyed by Tianyancha path ID.",
					"additionalProperties": true,
				},
			}, nil),
		},
	}
}

func riskActions() []registry.DefinitionAction {
	return []registry.DefinitionAction{
		{
			Name:        ActionGetCompanyRisk,
			Summary:     "Get a company's risk summary",
			Description: "Get a company's Tianyancha risk summary, including own, related, historical, and alert risks.",
			Safety:      readSafety(),
			InputSchema: objectSchema(map[string]any{
				"keyword": keywordSchema,
			}, []string{"keyword"}),
			OutputSchema: objectSchema(map[string]any{
				"riskLevel": nullableStringSchema("The overall risk level returned by Tianyancha."),
				"risks":     rawRecordArraySchema("Risk groups returned by Tianyancha, each with nested risk categories and entries."),
			}, nil),
		},
		{
			Name:        ActionGetCompanyRiskDetail,
			Summary:     "Get one page of detailed company risk records",
			Description: "Get one page of detailed Tianyancha risk records using a risk ID and type from get-company-risk.",
			Safety:      readSafety(),
			InputSchema: objectSchema(map[string]any{
				"riskId":   positiveIntegerSchema("The risk ID returned by get-company-risk."),
				"riskType": positiveIntegerSchema("The Tianyancha risk type code returned by get-company-risk."),
				"pageNum":  pageNumSchema,
				"pageSize": pageSizeSchema,
			}, []string{"riskId", "riskType"}),
			OutputSchema: objectSchema(map[string]any{
				"riskType": positiveIntegerSchema("The Tianyancha risk type code."),
				"records":  rawRecordArraySchema("The detailed risk records returned by Tianyancha."),
				"total":    nonNegativeIntegerSchema("The total number of detailed risk records."),
				"pageNum":  positiveIntegerSchema("The returned one-based page number."),
				"pageSize": positiveIntegerSchema("The requested number of records per page."),
			}, nil),
		},
		{
			Name:        ActionGetCompanyJudicialRisk,
			Summary:     "Get a company's judicial risk data",
			Description: "Get a company's combined legal cases, hearings, court notices, enforcement, dishonesty, filing, and service-announcement data.",
			Safety:      readSafety(),
			InputSchema: objectSchema(map[string]any{
				"keyword": keywordSchema,
			}, []string{"keyword"}),
			OutputSchema: objectSchema(map[string]any{
				"judicialRisk": nullableRawObjectSchema("The combined judicial risk payload."),
			}, nil),
		},
	}
}

func relationshipActions() []registry.DefinitionAction {
	newsInput := objectSchema(map[string]any{
		"companyName": nonWhitespaceSchema("The exact company name."),
		"companyId":   positiveIntegerSchema("The Tianyancha company ID."),
		"startDate":   dateSchema("The earliest publication date to include, in YYYY-MM-DD format."),
		"endDate":     dateSchema("The latest publication date to include, in YYYY-MM-DD format."),
		"tags":        stringArraySchema("The Tianyancha news tags to include."),
		"pageNum":     pageNumSchema,
		"pageSize":    pageSizeSchema,
	}, nil)
	newsInput["anyOf"] = []any{
		map[string]any{"required": []any{"companyName"}},
		map[string]any{"required": []any{"companyId"}},
	}

	relationshipInput := objectSchema(map[string]any{
		"sourceCompanyName": nonWhitespaceSchema("The exact source company name."),
		"sourceCompanyId":   positiveIntegerSchema("The Tianyancha source company ID."),
		"targetCompanyName": nonWhitespaceSchema("The exact target company name."),
		"targetCompanyId":   positiveIntegerSchema("The Tianyancha target company ID."),
		"relationshipTypes": enumArraySchema(
			"The relationship categories to search. Omit this field to use Tianyancha's default categories.",
			"A relationship category.", relationshipTypeValues),
		"maxDepth": maxDepthSchema,
	}, nil)
	relationshipInput["allOf"] = []any{
		map[string]any{"anyOf": []any{
			map[string]any{"required": []any{"sourceCompanyName"}},
			map[string]any{"required": []any{"sourceCompanyId"}},
		}},
		map[string]any{"anyOf": []any{
			map[string]any{"required": []any{"targetCompanyName"}},
			map[string]any{"required": []any{"targetCompanyId"}},
		}},
	}

	return []registry.DefinitionAction{
		{
			Name:        ActionFindCompanyRelationshipPaths,
			Summary:     "Find relationship paths between two companies",
			Description: "Find the shortest ownership, employment, legal, business, or historical relationship paths between two companies.",
			Safety:      readSafety(),
			InputSchema: relationshipInput,
			OutputSchema: objectSchema(map[string]any{
				"hasRelationship": map[string]any{
					"type":        "boolean",
					"description": "Whether Tianyancha found at least one relationship path.",
				},
				"paths": map[string]any{
					"type":        "array",
					"items":       rawObjectSchema("A shortest relationship path returned by Tianyancha."),
					"description": "The shortest relationship paths returned by Tianyancha.",
				},
			}, nil),
		},
		{
			Name:        ActionSearchTenders,
			Summary:     "Search tender and bid notices",
			Description: "Search Tianyancha tender and bid notices by keyword, participant role, notice type, province, and publication date.",
			Safety:      readSafety(),
			InputSchema: objectSchema(map[string]any{
				"keyword":          nonWhitespaceSchema("A keyword matched against titles, purchasers, or suppliers."),
				"searchTypes":      enumArraySchema("The fields to search. Omit this field to use Tianyancha's default.", "A tender search field.", []string{"title", "purchaser", "supplier"}),
				"noticeType":       enumSchema("The tender notice type to include.", []string{"forecast", "announcement", "result"}),
				"province":         nonWhitespaceSchema("The province used to filter results."),
				"publishStartDate": dateSchema("The earliest publication date to include, in YYYY-MM-DD format."),
				"publishEndDate":   dateSchema("The latest publication date to include, in YYYY-MM-DD format."),
				"pageNum":          pageNumSchema,
				"pageSize":         pageSizeSchema,
			}, []string{"keyword"}),
			OutputSchema: pagedOutputSchema("tenders", "The tender and bid notices matching the search filters."),
		},
		{
			Name:         ActionListCompanyNews,
			Summary:      "List news for one company",
			Description:  "List Tianyancha news for an exact company, optionally filtered by publication date and tags.",
			Safety:       readSafety(),
			InputSchema:  newsInput,
			OutputSchema: pagedOutputSchema("news", "The company news records returned by Tianyancha."),
		},
	}
}

var shareholderSourceSchema = map[string]any{
	"type":        "integer",
	"minimum":     1,
	"maximum":     2,
	"description": "The shareholder data source: 1 for business registration or 2 for the latest public filing. Defaults to 1.",
}

var maxDepthSchema = map[string]any{
	"type":        "integer",
	"minimum":     1,
	"maximum":     10,
	"description": "The maximum relationship search depth. Defaults to 2.",
}

func keywordPageInputSchema(extra map[string]any) map[string]any {
	properties := map[string]any{
		"keyword":  keywordSchema,
		"pageNum":  pageNumSchema,
		"pageSize": pageSizeSchema,
	}
	for name, schema := range extra {
		properties[name] = schema
	}
	return objectSchema(properties, []string{"keyword"})
}

var relationshipTypeValues = []string{
	"legal_representative",
	"employment",
	"investment",
	"branch",
	"legal",
	"pledge_and_mortgage",
	"business_cooperation",
	"shared_contact",
	"historical_shareholder",
	"historical_employment",
	"historical_legal_representative",
	"all",
}
