package aifin

import (
	"bytes"
	"encoding/json"
	"math"
	"regexp"
	"strings"
)

var (
	leadingZeroHKCode = regexp.MustCompile(`^0\d{4}\.HK$`)
	knownWindCode     = regexp.MustCompile(
		`^(?:\d{4}\.HK|\d{6}\.(?:SH|SZ|BJ|OF)|[A-Z]{1,5}\.(?:O|N|A|HK|SH|SZ|BJ))$`,
	)
)

var backendKlinePeriod = map[string]string{
	"1min":   "1",
	"5min":   "3",
	"10min":  "4",
	"15min":  "5",
	"30min":  "6",
	"60min":  "7",
	"120min": "8",
	"240min": "9",
	"1d":     "10",
	"1w":     "11",
	"1mo":    "12",
	"1y":     "13",
	"1q":     "14",
	"6mo":    "15",
}

func mapToolArguments(
	action string,
	raw json.RawMessage,
) (map[string]any, error) {
	switch action {
	case ActionGetStockPriceIndicators,
		ActionGetFundPriceIndicators,
		ActionGetIndexPriceIndicators:
		return mapPriceArguments(raw)
	case ActionGetStockKline,
		ActionGetFundKline,
		ActionGetIndexKline:
		return mapKlineArguments(raw)
	case ActionSearchStocks,
		ActionGetStockFundamentals,
		ActionSearchFunds,
		ActionGetFundPerformance:
		return mapQuestionArguments(raw)
	case ActionGetCompanyAnnouncements, ActionGetFinancialNews:
		return mapDocumentArguments(raw)
	case ActionNaturalLanguageGetEDBData:
		return mapEconomicArguments(raw)
	default:
		return nil, actionNotFound()
	}
}

func mapPriceArguments(raw json.RawMessage) (map[string]any, error) {
	var input struct {
		Windcode string  `json:"windcode"`
		Indexes  *string `json:"indexes"`
	}
	if err := json.Unmarshal(raw, &input); err != nil {
		return nil, invalidInput()
	}
	windcode, err := normalizeWindcode(input.Windcode)
	if err != nil {
		return nil, err
	}
	arguments := map[string]any{"windcode": windcode}
	if input.Indexes != nil {
		arguments["indexes"] = *input.Indexes
	}
	return arguments, nil
}

func mapKlineArguments(raw json.RawMessage) (map[string]any, error) {
	var input struct {
		Windcode string  `json:"windcode"`
		Begin    string  `json:"begin_date"`
		End      string  `json:"end_date"`
		Period   *string `json:"period"`
		AFType   *string `json:"aftype"`
		IsSusp   *string `json:"issusp"`
		AFDate   *string `json:"afdate"`
	}
	if err := json.Unmarshal(raw, &input); err != nil {
		return nil, invalidInput()
	}
	if input.Begin > input.End {
		return nil, invalidInput()
	}
	windcode, err := normalizeWindcode(input.Windcode)
	if err != nil {
		return nil, err
	}
	period := "1d"
	if input.Period != nil {
		period = *input.Period
	}
	mappedPeriod, ok := backendKlinePeriod[period]
	if !ok {
		return nil, invalidInput()
	}
	arguments := map[string]any{
		"windcode":   windcode,
		"begin_date": input.Begin,
		"end_date":   input.End,
		"period":     mappedPeriod,
	}
	if input.AFType != nil {
		arguments["aftype"] = *input.AFType
	}
	if input.IsSusp != nil {
		arguments["issusp"] = *input.IsSusp
	}
	if input.AFDate != nil {
		arguments["afdate"] = *input.AFDate
	}
	return arguments, nil
}

func mapQuestionArguments(raw json.RawMessage) (map[string]any, error) {
	var input struct {
		Question string `json:"question"`
	}
	if err := json.Unmarshal(raw, &input); err != nil {
		return nil, invalidInput()
	}
	return map[string]any{"question": input.Question}, nil
}

func mapDocumentArguments(raw json.RawMessage) (map[string]any, error) {
	var input struct {
		Query string       `json:"query"`
		TopK  *json.Number `json:"top_k"`
	}
	decoder := json.NewDecoder(bytes.NewReader(raw))
	decoder.UseNumber()
	if err := decoder.Decode(&input); err != nil {
		return nil, invalidInput()
	}
	arguments := map[string]any{"query": input.Query}
	if input.TopK != nil {
		value, err := input.TopK.Float64()
		if err != nil ||
			value != math.Trunc(value) ||
			value < 1 ||
			value > 100 {
			return nil, invalidInput()
		}
		arguments["top_k"] = int(value)
	}
	return arguments, nil
}

func mapEconomicArguments(raw json.RawMessage) (map[string]any, error) {
	var input struct {
		ExecutionMode string  `json:"executionMode"`
		Question      string  `json:"question"`
		BeginDate     *string `json:"beginDate"`
		EndDate       *string `json:"endDate"`
		Observation   *string `json:"observation"`
	}
	if err := json.Unmarshal(raw, &input); err != nil {
		return nil, invalidInput()
	}
	hasBegin := input.BeginDate != nil
	hasEnd := input.EndDate != nil
	if hasBegin != hasEnd ||
		(input.Observation != nil && (hasBegin || hasEnd)) {
		return nil, invalidInput()
	}
	if hasBegin && *input.BeginDate > *input.EndDate {
		return nil, invalidInput()
	}
	mode := mapEconomicMode(input.ExecutionMode)
	if (mode == "fetch" || mode == "searchFetch") &&
		input.Observation == nil &&
		!hasBegin {
		return nil, invalidInput()
	}
	arguments := map[string]any{
		"executionMode": mode,
		"question":      input.Question,
	}
	if hasBegin {
		arguments["beginDate"] = *input.BeginDate
		arguments["endDate"] = *input.EndDate
	}
	if input.Observation != nil {
		arguments["observation"] = *input.Observation
	}
	return arguments, nil
}

func mapEconomicMode(mode string) string {
	switch mode {
	case "仅搜索":
		return "search"
	case "仅提数":
		return "fetch"
	case "搜索并提数":
		return "searchFetch"
	default:
		return mode
	}
}

func normalizeWindcode(value string) (string, error) {
	rawValues := strings.Split(value, ",")
	normalized := make([]string, 0, len(rawValues))
	for _, value := range rawValues {
		raw := strings.TrimSpace(value)
		if raw == "" {
			return "", invalidInput()
		}
		upper := strings.ToUpper(raw)
		switch {
		case leadingZeroHKCode.MatchString(upper):
			normalized = append(normalized, upper[1:])
		case knownWindCode.MatchString(upper):
			normalized = append(normalized, upper)
		default:
			normalized = append(normalized, raw)
		}
	}
	if len(normalized) == 0 || len(normalized) > 50 {
		return "", invalidInput()
	}
	return strings.Join(normalized, ","), nil
}