package openai

import (
	"encoding/json"
	"fmt"
	"math"
	"sort"
	"strings"
)

const (
	maxToolValidationAttempts = 2

	runtimeMetadataToolValidationAttempt = "openai_tool_validation_attempt"
	runtimeMetadataToolValidationRetryOf = "openai_tool_validation_retry_of"
	runtimeMetadataToolValidationFailure = "openai_tool_validation_failure"

	toolCallOriginNative = "native"
	toolCallOriginText   = "text"
)

type toolValidationContract struct {
	enabled bool
	specs   map[string]textToolSpec
}

type toolCallCandidate struct {
	name      string
	arguments any
	origin    string
	raw       any
}

func toolValidationContractFromRequest(req chatCompletionRequest) toolValidationContract {
	if len(req.Tools) == 0 || toolChoiceDisablesToolValidation(req.ToolChoice) {
		return toolValidationContract{}
	}
	specs := requestedToolSpecs(req.Tools)
	if len(specs) == 0 {
		return toolValidationContract{}
	}
	return toolValidationContract{enabled: true, specs: specs}
}

func toolChoiceDisablesToolValidation(toolChoice any) bool {
	choice, ok := toolChoice.(string)
	return ok && strings.EqualFold(strings.TrimSpace(choice), "none")
}

func toolValidationAttemptMetadata(metadata map[string]string, attempt int, retryOf, failureReason string) map[string]string {
	out := make(map[string]string, len(metadata)+3)
	for key, value := range metadata {
		out[key] = value
	}
	out[runtimeMetadataToolValidationAttempt] = fmt.Sprintf("%d", attempt)
	if strings.TrimSpace(retryOf) != "" {
		out[runtimeMetadataToolValidationRetryOf] = retryOf
	}
	if strings.TrimSpace(failureReason) != "" {
		out[runtimeMetadataToolValidationFailure] = failureReason
	}
	return out
}

func validateToolCallResponse(contract toolValidationContract, toolCalls []any, origin string) error {
	if !contract.enabled || len(toolCalls) == 0 {
		return nil
	}
	candidates, err := toolCallCandidatesFromResponse(toolCalls, origin)
	if err != nil {
		return err
	}
	for i, candidate := range candidates {
		if err := validateToolCallCandidate(contract, candidate); err != nil {
			return fmt.Errorf("tool call %d: %w", i, err)
		}
	}
	return nil
}

func toolCallCandidatesFromResponse(toolCalls []any, origin string) ([]toolCallCandidate, error) {
	candidates := make([]toolCallCandidate, 0, len(toolCalls))
	for i, raw := range toolCalls {
		call, ok := raw.(map[string]any)
		if !ok {
			return nil, fmt.Errorf("tool call %d is not an object", i)
		}
		fn, ok := call["function"].(map[string]any)
		if !ok {
			return nil, fmt.Errorf("tool call %d function is not an object", i)
		}
		name, _ := fn["name"].(string)
		name = strings.TrimSpace(name)
		if name == "" {
			return nil, fmt.Errorf("tool call %d function.name is required", i)
		}
		arguments, err := parseToolCallArguments(fn["arguments"])
		if err != nil {
			return nil, fmt.Errorf("tool call %d function.arguments: %w", i, err)
		}
		candidates = append(candidates, toolCallCandidate{
			name:      name,
			arguments: arguments,
			origin:    origin,
			raw:       raw,
		})
	}
	return candidates, nil
}

func parseToolCallArguments(raw any) (any, error) {
	switch v := raw.(type) {
	case string:
		trimmed := strings.TrimSpace(v)
		if trimmed == "" {
			return nil, fmt.Errorf("must be valid JSON")
		}
		var decoded any
		if err := json.Unmarshal([]byte(trimmed), &decoded); err != nil {
			return nil, fmt.Errorf("must be valid JSON: %w", err)
		}
		return decoded, nil
	case map[string]any:
		return v, nil
	default:
		return nil, fmt.Errorf("must be a JSON string or object")
	}
}

func validateToolCallCandidate(contract toolValidationContract, candidate toolCallCandidate) error {
	spec, ok := contract.specs[candidate.name]
	if !ok {
		return fmt.Errorf("tool %q is not in request tools", candidate.name)
	}
	if err := validateJSONSchemaSubset(candidate.arguments, spec.parameters, "$.arguments"); err != nil {
		return fmt.Errorf("tool %q %s", candidate.name, err.Error())
	}
	return nil
}

func validateJSONSchemaSubset(value any, schema any, path string) error {
	m, ok := schema.(map[string]any)
	if !ok || len(m) == 0 {
		return nil
	}
	if branches := schemaBranches(m["anyOf"]); len(branches) > 0 {
		return validateAnySchemaBranch(value, branches, path, "anyOf")
	}
	if branches := schemaBranches(m["oneOf"]); len(branches) > 0 {
		return validateAnySchemaBranch(value, branches, path, "oneOf")
	}
	for _, branch := range schemaBranches(m["allOf"]) {
		if err := validateJSONSchemaSubset(value, branch, path); err != nil {
			return err
		}
	}
	if values := schemaEnumValues(m["enum"]); len(values) > 0 && !valueInJSONEnum(value, values) {
		return fmt.Errorf("%s value %s is not in enum", path, jsonValuePreview(value))
	}
	types := schemaTypeNames(m["type"])
	if len(types) > 0 && !jsonValueMatchesAnyType(value, types) {
		return fmt.Errorf("%s type %s does not match schema type %s", path, jsonValueTypeName(value), strings.Join(types, "|"))
	}
	if shouldValidateObjectSubset(value, m, types) {
		obj, ok := value.(map[string]any)
		if !ok {
			return fmt.Errorf("%s type %s does not match schema type object", path, jsonValueTypeName(value))
		}
		if err := validateJSONObjectSubset(obj, m, path); err != nil {
			return err
		}
	}
	if shouldValidateArraySubset(value, m, types) {
		items, ok := value.([]any)
		if !ok {
			return fmt.Errorf("%s type %s does not match schema type array", path, jsonValueTypeName(value))
		}
		if itemSchema := schemaArrayItemSchema(m); itemSchema != nil {
			for i, item := range items {
				if err := validateJSONSchemaSubset(item, itemSchema, fmt.Sprintf("%s[%d]", path, i)); err != nil {
					return err
				}
			}
		}
	}
	return nil
}

func validateAnySchemaBranch(value any, branches []any, path, branchType string) error {
	var reasons []string
	for _, branch := range branches {
		if err := validateJSONSchemaSubset(value, branch, path); err == nil {
			return nil
		} else {
			reasons = append(reasons, err.Error())
		}
	}
	sort.Strings(reasons)
	if len(reasons) == 0 {
		return fmt.Errorf("%s does not match %s", path, branchType)
	}
	return fmt.Errorf("%s does not match %s: %s", path, branchType, reasons[0])
}

func shouldValidateObjectSubset(value any, schema map[string]any, types []string) bool {
	return containsString(types, "object") || len(schemaObjectProperties(schema)) > 0 || len(schemaStringList(schema["required"])) > 0 || isAdditionalPropertiesFalse(schema)
}

func validateJSONObjectSubset(obj map[string]any, schema map[string]any, path string) error {
	props := schemaObjectProperties(schema)
	for _, key := range schemaStringList(schema["required"]) {
		if _, ok := obj[key]; !ok {
			return fmt.Errorf("%s.%s is required", path, key)
		}
	}
	if isAdditionalPropertiesFalse(schema) {
		for key := range obj {
			if _, ok := props[key]; !ok {
				return fmt.Errorf("%s.%s is not allowed by additionalProperties=false", path, key)
			}
		}
	}
	for key, propSchema := range props {
		child, ok := obj[key]
		if !ok {
			continue
		}
		if err := validateJSONSchemaSubset(child, propSchema, path+"."+key); err != nil {
			return err
		}
	}
	return nil
}

func shouldValidateArraySubset(value any, schema map[string]any, types []string) bool {
	return containsString(types, "array") || schemaArrayItemSchema(schema) != nil
}

func isAdditionalPropertiesFalse(schema map[string]any) bool {
	allowed, ok := schema["additionalProperties"].(bool)
	return ok && !allowed
}

func schemaTypeNames(value any) []string {
	switch v := value.(type) {
	case string:
		if strings.TrimSpace(v) == "" {
			return nil
		}
		return []string{strings.TrimSpace(v)}
	case []any:
		out := make([]string, 0, len(v))
		for _, item := range v {
			if s, ok := item.(string); ok && strings.TrimSpace(s) != "" {
				out = append(out, strings.TrimSpace(s))
			}
		}
		return out
	default:
		return nil
	}
}

func schemaStringList(value any) []string {
	raw, ok := value.([]any)
	if !ok {
		return nil
	}
	out := make([]string, 0, len(raw))
	for _, item := range raw {
		if s, ok := item.(string); ok && strings.TrimSpace(s) != "" {
			out = append(out, strings.TrimSpace(s))
		}
	}
	return out
}

func schemaEnumValues(value any) []any {
	raw, ok := value.([]any)
	if !ok {
		return nil
	}
	return raw
}

func valueInJSONEnum(value any, allowed []any) bool {
	for _, item := range allowed {
		if jsonValuesEqual(value, item) {
			return true
		}
	}
	return false
}

func jsonValuesEqual(a, b any) bool {
	encodedA, errA := json.Marshal(a)
	encodedB, errB := json.Marshal(b)
	return errA == nil && errB == nil && string(encodedA) == string(encodedB)
}

func jsonValueMatchesAnyType(value any, types []string) bool {
	for _, t := range types {
		if jsonValueMatchesType(value, t) {
			return true
		}
	}
	return false
}

func jsonValueMatchesType(value any, typ string) bool {
	switch typ {
	case "object":
		_, ok := value.(map[string]any)
		return ok
	case "array":
		_, ok := value.([]any)
		return ok
	case "string":
		_, ok := value.(string)
		return ok
	case "boolean":
		_, ok := value.(bool)
		return ok
	case "number":
		return isJSONNumber(value)
	case "integer":
		return isJSONInteger(value)
	case "null":
		return value == nil
	default:
		return true
	}
}

func isJSONNumber(value any) bool {
	switch value.(type) {
	case float64, float32, int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64:
		return true
	default:
		return false
	}
}

func isJSONInteger(value any) bool {
	switch v := value.(type) {
	case float64:
		return math.Trunc(v) == v
	case float32:
		return math.Trunc(float64(v)) == float64(v)
	case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64:
		return true
	default:
		return false
	}
}

func jsonValueTypeName(value any) string {
	switch value.(type) {
	case nil:
		return "null"
	case map[string]any:
		return "object"
	case []any:
		return "array"
	case string:
		return "string"
	case bool:
		return "boolean"
	default:
		if isJSONInteger(value) {
			return "integer"
		}
		if isJSONNumber(value) {
			return "number"
		}
		return fmt.Sprintf("%T", value)
	}
}

func jsonValuePreview(value any) string {
	encoded, err := json.Marshal(value)
	if err != nil {
		return fmt.Sprintf("%v", value)
	}
	return string(encoded)
}

func containsString(values []string, want string) bool {
	for _, value := range values {
		if value == want {
			return true
		}
	}
	return false
}