package jin10

import (
	"context"
	"encoding/json"
	"sort"

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

type inputKind uint8

const (
	inputEmpty inputKind = iota
	inputCode
	inputKline
	inputCursor
	inputKeyword
	inputKeywordCursor
	inputNewsID
)

type toolBinding struct {
	tool  string
	input inputKind
}

var toolBindings = map[string]toolBinding{
	ActionGetQuote:     {tool: "get_quote", input: inputCode},
	ActionGetKline:     {tool: "get_kline", input: inputKline},
	ActionListFlash:    {tool: "list_flash", input: inputCursor},
	ActionSearchFlash:  {tool: "search_flash", input: inputKeyword},
	ActionListNews:     {tool: "list_news", input: inputCursor},
	ActionSearchNews:   {tool: "search_news", input: inputKeywordCursor},
	ActionGetNews:      {tool: "get_news", input: inputNewsID},
	ActionListCalendar: {tool: "list_calendar", input: inputEmpty},
}

func allowedToolNames() []string {
	names := make([]string, 0, len(toolBindings))
	for _, binding := range toolBindings {
		names = append(names, binding.tool)
	}
	sort.Strings(names)
	return names
}

func invokeAction(
	ctx context.Context,
	client *mcp.Client,
	credentials runner.Credentials,
	action string,
	input json.RawMessage,
) (json.RawMessage, error) {
	if action == ActionListQuoteCodes {
		result, err := client.ReadResource(
			ctx,
			credentials,
			QuoteCodesResourceURI,
		)
		if err != nil {
			return nil, err
		}
		return normalizeResourceResult(result)
	}
	binding, ok := toolBindings[action]
	if !ok {
		return nil, runner.Errorf(
			runner.ErrorNotFound,
			"Jin10 action is not served",
			nil,
		)
	}
	arguments, err := decodeToolArguments(binding.input, input)
	if err != nil {
		return nil, err
	}
	result, err := client.CallTool(
		ctx,
		credentials,
		binding.tool,
		arguments,
	)
	if err != nil {
		return nil, err
	}
	return normalizeToolResult(result)
}

func decodeToolArguments(
	kind inputKind,
	raw json.RawMessage,
) (map[string]any, error) {
	switch kind {
	case inputEmpty:
		return map[string]any{}, nil
	case inputCode:
		var input struct {
			Code string `json:"code"`
		}
		if err := json.Unmarshal(raw, &input); err != nil {
			return nil, invalidInput()
		}
		return map[string]any{"code": input.Code}, nil
	case inputKline:
		return decodeKlineArguments(raw)
	case inputCursor:
		var input struct {
			Cursor *string `json:"cursor"`
		}
		if err := json.Unmarshal(raw, &input); err != nil {
			return nil, invalidInput()
		}
		arguments := map[string]any{}
		if input.Cursor != nil {
			arguments["cursor"] = *input.Cursor
		}
		return arguments, nil
	case inputKeyword:
		var input struct {
			Keyword string `json:"keyword"`
		}
		if err := json.Unmarshal(raw, &input); err != nil {
			return nil, invalidInput()
		}
		return map[string]any{"keyword": input.Keyword}, nil
	case inputKeywordCursor:
		return decodeKeywordCursorArguments(raw)
	case inputNewsID:
		var input struct {
			ID string `json:"id"`
		}
		if err := json.Unmarshal(raw, &input); err != nil {
			return nil, invalidInput()
		}
		return map[string]any{"id": input.ID}, nil
	default:
		return nil, invalidInput()
	}
}

func decodeKlineArguments(raw json.RawMessage) (map[string]any, error) {
	var input struct {
		Code  string       `json:"code"`
		Time  *json.Number `json:"time"`
		Count *json.Number `json:"count"`
	}
	if err := json.Unmarshal(raw, &input); err != nil {
		return nil, invalidInput()
	}
	arguments := map[string]any{"code": input.Code}
	if input.Time != nil {
		arguments["time"] = json.Number(input.Time.String())
	}
	if input.Count != nil {
		arguments["count"] = json.Number(input.Count.String())
	}
	return arguments, nil
}

func decodeKeywordCursorArguments(
	raw json.RawMessage,
) (map[string]any, error) {
	var input struct {
		Keyword string  `json:"keyword"`
		Cursor  *string `json:"cursor"`
	}
	if err := json.Unmarshal(raw, &input); err != nil {
		return nil, invalidInput()
	}
	arguments := map[string]any{"keyword": input.Keyword}
	if input.Cursor != nil {
		arguments["cursor"] = *input.Cursor
	}
	return arguments, nil
}

func invalidInput() error {
	return runner.Errorf(
		runner.ErrorValidation,
		"Jin10 action input was invalid",
		nil,
	)
}