package jin10

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

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

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

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/jin10/" + name + ".json"
	if err := compiler.AddResource(resource, document); err != nil {
		return nil, fmt.Errorf("compile Jin10 %s schema resource: %w", name, err)
	}
	compiled, err := compiler.Compile(resource)
	if err != nil {
		return nil, fmt.Errorf("compile Jin10 %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, runner.Errorf(
			runner.ErrorNotFound,
			"Jin10 action is not served",
			nil,
		)
	}
	return validateJSONDocument(
		compiled,
		raw,
		runner.ErrorValidation,
		"Jin10 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, runner.Errorf(
			runner.ErrorNotFound,
			"Jin10 action is not served",
			nil,
		)
	}
	return validateJSONDocument(
		compiled,
		raw,
		runner.ErrorUpstreamFailure,
		"Jin10 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)
	}
	compact, value, err := runner.DecodeUniqueJSONDocument(raw)
	if err != nil {
		return nil, runner.Errorf(code, message, nil)
	}
	if err := compiled.Validate(value); err != nil {
		return nil, runner.Errorf(code, message, nil)
	}
	return compact, nil
}

func normalizeToolResult(
	result mcp.ToolResult,
) (json.RawMessage, error) {
	structured := bytes.TrimSpace(result.StructuredContent)
	if result.StructuredContentPresent ||
		len(result.StructuredContent) > 0 {
		if len(structured) == 0 ||
			bytes.Equal(structured, []byte("null")) {
			return nil, invalidUpstreamResult()
		}
		return append(json.RawMessage(nil), structured...), nil
	}
	if len(result.Content) != 1 {
		return nil, invalidUpstreamResult()
	}
	var content struct {
		Type string  `json:"type"`
		Text *string `json:"text"`
	}
	if err := json.Unmarshal(result.Content[0], &content); err != nil ||
		content.Type != "text" ||
		content.Text == nil ||
		*content.Text == "" {
		return nil, invalidUpstreamResult()
	}
	return json.RawMessage(*content.Text), nil
}

func normalizeResourceResult(
	result mcp.ResourceResult,
) (json.RawMessage, error) {
	if len(result.Contents) != 1 {
		return nil, invalidUpstreamResult()
	}
	content := result.Contents[0]
	if content.URI != QuoteCodesResourceURI ||
		content.Text == "" ||
		len(content.Blob) != 0 {
		return nil, invalidUpstreamResult()
	}
	return json.RawMessage(content.Text), nil
}

func invalidUpstreamResult() error {
	return runner.Errorf(
		runner.ErrorUpstreamFailure,
		"Jin10 returned an invalid MCP result",
		nil,
	)
}