package arxiv

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"

	"github.com/tidbcloud/AgentTeam/internal/extension/registry"
	"github.com/tidbcloud/AgentTeam/internal/service/extensionflow"
)

type Invoker interface {
	Search(context.Context, SearchRequest) (SearchResult, error)
	Get(context.Context, GetRequest) (PaperResult, error)
}

type Executor struct {
	invoker Invoker
}

func NewExecutor(invoker Invoker) *Executor {
	return &Executor{invoker: invoker}
}

func (*Executor) Strategy() string {
	return "native"
}

func (executor *Executor) Invoke(
	ctx context.Context,
	invocation registry.BuiltinInvocation,
) (registry.BuiltinResult, error) {
	if executor == nil || executor.invoker == nil {
		return registry.BuiltinResult{}, fmt.Errorf(
			"arXiv Builtin executor is not configured",
		)
	}
	var (
		result any
		err    error
	)
	switch invocation.Action {
	case ActionSearch:
		var input struct {
			Query     string   `json:"query"`
			IDList    []string `json:"idList"`
			Offset    int      `json:"offset"`
			Limit     int      `json:"limit"`
			SortBy    string   `json:"sortBy"`
			SortOrder string   `json:"sortOrder"`
		}
		if err := json.Unmarshal(invocation.Input, &input); err != nil {
			return registry.BuiltinResult{}, invalidInput()
		}
		result, err = executor.invoker.Search(
			ctx,
			SearchRequest{
				Query:     input.Query,
				IDList:    input.IDList,
				Offset:    input.Offset,
				Limit:     input.Limit,
				SortBy:    input.SortBy,
				SortOrder: input.SortOrder,
			},
		)
	case ActionGet:
		input, decodeErr := decodePaperInput(invocation.Input)
		if decodeErr != nil {
			return registry.BuiltinResult{}, decodeErr
		}
		result, err = executor.invoker.Get(
			ctx,
			GetRequest{ID: input.ID},
		)
	default:
		return registry.BuiltinResult{}, &extensionflow.InvocationError{
			Code:    extensionflow.InvocationNotFound,
			Message: "arXiv action was not found",
		}
	}
	if err != nil {
		return registry.BuiltinResult{}, classifyError(err)
	}
	encoded, err := json.Marshal(result)
	if err != nil {
		return registry.BuiltinResult{}, fmt.Errorf(
			"encode arXiv result: %w",
			err,
		)
	}
	return registry.BuiltinResult{Data: encoded}, nil
}

type paperInput struct {
	ID string `json:"id"`
}

func decodePaperInput(raw json.RawMessage) (paperInput, error) {
	var input paperInput
	if err := json.Unmarshal(raw, &input); err != nil {
		return paperInput{}, invalidInput()
	}
	return input, nil
}

func invalidInput() error {
	return &extensionflow.InvocationError{
		Code:    extensionflow.InvocationValidation,
		Message: "arXiv input was invalid",
	}
}

func classifyError(err error) error {
	if err == nil ||
		errors.Is(err, context.Canceled) ||
		errors.Is(err, context.DeadlineExceeded) {
		return err
	}
	var code extensionflow.InvocationErrorCode
	var message string
	switch CodeOf(err) {
	case ErrorValidation:
		code = extensionflow.InvocationValidation
		message = "arXiv input was invalid"
	case ErrorNotFound:
		code = extensionflow.InvocationNotFound
		message = "arXiv paper was not found"
	case ErrorRateLimited:
		code = extensionflow.InvocationRateLimited
		message = "arXiv rate limited the request"
	case ErrorUpstreamFailure:
		code = extensionflow.InvocationUpstreamFailure
		message = "arXiv request failed"
	default:
		return err
	}
	return &extensionflow.InvocationError{
		Code:    code,
		Message: message,
	}
}