package run

import (
	"encoding/json"
	"fmt"
	"io/fs"
	"sort"
	"strings"

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

// result_schema.go — kernel-native structured output for single-loop runs.
//
// When tool.Invocation.ResultSchema is set, the executor validates the loop's
// final answer against the schema and, on violation, sends the loop back for
// bounded corrective rounds ON THE SAME CONVERSATION — the model repairs its
// own answer with full task context, which beats any post-hoc host-side
// reformat (the host can't know what the model meant). Mirrors the FinalGuard
// extra-round mechanics (same model, same observers, small step cap), with one
// deliberate contrast: FinalGuard discards the extra round's text, while a
// schema round's text IS the new candidate answer.
//
// Fail-open posture: a schema that doesn't compile logs a WARN and skips
// validation entirely — the host's pre-dispatch compile is the authoritative
// gate, and the kernel must not brick runs on a host bug. Persistent
// violations keep the last candidate output; hosts layer their own fallback
// validator/envelope on top.

const (
	// resultSchemaMaxBytes caps the schema the kernel will compile.
	resultSchemaMaxBytes = 64 << 10
	// resultSchemaMaxRounds caps corrective rounds per run.
	resultSchemaMaxRounds = 2
	// resultSchemaRoundMaxSteps caps each corrective round. Producing a
	// corrected JSON document needs no tools; 2 steps tolerates one stray
	// tool call before the reply.
	resultSchemaRoundMaxSteps = 2
	// resultSchemaExtractCap bounds the output size scanned for a JSON
	// document (the balanced-span scan is quadratic in the worst case).
	resultSchemaExtractCap = 64 << 10
)

// denyLoader fails every $ref resolution. jsonschema/v6's default compiler
// resolves file:// refs from local disk (and can be configured for http) —
// a caller-supplied schema must never reach either.
type denyLoader struct{}

func (denyLoader) Load(url string) (any, error) {
	return nil, fmt.Errorf("external $ref %q refused: remote/file schema loading is disabled: %w", url, fs.ErrNotExist)
}

// compileResultSchema compiles a raw schema with external loading disabled.
func compileResultSchema(raw json.RawMessage) (*jsonschema.Schema, error) {
	if len(raw) > resultSchemaMaxBytes {
		return nil, fmt.Errorf("schema is %d bytes (cap %d)", len(raw), resultSchemaMaxBytes)
	}
	doc, err := jsonschema.UnmarshalJSON(strings.NewReader(string(raw)))
	if err != nil {
		return nil, fmt.Errorf("schema is not valid JSON: %w", err)
	}
	c := jsonschema.NewCompiler()
	c.UseLoader(denyLoader{})
	if err := c.AddResource("inline://result-schema", doc); err != nil {
		return nil, err
	}
	return c.Compile("inline://result-schema")
}

// extractResultJSON finds the answer's JSON document: the whole string, a
// fenced block, or the first parseable balanced {...}/[...] span scanning
// left to right. Input is capped at resultSchemaExtractCap.
func extractResultJSON(s string) (json.RawMessage, bool) {
	s = strings.TrimSpace(s)
	if len(s) > resultSchemaExtractCap {
		s = s[:resultSchemaExtractCap]
	}
	if json.Valid([]byte(s)) && len(s) > 0 {
		return json.RawMessage(s), true
	}
	// Fenced block: ```json ... ``` or bare ``` ... ```.
	for _, fence := range []string{"```json", "```"} {
		if i := strings.Index(s, fence); i >= 0 {
			rest := s[i+len(fence):]
			if j := strings.Index(rest, "```"); j >= 0 {
				candidate := strings.TrimSpace(rest[:j])
				if json.Valid([]byte(candidate)) && candidate != "" {
					return json.RawMessage(candidate), true
				}
			}
		}
	}
	// First parseable balanced span.
	for i := 0; i < len(s); i++ {
		if s[i] != '{' && s[i] != '[' {
			continue
		}
		if end, ok := balancedSpanEnd(s, i); ok {
			candidate := s[i : end+1]
			if json.Valid([]byte(candidate)) {
				return json.RawMessage(candidate), true
			}
		}
	}
	return nil, false
}

// balancedSpanEnd returns the index of the bracket closing the span opened
// at start, honoring JSON string/escape rules.
func balancedSpanEnd(s string, start int) (int, bool) {
	depth := 0
	inStr := false
	esc := false
	for i := start; i < len(s); i++ {
		c := s[i]
		if inStr {
			switch {
			case esc:
				esc = false
			case c == '\\':
				esc = true
			case c == '"':
				inStr = false
			}
			continue
		}
		switch c {
		case '"':
			inStr = true
		case '{', '[':
			depth++
		case '}', ']':
			depth--
			if depth == 0 {
				return i, true
			}
		}
	}
	return 0, false
}

// validateResultJSON returns flattened error paths, nil when the doc conforms.
func validateResultJSON(sch *jsonschema.Schema, doc json.RawMessage) []string {
	inst, err := jsonschema.UnmarshalJSON(strings.NewReader(string(doc)))
	if err != nil {
		return []string{"document is not valid JSON: " + err.Error()}
	}
	err = sch.Validate(inst)
	if err == nil {
		return nil
	}
	var ve *jsonschema.ValidationError
	if ok := asValidationError(err, &ve); !ok {
		return []string{err.Error()}
	}
	var out []string
	flattenValidationError(ve, &out)
	sort.Strings(out)
	if len(out) == 0 {
		out = []string{ve.Error()}
	}
	return out
}

func asValidationError(err error, target **jsonschema.ValidationError) bool {
	ve, ok := err.(*jsonschema.ValidationError)
	if ok {
		*target = ve
	}
	return ok
}

// flattenValidationError collects leaf causes as "/path: message" lines.
func flattenValidationError(ve *jsonschema.ValidationError, out *[]string) {
	if len(ve.Causes) == 0 {
		loc := "/" + strings.Join(ve.InstanceLocation, "/")
		if len(ve.InstanceLocation) == 0 {
			loc = "/"
		}
		*out = append(*out, fmt.Sprintf("%s: %v", loc, ve.ErrorKind))
		return
	}
	for _, c := range ve.Causes {
		flattenValidationError(c, out)
	}
}

// resultSchemaCorrection builds the corrective user turn for one round.
func resultSchemaCorrection(raw json.RawMessage, errs []string) string {
	return fmt.Sprintf(`Your final answer must be a SINGLE JSON document conforming to this JSON Schema (no prose, no code fences):
%s

Your previous answer failed validation:
- %s

Reply now with ONLY the corrected JSON document. Preserve your answer's content; do not invent data; use null or empty values for anything missing.`,
		string(raw), strings.Join(errs, "\n- "))
}