package aifin

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"

	jsonschema "github.com/santhosh-tekuri/jsonschema/v6"

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

type catalogValidators struct {
	input  map[string]*jsonschema.Schema
	output map[string]*jsonschema.Schema
}

func compileCatalogValidators(
	definition registry.Definition,
) (catalogValidators, error) {
	validators := catalogValidators{
		input:  make(map[string]*jsonschema.Schema, len(definition.Actions)),
		output: make(map[string]*jsonschema.Schema, len(definition.Actions)),
	}
	for _, action := range definition.Actions {
		input, err := compileSchema(action.Name+"-input", action.InputSchema)
		if err != nil {
			return catalogValidators{}, err
		}
		output, err := compileSchema(
			action.Name+"-output",
			action.OutputSchema,
		)
		if err != nil {
			return catalogValidators{}, err
		}
		validators.input[action.Name] = input
		validators.output[action.Name] = output
	}
	return validators, nil
}

func compileSchema(
	name string,
	document map[string]any,
) (*jsonschema.Schema, error) {
	compiler := jsonschema.NewCompiler()
	compiler.DefaultDraft(jsonschema.Draft2020)
	resource := "https://agentteam.invalid/aifinmarket/" + name + ".json"
	if err := compiler.AddResource(resource, document); err != nil {
		return nil, fmt.Errorf(
			"compile Wind AIFin Market %s schema resource: %w",
			name,
			err,
		)
	}
	compiled, err := compiler.Compile(resource)
	if err != nil {
		return nil, fmt.Errorf(
			"compile Wind AIFin Market %s schema: %w",
			name,
			err,
		)
	}
	return compiled, nil
}

func (validators catalogValidators) validateInput(
	action string,
	raw json.RawMessage,
) (json.RawMessage, error) {
	compiled, ok := validators.input[action]
	if !ok {
		return nil, actionNotFound()
	}
	return validateJSONDocument(
		compiled,
		raw,
		runner.ErrorValidation,
		"Wind AIFin Market action input was invalid",
	)
}

func (validators catalogValidators) validateOutput(
	action string,
	raw json.RawMessage,
) (json.RawMessage, error) {
	compiled, ok := validators.output[action]
	if !ok {
		return nil, actionNotFound()
	}
	return validateJSONDocument(
		compiled,
		raw,
		runner.ErrorUpstreamFailure,
		"Wind AIFin Market returned an invalid response",
	)
}

func validateJSONDocument(
	compiled *jsonschema.Schema,
	raw json.RawMessage,
	code runner.ErrorCode,
	message string,
) (json.RawMessage, error) {
	if compiled == nil ||
		len(raw) == 0 ||
		len(raw) > runner.MaxResponseBytes {
		return nil, runner.Errorf(code, message, nil)
	}
	decoder := json.NewDecoder(bytes.NewReader(raw))
	decoder.UseNumber()
	var value any
	if err := decoder.Decode(&value); err != nil {
		return nil, runner.Errorf(code, message, nil)
	}
	var trailing any
	if err := decoder.Decode(&trailing); err != io.EOF {
		return nil, runner.Errorf(code, message, nil)
	}
	if err := compiled.Validate(value); err != nil {
		return nil, runner.Errorf(code, message, nil)
	}
	var compact bytes.Buffer
	if err := json.Compact(&compact, raw); err != nil ||
		compact.Len() > runner.MaxResponseBytes {
		return nil, runner.Errorf(code, message, nil)
	}
	return json.RawMessage(compact.Bytes()), nil
}

func actionNotFound() error {
	return runner.Errorf(
		runner.ErrorNotFound,
		"Wind AIFin Market action is not served",
		nil,
	)
}

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