fq/pkg/interp/interp.go

package interp

import (
	"bytes"
	"context"
	"crypto/md5"
	"embed"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"math/big"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/mitchellh/mapstructure"
	"github.com/wader/fq/format/registry"
	"github.com/wader/fq/internal/ansi"
	"github.com/wader/fq/internal/colorjson"
	"github.com/wader/fq/internal/ctxstack"
	"github.com/wader/fq/internal/ioextra"
	"github.com/wader/fq/internal/num"
	"github.com/wader/fq/internal/pos"
	"github.com/wader/fq/pkg/bitio"
	"github.com/wader/fq/pkg/decode"

	"github.com/wader/gojq"
)

//go:embed interp.jq
//go:embed internal.jq
//go:embed funcs.jq
//go:embed grep.jq
//go:embed options.jq
//go:embed args.jq
//go:embed query.jq
//go:embed repl.jq
//go:embed formats.jq
var builtinFS embed.FS

var initSource = `include "@builtin/interp";`

type valueError struct {
	v interface{}
}

func (v valueError) Error() string      { return fmt.Sprintf("error: %v", v.v) }
func (v valueError) Value() interface{} { return v.v }

type compileError struct {
	err      error
	what     string
	filename string
	pos      pos.Pos
}

func (ce compileError) Value() interface{} {
	return map[string]interface{}{
		"error":    ce.err.Error(),
		"what":     ce.what,
		"filename": ce.filename,
		"line":     ce.pos.Line,
		"column":   ce.pos.Column,
	}
}
func (ce compileError) Error() string {
	filename := ce.filename
	if filename == "" {
		filename = "src"
	}
	return fmt.Sprintf("%s:%d:%d: %s: %s", filename, ce.pos.Line, ce.pos.Column, ce.what, ce.err.Error())
}

var ErrEOF = io.EOF
var ErrInterrupt = errors.New("Interrupt")

// gojq errors can implement this to signal exit code
type Exiter interface {
	ExitCode() int
}

// gojq halt_error uses this
type IsEmptyErrorer interface {
	IsEmptyError() bool
}

type Terminal interface {
	Size() (int, int)
	IsTerminal() bool
}

type Input interface {
	fs.File
	Terminal
}

type Output interface {
	io.Writer
	Terminal
}

type OS interface {
	Stdin() Input
	Stdout() Output
	Stderr() Output
	InterruptChan() chan struct{}
	Args() []string
	Environ() []string
	ConfigDir() (string, error)
	// FS.File returned by FS().Open() can optionally implement io.Seeker
	FS() fs.FS
	Readline(prompt string, complete func(line string, pos int) (newLine []string, shared int)) (string, error)
	History() ([]string, error)
}

type FixedFileInfo struct {
	FName    string
	FSize    int64
	FMode    fs.FileMode
	FModTime time.Time
	FIsDir   bool
	FSys     interface{}
}

func (ffi FixedFileInfo) Name() string       { return ffi.FName }
func (ffi FixedFileInfo) Size() int64        { return ffi.FSize }
func (ffi FixedFileInfo) Mode() fs.FileMode  { return ffi.FMode }
func (ffi FixedFileInfo) ModTime() time.Time { return ffi.FModTime }
func (ffi FixedFileInfo) IsDir() bool        { return ffi.FIsDir }
func (ffi FixedFileInfo) Sys() interface{}   { return ffi.FSys }

type FileReader struct {
	R        io.Reader
	FileInfo FixedFileInfo
}

func (rf FileReader) Stat() (fs.FileInfo, error) { return rf.FileInfo, nil }
func (rf FileReader) Read(p []byte) (int, error) { return rf.R.Read(p) }
func (FileReader) Close() error                  { return nil }

type Value interface {
	gojq.JQValue

	ExtKeys() []string
}

type Display interface {
	Display(w io.Writer, opts Options) error
}

type ToBufferView interface {
	ToBufferView() (BufferRange, error)
}

type JQValueEx interface {
	JQValueToGoJQEx(optsFn func() Options) interface{}
}

func valuePath(v *decode.Value) []interface{} {
	var parts []interface{}

	for v.Parent != nil {
		switch vv := v.Parent.V.(type) {
		case decode.Compound:
			if vv.IsArray {
				parts = append([]interface{}{v.Index}, parts...)
			} else {
				parts = append([]interface{}{v.Name}, parts...)
			}
		}
		v = v.Parent
	}

	return parts
}

func valuePathDecorated(v *decode.Value, d Decorator) string {
	var parts []string

	for _, p := range valuePath(v) {
		switch p := p.(type) {
		case string:
			parts = append(parts, ".", d.ObjectKey.Wrap(p))
		case int:
			indexStr := strconv.Itoa(p)
			parts = append(parts, fmt.Sprintf("%s%s%s", d.Index.F("["), d.Number.F(indexStr), d.Index.F("]")))
		}
	}

	if len(parts) == 0 {
		return "."
	}

	return strings.Join(parts, "")
}

type iterFn func() (interface{}, bool)

func (i iterFn) Next() (interface{}, bool) { return i() }

type loadModule struct {
	init func() ([]*gojq.Query, error)
	load func(name string) (*gojq.Query, error)
}

func (l loadModule) LoadInitModules() ([]*gojq.Query, error)     { return l.init() }
func (l loadModule) LoadModule(name string) (*gojq.Query, error) { return l.load(name) }

func toString(v interface{}) (string, error) {
	switch v := v.(type) {
	case string:
		return v, nil
	case gojq.JQValue:
		return toString(v.JQValueToGoJQ())
	default:
		b, err := toBytes(v)
		if err != nil {
			return "", fmt.Errorf("value can't be a string")
		}

		return string(b), nil
	}
}

func toBigInt(v interface{}) (*big.Int, error) {
	switch v := v.(type) {
	case int:
		return new(big.Int).SetInt64(int64(v)), nil
	case float64:
		return new(big.Int).SetInt64(int64(v)), nil
	case *big.Int:
		return v, nil
	default:
		return nil, fmt.Errorf("value is not a number")
	}
}

func toBytes(v interface{}) ([]byte, error) {
	switch v := v.(type) {
	// TODO: remove?
	case []byte:
		return v, nil
	default:
		bb, err := toBuffer(v)
		if err != nil {
			return nil, fmt.Errorf("value is not bytes")
		}
		buf := &bytes.Buffer{}
		if _, err := io.Copy(buf, bb); err != nil {
			return nil, err
		}

		return buf.Bytes(), nil
	}
}

func toBuffer(v interface{}) (*bitio.Buffer, error) {
	return toBufferEx(v, false)
}

func toBufferEx(v interface{}, inArray bool) (*bitio.Buffer, error) {
	switch vv := v.(type) {
	case ToBufferView:
		bv, err := vv.ToBufferView()
		if err != nil {
			return nil, err
		}
		return bv.bb.BitBufRange(bv.r.Start, bv.r.Len)
	case string:
		return bitio.NewBufferFromBytes([]byte(vv), -1), nil
	case []byte:
		return bitio.NewBufferFromBytes(vv, -1), nil
	case int, float64, *big.Int:
		bi, err := toBigInt(v)
		if err != nil {
			return nil, err
		}

		if inArray {
			if bi.Cmp(big.NewInt(255)) > 0 || bi.Cmp(big.NewInt(0)) < 0 {
				return nil, fmt.Errorf("buffer byte list must be bytes (0-255) got %v", bi)
			}
			n := bi.Uint64()
			b := [1]byte{byte(n)}
			return bitio.NewBufferFromBytes(b[:], -1), nil
		}

		// TODO: how should this work? "0xf | tobytes" 4bits or 8bits? now 4
		//padBefore := (8 - (bi.BitLen() % 8)) % 8
		padBefore := 0
		bb, err := bitio.NewBufferFromBytes(bi.Bytes(), -1).BitBufRange(int64(padBefore), int64(bi.BitLen()))
		if err != nil {
			return nil, err
		}
		return bb, nil
	case []interface{}:
		var rr []bitio.BitReadAtSeeker
		// TODO: optimize byte array case, flatten into one slice
		for _, e := range vv {
			eBB, eErr := toBufferEx(e, true)
			if eErr != nil {
				return nil, eErr
			}
			rr = append(rr, eBB)
		}

		mb, err := bitio.NewMultiBitReader(rr)
		if err != nil {
			return nil, err
		}

		bb, err := bitio.NewBufferFromBitReadSeeker(mb)
		if err != nil {
			return nil, err
		}

		return bb, nil
	default:
		return nil, fmt.Errorf("value can't be a buffer")
	}
}

func toBufferView(v interface{}) (BufferRange, error) {
	switch vv := v.(type) {
	case ToBufferView:
		return vv.ToBufferView()
	default:
		bb, err := toBuffer(v)
		if err != nil {
			return BufferRange{}, err
		}
		return newBufferRangeFromBuffer(bb, 8), nil
	}
}

func queryErrorPosition(src string, v error) pos.Pos {
	var offset int

	if tokIf, ok := v.(interface{ Token() (string, int) }); ok { //nolint:errorlint
		_, offset = tokIf.Token()
	}
	if offset >= 0 {
		return pos.NewFromOffset(src, offset)
	}
	return pos.Pos{}
}

type Variable struct {
	Name  string
	Value interface{}
}

type Function struct {
	Name     string
	MinArity int
	MaxArity int
	Fn       func(interface{}, []interface{}) interface{}
	IterFn   func(interface{}, []interface{}) gojq.Iter
}

type RunMode int

const (
	ScriptMode RunMode = iota
	REPLMode
	CompletionMode
)

type evalContext struct {
	// structcheck has problems with embedding https://gitlab.com/opennota/check#known-limitations
	ctx    context.Context
	output io.Writer
}

type Interp struct {
	registry       *registry.Registry
	os             OS
	initFqQuery    *gojq.Query
	includeCache   map[string]*gojq.Query
	interruptStack *ctxstack.Stack
	// global state, is ref as Interp i cloned per eval
	state *interface{}

	// new for each run, other values are copied by value
	evalContext evalContext
}

func New(os OS, registry *registry.Registry) (*Interp, error) {
	var err error

	i := &Interp{
		os:       os,
		registry: registry,
	}

	i.includeCache = map[string]*gojq.Query{}
	i.initFqQuery, err = gojq.Parse(initSource)
	if err != nil {
		return nil, fmt.Errorf("init:%s: %w", queryErrorPosition(initSource, err), err)
	}
	// TODO: refactor ctxstack have a CancelTop and return c context to Stop?
	i.interruptStack = ctxstack.New(func(stopCh chan struct{}) {
		select {
		case <-stopCh:
			return
		case <-os.InterruptChan():
			return
		}
	})
	i.state = new(interface{})

	return i, nil
}

func (i *Interp) Stop() {
	// TODO: cancel all run instances?
	i.interruptStack.Stop()
}

func (i *Interp) Main(ctx context.Context, output Output, version string) error {
	var args []interface{}
	for _, a := range i.os.Args() {
		args = append(args, a)
	}

	input := map[string]interface{}{
		"args":    args,
		"version": version,
	}

	iter, err := i.EvalFunc(ctx, input, "_main", nil, output)
	if err != nil {
		fmt.Fprintln(i.os.Stderr(), err)
		return err
	}
	for {
		v, ok := iter.Next()
		if !ok {
			break
		}

		switch v := v.(type) {
		case error:
			if emptyErr, ok := v.(IsEmptyErrorer); ok && emptyErr.IsEmptyError() { //nolint:errorlint
				// no output
			} else if errors.Is(v, context.Canceled) {
				// ignore context cancel here for now, which means user somehow interrupted the interpreter
				// TODO: handle this inside interp.jq instead but then we probably have to do nested
				// eval and or also use different contexts for the interpreter and reading/decoding
			} else {
				fmt.Fprintln(i.os.Stderr(), v)
			}
			return v
		case [2]interface{}:
			fmt.Fprintln(i.os.Stderr(), v[:]...)
		default:
			// TODO: can this happen?
			fmt.Fprintln(i.os.Stderr(), v)
		}
	}

	return nil
}

func (i *Interp) Eval(ctx context.Context, c interface{}, src string, srcFilename string, output io.Writer) (gojq.Iter, error) {
	gq, err := gojq.Parse(src)
	if err != nil {
		p := queryErrorPosition(src, err)
		return nil, compileError{
			err:      err,
			what:     "parse",
			filename: srcFilename,
			pos:      p,
		}
	}

	// make copy of interp
	ci := *i
	ni := &ci
	ni.evalContext = evalContext{}

	var variableNames []string
	var variableValues []interface{}
	for k, v := range i.variables() {
		variableNames = append(variableNames, "$"+k)
		variableValues = append(variableValues, v)
	}

	var funcCompilerOpts []gojq.CompilerOption
	for _, f := range ni.makeFunctions() {
		if f.IterFn != nil {
			funcCompilerOpts = append(funcCompilerOpts,
				gojq.WithIterFunction(f.Name, f.MinArity, f.MaxArity, f.IterFn))
		} else {
			funcCompilerOpts = append(funcCompilerOpts,
				gojq.WithFunction(f.Name, f.MinArity, f.MaxArity, f.Fn))
		}
	}

	compilerOpts := append([]gojq.CompilerOption{}, funcCompilerOpts...)
	compilerOpts = append(compilerOpts, gojq.WithEnvironLoader(ni.os.Environ))
	compilerOpts = append(compilerOpts, gojq.WithVariables(variableNames))
	compilerOpts = append(compilerOpts, gojq.WithModuleLoader(loadModule{
		init: func() ([]*gojq.Query, error) {
			return []*gojq.Query{i.initFqQuery}, nil
		},
		load: func(name string) (*gojq.Query, error) {
			if err := ctx.Err(); err != nil {
				return nil, err
			}

			var filename string
			// support include "nonexisting?" to ignore include error
			var isTry bool
			if strings.HasSuffix(name, "?") {
				isTry = true
				filename = name[0 : len(name)-1]
			} else {
				filename = name
			}
			filename = filename + ".jq"

			pathPrefixes := []struct {
				prefix string
				fn     func(filename string) (io.ReadCloser, error)
			}{
				{
					"@builtin/", func(filename string) (io.ReadCloser, error) {
						return builtinFS.Open(filename)
					},
				},
				{
					"@config/", func(filename string) (io.ReadCloser, error) {
						configDir, err := i.os.ConfigDir()
						if err != nil {
							return nil, err
						}
						return i.os.FS().Open(filepath.Join(configDir, filename))
					},
				},
				{
					"", func(filename string) (io.ReadCloser, error) {
						// TODO: jq $ORIGIN

						if filepath.IsAbs(filename) {
							return i.os.FS().Open(filename)
						}

						for _, path := range append([]string{"./"}, i.includePaths()...) {
							if f, err := i.os.FS().Open(filepath.Join(path, filename)); err == nil {
								return f, nil
							}
						}

						return nil, &fs.PathError{Op: "open", Path: filename, Err: fs.ErrNotExist}
					},
				},
			}

			for _, p := range pathPrefixes {
				if !strings.HasPrefix(filename, p.prefix) {
					continue
				}

				if q, ok := ni.includeCache[filename]; ok {
					return q, nil
				}

				filenamePart := strings.TrimPrefix(filename, p.prefix)
				f, err := p.fn(filenamePart)
				if err != nil {
					if !isTry {
						return nil, err
					}
					f = io.NopCloser(&bytes.Buffer{})
				}
				defer f.Close()

				b, err := io.ReadAll(f)
				if err != nil {
					return nil, err
				}
				s := string(b)
				q, err := gojq.Parse(s)
				if err != nil {
					p := queryErrorPosition(s, err)
					return nil, compileError{
						err:      err,
						what:     "parse",
						filename: filenamePart,
						pos:      p,
					}
				}

				// not identity body means it returns something, threat as dynamic include
				if q.Term.Type != gojq.TermTypeIdentity {
					gc, err := gojq.Compile(q, funcCompilerOpts...)
					if err != nil {
						return nil, err
					}
					iter := gc.RunWithContext(context.Background(), nil)
					var vs []interface{}
					for {
						v, ok := iter.Next()
						if !ok {
							break
						}
						if err, ok := v.(error); ok {
							return nil, err
						}
						vs = append(vs, v)
					}
					if len(vs) != 1 {
						return nil, fmt.Errorf("dynamic include: must output one string, got: %#v", vs)
					}
					s, sOk := vs[0].(string)
					if !sOk {
						return nil, fmt.Errorf("dynamic include: must be string, got %#v", s)
					}
					q, err = gojq.Parse(s)
					if err != nil {
						p := queryErrorPosition(s, err)
						return nil, compileError{
							err:      err,
							what:     "parse",
							filename: filenamePart,
							pos:      p,
						}
					}
				}

				// TODO: some better way of handling relative includes that
				// works with @builtin etc
				basePath := filepath.Dir(name)
				for _, i := range q.Imports {
					rewritePath := func(base, path string) string {
						if strings.HasPrefix(i.IncludePath, "@") {
							return path
						}
						if filepath.IsAbs(i.IncludePath) {
							return path
						}
						return filepath.Join(base, path)
					}
					i.IncludePath = rewritePath(basePath, i.IncludePath)
					i.ImportPath = rewritePath(basePath, i.ImportPath)
				}

				i.includeCache[filename] = q

				return q, nil
			}

			panic("unreachable")
		},
	}))

	gc, err := gojq.Compile(gq, compilerOpts...)
	if err != nil {
		p := queryErrorPosition(src, err)
		return nil, compileError{
			err:      err,
			what:     "compile",
			filename: srcFilename,
			pos:      p,
		}
	}

	runCtx, runCtxCancelFn := i.interruptStack.Push(ctx)
	ni.evalContext.ctx = runCtx
	ni.evalContext.output = ioextra.CtxWriter{Writer: output, Ctx: runCtx}
	iter := gc.RunWithContext(runCtx, c, variableValues...)

	iterWrapper := iterFn(func() (interface{}, bool) {
		v, ok := iter.Next()
		// gojq ctx cancel will not return ok=false, just cancelled error
		if !ok {
			runCtxCancelFn()
		}
		return v, ok
	})

	return iterWrapper, nil
}

func (i *Interp) EvalFunc(ctx context.Context, c interface{}, name string, args []interface{}, output io.Writer) (gojq.Iter, error) {
	var argsExpr []string
	for i := range args {
		argsExpr = append(argsExpr, fmt.Sprintf("$_args[%d]", i))
	}
	argExpr := ""
	if len(argsExpr) > 0 {
		argExpr = "(" + strings.Join(argsExpr, ";") + ")"
	}

	trampolineInput := map[string]interface{}{
		"input": c,
		"args":  args,
	}
	// _args to mark variable as internal and hide it from completion
	// {input: ..., args: [...]} | .args as {args: $_args} | .input | name[($_args[0]; ...)]
	trampolineExpr := fmt.Sprintf(". as {args: $_args} | .input | %s%s", name, argExpr)
	iter, err := i.Eval(ctx, trampolineInput, trampolineExpr, "", output)
	if err != nil {
		return nil, err
	}
	return iter, nil
}

func (i *Interp) EvalFuncValues(ctx context.Context, c interface{}, name string, args []interface{}, output io.Writer) ([]interface{}, error) {
	iter, err := i.EvalFunc(ctx, c, name, args, output)
	if err != nil {
		return nil, err
	}

	var vs []interface{}
	for {
		v, ok := iter.Next()
		if !ok {
			break
		}
		vs = append(vs, v)
	}

	return vs, nil
}

type Options struct {
	Depth          int    `mapstructure:"depth"`
	ArrayTruncate  int    `mapstructure:"array_truncate"`
	Verbose        bool   `mapstructure:"verbose"`
	DecodeProgress bool   `mapstructure:"decode_progress"`
	Color          bool   `mapstructure:"color"`
	Colors         string `mapstructure:"colors"`
	ByteColors     string `mapstructure:"byte_colors"`
	Unicode        bool   `mapstructure:"unicode"`
	RawOutput      bool   `mapstructure:"raw_output"`
	REPL           bool   `mapstructure:"repl"`
	RawString      bool   `mapstructure:"raw_string"`
	JoinString     string `mapstructure:"join_string"`
	Compact        bool   `mapstructure:"compact"`
	BitsFormat     string `mapstructure:"bits_format"`
	LineBytes      int    `mapstructure:"line_bytes"`
	DisplayBytes   int    `mapstructure:"display_bytes"`
	AddrBase       int    `mapstructure:"addrbase"`
	SizeBase       int    `mapstructure:"sizebase"`

	Decorator    Decorator
	BitsFormatFn func(bb *bitio.Buffer) (interface{}, error)
}

func bitsFormatFnFromOptions(opts Options) func(bb *bitio.Buffer) (interface{}, error) {
	switch opts.BitsFormat {
	case "md5":
		return func(bb *bitio.Buffer) (interface{}, error) {
			d := md5.New()
			if _, err := io.Copy(d, bb); err != nil {
				return "", err
			}
			return hex.EncodeToString(d.Sum(nil)), nil
		}
	case "base64":
		return func(bb *bitio.Buffer) (interface{}, error) {
			b := &bytes.Buffer{}
			e := base64.NewEncoder(base64.StdEncoding, b)
			if _, err := io.Copy(e, bb); err != nil {
				return "", err
			}
			e.Close()
			return b.String(), nil
		}
	case "truncate":
		// TODO: configure
		return func(bb *bitio.Buffer) (interface{}, error) {
			b := &bytes.Buffer{}
			if _, err := io.Copy(b, io.LimitReader(bb, 1024)); err != nil {
				return "", err
			}
			return b.String(), nil
		}
	case "string":
		return func(bb *bitio.Buffer) (interface{}, error) {
			b := &bytes.Buffer{}
			if _, err := io.Copy(b, bb); err != nil {
				return "", err
			}
			return b.String(), nil
		}
	case "snippet":
		fallthrough
	default:
		return func(bb *bitio.Buffer) (interface{}, error) {
			b := &bytes.Buffer{}
			e := base64.NewEncoder(base64.StdEncoding, b)
			if _, err := io.Copy(e, io.LimitReader(bb, 256)); err != nil {
				return "", err
			}
			e.Close()
			return fmt.Sprintf("<%s>%s", num.Bits(bb.Len()).StringByteBits(opts.SizeBase), b.String()), nil
		}
	}
}

func (i *Interp) lookupState(key string) interface{} {
	if i.state == nil {
		return nil
	}
	m, ok := (*i.state).(map[string]interface{})
	if !ok {
		return nil
	}
	return m[key]
}

func (i *Interp) includePaths() []string {
	pathsAny, _ := i.lookupState("include_paths").([]interface{})
	var paths []string
	for _, pathAny := range pathsAny {
		paths = append(paths, pathAny.(string))
	}
	return paths
}

func (i *Interp) variables() map[string]interface{} {
	variablesAny, _ := i.lookupState("variables").(map[string]interface{})
	return variablesAny
}

func (i *Interp) Options(v interface{}) Options {
	var opts Options
	_ = mapstructure.Decode(v, &opts)
	opts.ArrayTruncate = num.MaxInt(0, opts.ArrayTruncate)
	opts.Depth = num.MaxInt(0, opts.Depth)
	opts.AddrBase = num.ClampInt(2, 36, opts.AddrBase)
	opts.SizeBase = num.ClampInt(2, 36, opts.SizeBase)
	opts.LineBytes = num.MaxInt(0, opts.LineBytes)
	opts.DisplayBytes = num.MaxInt(0, opts.DisplayBytes)
	opts.Decorator = decoratorFromOptions(opts)
	opts.BitsFormatFn = bitsFormatFnFromOptions(opts)

	return opts
}

func (i *Interp) NewColorJSON(opts Options) (*colorjson.Encoder, error) {
	indent := 2
	if opts.Compact {
		indent = 0
	}

	return colorjson.NewEncoder(
		opts.Color,
		false,
		indent,
		func(v interface{}) interface{} {
			if v, ok := toValue(func() Options { return opts }, v); ok {
				return v
			}
			panic(fmt.Sprintf("toValue not a JQValue value: %#v", v))
		},
		colorjson.Colors{
			Reset:     []byte(ansi.Reset.SetString),
			Null:      []byte(opts.Decorator.Null.SetString),
			False:     []byte(opts.Decorator.False.SetString),
			True:      []byte(opts.Decorator.True.SetString),
			Number:    []byte(opts.Decorator.Number.SetString),
			String:    []byte(opts.Decorator.String.SetString),
			ObjectKey: []byte(opts.Decorator.ObjectKey.SetString),
			Array:     []byte(opts.Decorator.Array.SetString),
			Object:    []byte(opts.Decorator.Object.SetString),
		},
	), nil
}