package exec

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

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

const draft2020URI = "https://json-schema.org/draft/2020-12/schema"

// decodeStrict decodes JSON with UseNumber so integers beyond 2^53 keep full
// precision (jsonschema/v6 understands json.Number natively), and rejects
// trailing JSON after the top-level value.
func decodeStrict(raw json.RawMessage) (any, error) {
	dec := json.NewDecoder(bytes.NewReader(raw))
	dec.UseNumber()
	var v any
	if err := dec.Decode(&v); err != nil {
		return nil, err
	}
	// second Decode must hit EOF: rejects any trailing JSON value.
	if err := dec.Decode(new(struct{})); err != io.EOF {
		return nil, fmt.Errorf("unexpected trailing JSON")
	}
	return v, nil
}

// assertDraft2020 rejects any $schema dialect other than Draft 2020-12, anywhere
// in the document — including nested resources like $defs (a root-only check is
// bypassed by a nested $id+$schema pinning draft-07, which would turn format
// into an assertion and break catalog-wide parity). $schema must be a string.
func assertDraft2020(doc any) error { return walkDraft2020(doc) }

// JSON Schema keywords grouped by value shape, so walkDraft2020 descends only
// into schema-valued positions — never into data keywords (const/enum/examples/
// default) where a literal "$schema" string is ordinary application data.
var (
	namedSchemaMapKeywords = []string{"$defs", "properties", "patternProperties", "dependentSchemas"} // value: map[name]schema
	singleSchemaKeywords   = []string{"additionalProperties", "items", "contains", "unevaluatedItems",
		"if", "then", "else", "not", "propertyNames", "unevaluatedProperties",
		"contentSchema"} // value: schema (or bool)
	schemaArrayKeywords = []string{"oneOf", "allOf", "anyOf", "prefixItems"} // value: []schema
)

func walkDraft2020(node any) error {
	m, ok := node.(map[string]any)
	if !ok {
		return nil // bool schemas and non-object values have no $schema
	}
	if raw, ok := m["$schema"]; ok {
		s, isStr := raw.(string)
		if !isStr {
			return fmt.Errorf("$schema must be a string")
		}
		if s != draft2020URI {
			return fmt.Errorf("schema declares $schema %q; only Draft 2020-12 (%s) is allowed", s, draft2020URI)
		}
	}
	// Reject pre-2020 dialect keywords outright: under Draft 2020-12 these are
	// replaced by $defs/dependentSchemas+dependentRequired/prefixItems, and an
	// older-dialect resource can silently switch dialect.
	for _, keyword := range []string{"definitions", "dependencies", "additionalItems"} {
		if _, ok := m[keyword]; ok {
			return fmt.Errorf("pre-2020 dialect keyword %q not allowed in Draft 2020-12 (use $defs/dependentSchemas/prefixItems)", keyword)
		}
	}
	for _, kw := range namedSchemaMapKeywords {
		if v, ok := m[kw].(map[string]any); ok {
			for _, sub := range v { // each named subschema
				if err := walkDraft2020(sub); err != nil {
					return err
				}
			}
		}
	}
	for _, kw := range singleSchemaKeywords {
		if v, ok := m[kw]; ok {
			if err := walkDraft2020(v); err != nil {
				return err
			}
		}
	}
	for _, kw := range schemaArrayKeywords {
		if arr, ok := m[kw].([]any); ok {
			for _, sub := range arr {
				if err := walkDraft2020(sub); err != nil {
					return err
				}
			}
		}
	}
	return nil
}

// compileSchema compiles a Draft 2020-12 JSON Schema and returns the compiled
// validator. Malformed, illegal, or wrong-dialect schemas fail descriptor
// catalog construction. Catalog reuses the compiled validator for
// route input validation, keeping construction and request validation identical
// with no hand-rolled subset. format stays annotation-only (no
// AssertFormat) to match current MCP input leniency.
func compileSchema(schema json.RawMessage) (*jsonschema.Schema, error) {
	doc, err := decodeStrict(schema)
	if err != nil {
		return nil, fmt.Errorf("schema is not valid JSON: %w", err)
	}
	if err := assertDraft2020(doc); err != nil {
		return nil, err
	}
	c := jsonschema.NewCompiler()
	c.DefaultDraft(jsonschema.Draft2020) // pin: lib default (draftLatest) may change across versions
	if err := c.AddResource("schema.json", doc); err != nil {
		return nil, fmt.Errorf("schema compile: %w", err)
	}
	s, err := c.Compile("schema.json")
	if err != nil {
		return nil, fmt.Errorf("schema compile: %w", err)
	}
	return s, nil
}

// validateInstance validates a JSON instance against a compiled schema, using
// UseNumber so large-integer const/enum/minMax/multipleOf constraints are exact.
func validateInstance(compiled *jsonschema.Schema, instance json.RawMessage) error {
	v, err := decodeStrict(instance)
	if err != nil {
		return fmt.Errorf("instance is not valid JSON: %w", err)
	}
	if err := compiled.Validate(v); err != nil {
		return fmt.Errorf("instance fails schema: %w", err)
	}
	return nil
}