package maintenance

import (
	"bytes"
	"crypto/ecdh"
	"encoding/binary"
	"encoding/json"
	"errors"
	"strings"
	"unicode/utf8"
)

const (
	MaxInputBytes    = 64 << 10
	MaxArguments     = 32
	MaxArgumentBytes = 8 << 10
	inputPrefix      = "CMI1"
	maxInputFrame    = len(inputPrefix) + 4 + 4*MaxArguments + MaxArgumentBytes + 4 + MaxInputBytes
)

var ErrInput = errors.New("maintenance: invalid private input")

// Input contains arguments for the pinned artifact and its exact JSON stdin.
// The worker supplies the existing database path; this value grants no authority
// to select an executable, shell, database, environment, or working directory.
type Input struct {
	Arguments []string
	Stdin     []byte
}

func validInput(in Input) bool {
	if len(in.Arguments) == 0 || len(in.Arguments) > MaxArguments || len(in.Stdin) > MaxInputBytes || !utf8.Valid(in.Stdin) || !json.Valid(in.Stdin) {
		return false
	}
	total := 0
	for _, arg := range in.Arguments {
		if len(arg) > MaxArgumentBytes-total || strings.ContainsRune(arg, 0) || !utf8.ValidString(arg) {
			return false
		}
		total += len(arg)
	}
	return true
}

// EncodeInput retains whitespace, escapes, and every other byte of signed input.
// JSON semantics, including allowed fields, remain the app command's concern.
func EncodeInput(in Input) ([]byte, error) {
	if !validInput(in) {
		return nil, ErrInput
	}
	out := binary.BigEndian.AppendUint32([]byte(inputPrefix), uint32(len(in.Arguments)))
	for _, arg := range in.Arguments {
		out = appendField(out, []byte(arg))
	}
	return appendField(out, in.Stdin), nil
}

// DecodeInput returns owned data only after checking the complete bounded frame.
func DecodeInput(frame []byte) (Input, error) {
	if len(frame) > maxInputFrame || len(frame) < 8 || !bytes.Equal(frame[:4], []byte(inputPrefix)) {
		return Input{}, ErrInput
	}
	count := binary.BigEndian.Uint32(frame[4:8])
	if count == 0 || count > MaxArguments {
		return Input{}, ErrInput
	}
	rest := frame[8:]
	take := func() ([]byte, bool) {
		if len(rest) < 4 {
			return nil, false
		}
		n := binary.BigEndian.Uint32(rest[:4])
		rest = rest[4:]
		if uint64(n) > uint64(len(rest)) {
			return nil, false
		}
		value := rest[:int(n)]
		rest = rest[int(n):]
		return value, true
	}
	in := Input{Arguments: make([]string, 0, int(count))}
	for range count {
		arg, ok := take()
		if !ok {
			return Input{}, ErrInput
		}
		in.Arguments = append(in.Arguments, string(arg))
	}
	stdin, ok := take()
	if !ok || len(rest) != 0 {
		return Input{}, ErrInput
	}
	in.Stdin = stdin
	if !validInput(in) {
		return Input{}, ErrInput
	}
	in.Stdin = bytes.Clone(stdin)
	return in, nil
}

func SealInput(pub *ecdh.PublicKey, binding Binding, in Input) (string, error) {
	if binding.Direction != DirectionInput {
		return "", ErrInput
	}
	frame, err := EncodeInput(in)
	if err != nil {
		return "", ErrInput
	}
	defer clear(frame)
	sealed, err := Seal(pub, binding, frame)
	if err != nil {
		return "", ErrInput
	}
	return sealed, nil
}

func OpenInput(priv *ecdh.PrivateKey, binding Binding, sealed string) (Input, error) {
	if binding.Direction != DirectionInput {
		return Input{}, ErrInput
	}
	frame, err := Open(priv, binding, sealed)
	if err != nil {
		return Input{}, ErrInput
	}
	defer clear(frame)
	return DecodeInput(frame)
}