promptkit/internal/validate/standard_validator.go

package validate

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"

	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
	"gitea.maximumdirect.net/eric/promptkit/internal/filecatalog"
	"github.com/santhosh-tekuri/jsonschema/v6"
)

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

const schemaReadChunkSize = 64 * 1024

// StandardValidator provides basic, JSON, and JSON Schema output validation.
type StandardValidator struct {
	schemaBaseDir string
}

type FSValidator struct {
	fsys fs.FS
	root string
}

func NewStandardValidator(schemaBaseDir string) Validator {
	return &StandardValidator{schemaBaseDir: schemaBaseDir}
}

func NewFSValidator(fsys fs.FS, root string) Validator {
	return &FSValidator{fsys: fsys, root: root}
}

func (v *StandardValidator) Validate(ctx context.Context, artifact *domain.Artifact, contract domain.OutputContract) (domain.ValidationResult, error) {
	return validateArtifact(ctx, artifact, contract, v.validateJSONSchema)
}

func (v *FSValidator) Validate(ctx context.Context, artifact *domain.Artifact, contract domain.OutputContract) (domain.ValidationResult, error) {
	return validateArtifact(ctx, artifact, contract, v.validateJSONSchema)
}

type preparedValidation struct {
	contract       domain.OutputContract
	schemaDocument any
	schema         schemaExecutor
}

func (p *preparedValidation) Validate(ctx context.Context, artifact *domain.Artifact) (domain.ValidationResult, error) {
	return validateArtifact(ctx, artifact, p.contract, p.validateJSONSchema)
}

func (p *preparedValidation) SchemaDocument() any {
	return p.schemaDocument
}

func (p *preparedValidation) validateJSONSchema(ctx context.Context, instance any, _ string) ([]string, error) {
	if p.schema == nil {
		return nil, errors.New("prepared JSON schema is unavailable")
	}
	return executeJSONSchema(ctx, p.schema, instance)
}

func (v *StandardValidator) PrepareValidation(ctx context.Context, contract domain.OutputContract) (PreparedValidation, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}

	prepared := &preparedValidation{contract: contract}
	if contract.ValidationMode != domain.ValidationJSONSchema {
		return prepared, nil
	}

	resolvedSchemaPath, err := v.resolveSchemaPath(ctx, contract.SchemaPath)
	if err != nil {
		return nil, err
	}
	schemaDocument, err := loadJSONSchemaFile(ctx, resolvedSchemaPath)
	if err != nil {
		return nil, fmt.Errorf("failed to compile JSON schema %q: %w", resolvedSchemaPath, err)
	}
	schemaRoot, err := v.schemaRoot()
	if err != nil {
		return nil, err
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	compiler := newSchemaCompiler(standardSchemaLoader{ctx: ctx, root: schemaRoot})
	resourceURL := fileSchemaResourceURL(resolvedSchemaPath)
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if err := compiler.AddResource(resourceURL.String(), schemaDocument); err != nil {
		if contextErr := ctx.Err(); contextErr != nil {
			return nil, contextErr
		}
		return nil, fmt.Errorf("failed to register JSON schema %q: %w", resolvedSchemaPath, err)
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	schema, err := compileJSONSchema(ctx, resourceURL.String(), compiler.Compile)
	if err != nil {
		return nil, fmt.Errorf("failed to compile JSON schema %q: %w", resolvedSchemaPath, err)
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}

	prepared.schemaDocument = schemaDocument
	prepared.schema = schema
	return prepared, nil
}

func (v *FSValidator) PrepareValidation(ctx context.Context, contract domain.OutputContract) (PreparedValidation, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}

	prepared := &preparedValidation{contract: contract}
	if contract.ValidationMode != domain.ValidationJSONSchema {
		return prepared, nil
	}

	schemaName, schemaDocument, err := v.loadSchemaDocument(ctx, contract.SchemaPath)
	if err != nil {
		return nil, err
	}
	resourceURL := fsSchemaResourceURL(schemaName)
	compiler := newSchemaCompiler(fsSchemaLoader{ctx: ctx, fsys: v.fsys, root: filecatalog.CleanFSRoot(v.root)})
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if err := compiler.AddResource(resourceURL.String(), schemaDocument); err != nil {
		if contextErr := ctx.Err(); contextErr != nil {
			return nil, contextErr
		}
		return nil, fmt.Errorf("failed to register JSON schema %q: %w", schemaName, err)
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	schema, err := compileJSONSchema(ctx, resourceURL.String(), compiler.Compile)
	if err != nil {
		return nil, fmt.Errorf("failed to compile JSON schema %q: %w", schemaName, err)
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}

	prepared.schemaDocument = schemaDocument
	prepared.schema = schema
	return prepared, nil
}

type schemaValidatorFunc func(ctx context.Context, instance any, schemaPath string) ([]string, error)

func validateArtifact(ctx context.Context, artifact *domain.Artifact, contract domain.OutputContract, validateSchema schemaValidatorFunc) (domain.ValidationResult, error) {
	select {
	case <-ctx.Done():
		return domain.ValidationResult{}, ctx.Err()
	default:
	}

	res := domain.ValidationResult{
		Mode:           contract.ValidationMode,
		SchemaPath:     contract.SchemaPath,
		RepairAttempts: contract.RepairAttempts,
	}

	if artifact == nil {
		return domain.ValidationResult{}, errors.New("artifact is required for validation")
	}

	switch contract.ValidationMode {
	case domain.ValidationNone:
		res.Status = domain.ValidationSkipped
		res.IsValid = true
		return res, nil
	case domain.ValidationBasic:
		empty := strings.TrimSpace(string(artifact.Body)) == ""
		if err := ctx.Err(); err != nil {
			return domain.ValidationResult{}, err
		}
		if empty {
			res.Status = domain.ValidationFailed
			res.IsValid = false
			res.Errors = []string{"output is empty"}
			return res, nil
		}
		res.Status = domain.ValidationPassed
		res.IsValid = true
		return res, nil
	case domain.ValidationJSON:
		valid := json.Valid(artifact.Body)
		if err := ctx.Err(); err != nil {
			return domain.ValidationResult{}, err
		}
		if !valid {
			res.Status = domain.ValidationFailed
			res.IsValid = false
			res.Errors = []string{"invalid JSON"}
			return res, nil
		}
		res.Status = domain.ValidationPassed
		res.IsValid = true
		return res, nil
	case domain.ValidationJSONSchema:
		instance, jsonErr := decodeJSONValue(ctx, artifact.Body)
		if jsonErr != nil {
			if contextErr := ctx.Err(); contextErr != nil {
				return domain.ValidationResult{}, contextErr
			}
			res.Status = domain.ValidationFailed
			res.IsValid = false
			res.Errors = []string{fmt.Sprintf("invalid JSON: %v", jsonErr)}
			return res, nil
		}

		validationErrors, err := validateSchema(ctx, instance, contract.SchemaPath)
		if err != nil {
			return domain.ValidationResult{}, err
		}
		if len(validationErrors) > 0 {
			res.Status = domain.ValidationFailed
			res.IsValid = false
			res.Errors = validationErrors
			return res, nil
		}

		res.Status = domain.ValidationPassed
		res.IsValid = true
		return res, nil
	default:
		return domain.ValidationResult{}, fmt.Errorf("unsupported validation mode: %q", contract.ValidationMode)
	}
}

func (v *StandardValidator) validateJSONSchema(ctx context.Context, instance any, schemaPath string) ([]string, error) {
	resolvedSchemaPath, err := v.resolveSchemaPath(ctx, schemaPath)
	if err != nil {
		return nil, err
	}

	schemaRoot, err := v.schemaRoot()
	if err != nil {
		return nil, err
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	compiler := newSchemaCompiler(standardSchemaLoader{ctx: ctx, root: schemaRoot})
	resourceURL := fileSchemaResourceURL(resolvedSchemaPath)
	schema, err := compileJSONSchema(ctx, resourceURL.String(), compiler.Compile)
	if err != nil {
		return nil, fmt.Errorf("failed to compile JSON schema %q: %w", resolvedSchemaPath, err)
	}

	return executeJSONSchema(ctx, schema, instance)
}

func (v *FSValidator) validateJSONSchema(ctx context.Context, instance any, schemaPath string) ([]string, error) {
	schemaName, schemaDoc, err := v.loadSchemaDocument(ctx, schemaPath)
	if err != nil {
		return nil, err
	}

	resourceURL := fsSchemaResourceURL(schemaName)
	if err := validateSchemaDialect(schemaDoc); err != nil {
		return nil, fmt.Errorf("failed to compile JSON schema %q: %w", schemaName, err)
	}
	compiler := newSchemaCompiler(fsSchemaLoader{ctx: ctx, fsys: v.fsys, root: filecatalog.CleanFSRoot(v.root)})
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if err := compiler.AddResource(resourceURL.String(), schemaDoc); err != nil {
		if contextErr := ctx.Err(); contextErr != nil {
			return nil, contextErr
		}
		return nil, fmt.Errorf("failed to register JSON schema %q: %w", schemaName, err)
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	schema, err := compileJSONSchema(ctx, resourceURL.String(), compiler.Compile)
	if err != nil {
		return nil, fmt.Errorf("failed to compile JSON schema %q: %w", schemaName, err)
	}

	return executeJSONSchema(ctx, schema, instance)
}

func decodeJSONValue(ctx context.Context, body []byte) (any, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	decoder := json.NewDecoder(bytes.NewReader(body))
	decoder.UseNumber()

	var value any
	decodeErr := decoder.Decode(&value)
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if decodeErr != nil {
		return nil, decodeErr
	}

	var trailing any
	trailingErr := decoder.Decode(&trailing)
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if errors.Is(trailingErr, io.EOF) {
		return value, nil
	} else if trailingErr != nil {
		return nil, trailingErr
	}
	return nil, errors.New("multiple JSON values")
}

func (v *StandardValidator) resolveSchemaPath(ctx context.Context, schemaPath string) (string, error) {
	if err := ctx.Err(); err != nil {
		return "", err
	}
	if strings.TrimSpace(schemaPath) == "" {
		return "", errors.New("schema path is required for json_schema validation")
	}

	root, err := v.schemaRoot()
	if err != nil {
		return "", err
	}
	if err := ctx.Err(); err != nil {
		return "", err
	}
	resolved, err := containedFilesystemPath(root, schemaPath)
	if err != nil {
		return "", err
	}
	if err := ctx.Err(); err != nil {
		return "", err
	}
	if _, err := os.Stat(resolved); err != nil {
		if contextErr := ctx.Err(); contextErr != nil {
			return "", contextErr
		}
		return "", fmt.Errorf("failed to access schema file %q: %w", resolved, err)
	}
	if err := ctx.Err(); err != nil {
		return "", err
	}

	return resolved, nil
}

func (v *StandardValidator) schemaRoot() (string, error) {
	root := v.schemaBaseDir
	if strings.TrimSpace(root) == "" {
		root = "."
	}
	absolute, err := filepath.Abs(root)
	if err != nil {
		return "", fmt.Errorf("failed to resolve schema source %q: %w", root, err)
	}
	resolved, err := filepath.EvalSymlinks(absolute)
	if err != nil {
		return "", fmt.Errorf("failed to access schema source %q: %w", root, err)
	}
	return resolved, nil
}

func (v *FSValidator) loadSchemaDocument(ctx context.Context, schemaPath string) (string, any, error) {
	resolved, err := v.resolveSchemaPath(ctx, schemaPath)
	if err != nil {
		return "", nil, err
	}

	raw, err := readSchemaFile(ctx, func() (fs.File, error) {
		return v.fsys.Open(resolved)
	})
	if err != nil {
		return "", nil, fmt.Errorf("failed to read schema file %q: %w", resolved, err)
	}

	doc, err := decodeJSONValue(ctx, raw)
	if err != nil {
		return "", nil, fmt.Errorf("failed to decode JSON schema %q: %w", resolved, err)
	}
	if err := validateSchemaDialect(doc); err != nil {
		if contextErr := ctx.Err(); contextErr != nil {
			return "", nil, contextErr
		}
		return "", nil, fmt.Errorf("failed to decode JSON schema %q: %w", resolved, err)
	}
	if err := ctx.Err(); err != nil {
		return "", nil, err
	}
	return resolved, doc, nil
}

func (v *FSValidator) resolveSchemaPath(ctx context.Context, schemaPath string) (string, error) {
	if err := ctx.Err(); err != nil {
		return "", err
	}
	if strings.TrimSpace(schemaPath) == "" {
		return "", errors.New("schema path is required for json_schema validation")
	}
	if v.fsys == nil {
		return "", errors.New("schema filesystem is nil")
	}

	cleanRoot := filecatalog.CleanFSRoot(v.root)
	rootInfo, err := fs.Stat(v.fsys, cleanRoot)
	if err != nil {
		if contextErr := ctx.Err(); contextErr != nil {
			return "", contextErr
		}
		return "", fmt.Errorf("failed to access schema source %q: %w", cleanRoot, err)
	}
	if err := ctx.Err(); err != nil {
		return "", err
	}

	var resolved string
	if rootInfo.IsDir() {
		resolvedPath, _, err := filecatalog.ResolveFSPath(cleanRoot, cleanRoot, schemaPath)
		if err != nil {
			return "", err
		}
		resolved = resolvedPath
	} else {
		cleanSchemaPath, err := cleanSchemaFSPath(schemaPath)
		if err != nil {
			return "", err
		}
		if cleanSchemaPath != strings.TrimSpace(path.Base(cleanRoot)) {
			return "", fmt.Errorf("schema path %q does not match schema file %q", cleanSchemaPath, path.Base(cleanRoot))
		}
		resolved = cleanRoot
	}

	if _, err := fs.Stat(v.fsys, resolved); err != nil {
		if contextErr := ctx.Err(); contextErr != nil {
			return "", contextErr
		}
		return "", fmt.Errorf("failed to access schema file %q: %w", resolved, err)
	}
	if err := ctx.Err(); err != nil {
		return "", err
	}
	return resolved, nil
}

func cleanSchemaFSPath(schemaPath string) (string, error) {
	cleaned := strings.TrimSpace(schemaPath)
	if cleaned == "" {
		return "", errors.New("schema path is required for json_schema validation")
	}
	cleaned = path.Clean(cleaned)
	if path.IsAbs(cleaned) {
		return "", fmt.Errorf("schema path %q must be relative", schemaPath)
	}
	return cleaned, nil
}

func fileSchemaResourceURL(schemaName string) *url.URL {
	filePath := filepath.ToSlash(schemaName)
	if runtime.GOOS == "windows" && !strings.HasPrefix(filePath, "/") {
		filePath = "/" + filePath
	}
	return &url.URL{Scheme: "file", Path: filePath}
}

func fsSchemaResourceURL(schemaName string) *url.URL {
	return &url.URL{
		Scheme: "promptkit-schema",
		Path:   "/" + strings.TrimPrefix(path.Clean(schemaName), "/"),
	}
}

func newSchemaCompiler(loader jsonschema.URLLoader) *jsonschema.Compiler {
	compiler := jsonschema.NewCompiler()
	compiler.DefaultDraft(jsonschema.Draft2020)
	compiler.UseLoader(loader)
	return compiler
}

type schemaExecutor interface {
	Validate(instance any) error
}

func compileJSONSchema(ctx context.Context, resourceURL string, compile func(string) (*jsonschema.Schema, error)) (*jsonschema.Schema, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	schema, compileErr := compile(resourceURL)
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	return schema, compileErr
}

func executeJSONSchema(ctx context.Context, schema schemaExecutor, instance any) ([]string, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	validationErr := schema.Validate(instance)
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if validationErr != nil {
		return []string{fmt.Sprintf("json schema validation failed: %v", validationErr)}, nil
	}
	return nil, nil
}

func validateSchemaDialect(doc any) error {
	object, ok := doc.(map[string]any)
	if !ok {
		return nil
	}
	value, ok := object["$schema"]
	if !ok {
		return nil
	}
	dialect, ok := value.(string)
	if !ok {
		return errors.New("$schema must be a string")
	}
	if dialect != jsonSchemaDraft2020 && dialect != jsonSchemaDraft2020+"#" {
		return fmt.Errorf("unsupported JSON Schema dialect %q; expected %q", dialect, jsonSchemaDraft2020)
	}
	return nil
}

type standardSchemaLoader struct {
	ctx  context.Context
	root string
}

func (l standardSchemaLoader) Load(resourceURL string) (any, error) {
	if err := l.ctx.Err(); err != nil {
		return nil, err
	}
	parsed, err := url.Parse(resourceURL)
	if err != nil || parsed.Scheme != "file" || parsed.Host != "" || parsed.RawQuery != "" || parsed.Opaque != "" {
		return nil, fmt.Errorf("schema reference %q is not a contained file reference", resourceURL)
	}
	fileName, err := (jsonschema.FileLoader{}).ToFile(resourceURL)
	if err != nil {
		return nil, fmt.Errorf("schema reference %q is not a contained file reference: %w", resourceURL, err)
	}
	resolved, err := containedFilesystemPath(l.root, fileName)
	if err != nil {
		if contextErr := l.ctx.Err(); contextErr != nil {
			return nil, contextErr
		}
		return nil, err
	}
	if err := l.ctx.Err(); err != nil {
		return nil, err
	}
	return loadJSONSchemaFile(l.ctx, resolved)
}

func containedFilesystemPath(root, name string) (string, error) {
	candidate := name
	if !filepath.IsAbs(candidate) {
		candidate = filepath.Join(root, candidate)
	}
	candidate, err := filepath.Abs(candidate)
	if err != nil {
		return "", fmt.Errorf("failed to resolve schema path %q: %w", name, err)
	}
	candidate, err = filepath.EvalSymlinks(candidate)
	if err != nil {
		return "", fmt.Errorf("failed to access schema file %q: %w", candidate, err)
	}
	relative, err := filepath.Rel(root, candidate)
	if err != nil {
		return "", fmt.Errorf("failed to compare schema path %q with source root: %w", candidate, err)
	}
	if relative == ".." || strings.HasPrefix(relative, ".."+string(filepath.Separator)) {
		return "", fmt.Errorf("schema path %q escapes source root", name)
	}
	return candidate, nil
}

func loadJSONSchemaFile(ctx context.Context, name string) (any, error) {
	raw, err := readSchemaFile(ctx, func() (fs.File, error) {
		return os.Open(name)
	})
	if err != nil {
		return nil, err
	}
	doc, err := decodeJSONValue(ctx, raw)
	if err != nil {
		return nil, err
	}
	if err := validateSchemaDialect(doc); err != nil {
		if contextErr := ctx.Err(); contextErr != nil {
			return nil, contextErr
		}
		return nil, err
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	return doc, nil
}

func readSchemaFile(ctx context.Context, open func() (fs.File, error)) ([]byte, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	file, openErr := open()
	if err := ctx.Err(); err != nil {
		if file != nil {
			_ = file.Close()
		}
		return nil, err
	}
	if openErr != nil {
		if file != nil {
			_ = file.Close()
		}
		return nil, openErr
	}
	if file == nil {
		return nil, errors.New("schema source returned a nil file")
	}
	defer file.Close()

	var contents []byte
	chunk := make([]byte, schemaReadChunkSize)
	for {
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		n, readErr := file.Read(chunk)
		if n > 0 {
			contents = append(contents, chunk[:n]...)
		}
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		if errors.Is(readErr, io.EOF) {
			return contents, nil
		}
		if readErr != nil {
			return nil, readErr
		}
		if n == 0 {
			return nil, io.ErrNoProgress
		}
	}
}

type fsSchemaLoader struct {
	ctx  context.Context
	fsys fs.FS
	root string
}

func (l fsSchemaLoader) Load(resourceURL string) (any, error) {
	if err := l.ctx.Err(); err != nil {
		return nil, err
	}
	parsed, err := url.Parse(resourceURL)
	if err != nil {
		return nil, fmt.Errorf("invalid schema reference %q: %w", resourceURL, err)
	}
	if parsed.Scheme != "promptkit-schema" || parsed.Host != "" || parsed.RawQuery != "" || parsed.Opaque != "" {
		return nil, fmt.Errorf("schema reference %q is not allowed", resourceURL)
	}
	name := strings.TrimPrefix(parsed.Path, "/")
	name = path.Clean(name)
	if l.root == "." {
		if strings.HasPrefix(name, "../") || name == ".." {
			return nil, fmt.Errorf("schema reference %q escapes source root", resourceURL)
		}
	} else if name != l.root && !strings.HasPrefix(name, l.root+"/") {
		return nil, fmt.Errorf("schema reference %q escapes source root", resourceURL)
	}

	rootInfo, err := fs.Stat(l.fsys, l.root)
	if err != nil {
		if contextErr := l.ctx.Err(); contextErr != nil {
			return nil, contextErr
		}
		return nil, err
	}
	if err := l.ctx.Err(); err != nil {
		return nil, err
	}
	if !rootInfo.IsDir() && name != l.root {
		return nil, fmt.Errorf("schema reference %q is outside the configured schema file", resourceURL)
	}

	raw, err := readSchemaFile(l.ctx, func() (fs.File, error) {
		return l.fsys.Open(name)
	})
	if err != nil {
		return nil, err
	}
	doc, err := decodeJSONValue(l.ctx, raw)
	if err != nil {
		return nil, err
	}
	if err := validateSchemaDialect(doc); err != nil {
		if contextErr := l.ctx.Err(); contextErr != nil {
			return nil, contextErr
		}
		return nil, err
	}
	if err := l.ctx.Err(); err != nil {
		return nil, err
	}
	return doc, nil
}