package proto

import (
	"fmt"

	"github.com/segmentio/encoding/json"
)

// Rewriter is an interface implemented by types that support rewriting protobuf
// messages.
type Rewriter interface {
	// The function is expected to append the new content to the byte slice
	// passed as argument. If it wasn't able to perform the rewrite, it must
	// return a non-nil error.
	Rewrite(out, in []byte) ([]byte, error)
}

type identity struct{}

func (identity) Rewrite(out, in []byte) ([]byte, error) {
	return append(out, in...), nil
}

// MultiRewriter constructs a Rewriter which applies all rewriters passed as
// arguments.
func MultiRewriter(rewriters ...Rewriter) Rewriter {
	if len(rewriters) == 1 {
		return rewriters[0]
	}
	m := &multiRewriter{rewriters: make([]Rewriter, len(rewriters))}
	copy(m.rewriters, rewriters)
	return m
}

type multiRewriter struct {
	rewriters []Rewriter
}

func (m *multiRewriter) Rewrite(out, in []byte) ([]byte, error) {
	var err error

	for _, rw := range m.rewriters {
		if out, err = rw.Rewrite(out, in); err != nil {
			return out, err
		}
	}

	return out, nil
}

// RewriteFunc is a function type implementing the Rewriter interface.
type RewriteFunc func([]byte, []byte) ([]byte, error)

// Rewrite satisfies the Rewriter interface.
func (r RewriteFunc) Rewrite(out, in []byte) ([]byte, error) {
	return r(out, in)
}

// MessageRewriter maps field numbers to rewrite rules, satisfying the Rewriter
// interace to support composing rewrite rules.
type MessageRewriter []Rewriter

// Rewrite applies the rewrite rule matching f in r, satisfies the Rewriter
// interface.
func (r MessageRewriter) Rewrite(out, in []byte) ([]byte, error) {
	seen := make(fieldset, 4)

	if n := seen.len(); len(r) >= n {
		seen = makeFieldset(len(r) + 1)
	}

	for len(in) != 0 {
		f, t, v, m, err := Parse(in)
		if err != nil {
			return out, err
		}

		if i := int(f); i >= 0 && i < len(r) && r[i] != nil {
			if !seen.has(i) {
				seen.set(i)
				if out, err = r[i].Rewrite(out, v); err != nil {
					return out, err
				}
			}
		} else {
			out = Append(out, f, t, v)
		}

		in = m
	}

	for i, f := range r {
		if f != nil && !seen.has(i) {
			b, err := r[i].Rewrite(out, nil)
			if err != nil {
				return b, err
			}
			out = b
		}
	}

	return out, nil
}

type fieldset []uint64

func makeFieldset(n int) fieldset {
	if (n % 64) != 0 {
		n = (n + 1) / 64
	} else {
		n /= 64
	}
	return make(fieldset, n)
}

func (f fieldset) len() int {
	return len(f) * 64
}

func (f fieldset) has(i int) bool {
	x, y := f.index(i)
	return ((f[x] >> y) & 1) != 0
}

func (f fieldset) set(i int) {
	x, y := f.index(i)
	f[x] |= 1 << y
}

func (f fieldset) unset(i int) {
	x, y := f.index(i)
	f[x] &= ^(1 << y)
}

func (f fieldset) index(i int) (int, int) {
	return i / 64, i % 64
}

// ParseRewriteTemplate constructs a Rewriter for a protobuf type using the
// given json template to describe the rewrite rules.
//
// The json template contains a representation of the message that is used as the
// source values to overwrite in the protobuf targeted by the resulting rewriter.
//
// The rules are an optional set of RewriterRules that can provide alternative
// Rewriters from the default used for the field type. These rules are given the
// json.RawMessage bytes from the template, and they are expected to create a
// Rewriter to be applied against the target protobuf.
func ParseRewriteTemplate(typ Type, jsonTemplate []byte, rules ...RewriterRules) (Rewriter, error) {
	switch typ.Kind() {
	case Struct:
		return parseRewriteTemplateStruct(typ, 0, jsonTemplate, rules...)
	default:
		return nil, fmt.Errorf("cannot construct a rewrite template from a non-struct type %s", typ.Name())
	}
}

func parseRewriteTemplate(t Type, f FieldNumber, j json.RawMessage, rule any) (Rewriter, error) {
	if rwer, ok := rule.(Rewriterer); ok {
		return rwer.Rewriter(t, f, j)
	}

	switch t.Kind() {
	case Bool:
		return parseRewriteTemplateBool(t, f, j)
	case Int32:
		return parseRewriteTemplateInt32(t, f, j)
	case Int64:
		return parseRewriteTemplateInt64(t, f, j)
	case Sint32:
		return parseRewriteTemplateSint32(t, f, j)
	case Sint64:
		return parseRewriteTemplateSint64(t, f, j)
	case Uint32:
		return parseRewriteTemplateUint64(t, f, j)
	case Uint64:
		return parseRewriteTemplateUint64(t, f, j)
	case Fix32:
		return parseRewriteTemplateFix32(t, f, j)
	case Fix64:
		return parseRewriteTemplateFix64(t, f, j)
	case Sfix32:
		return parseRewriteTemplateSfix32(t, f, j)
	case Sfix64:
		return parseRewriteTemplateSfix64(t, f, j)
	case Float:
		return parseRewriteTemplateFloat(t, f, j)
	case Double:
		return parseRewriteTemplateDouble(t, f, j)
	case String:
		return parseRewriteTemplateString(t, f, j)
	case Bytes:
		return parseRewriteTemplateBytes(t, f, j)
	case Map:
		return parseRewriteTemplateMap(t, f, j)
	case Struct:
		sub, n, ok := [1]RewriterRules{}, 0, false
		if sub[0], ok = rule.(RewriterRules); ok {
			n = 1
		}
		return parseRewriteTemplateStruct(t, f, j, sub[:n]...)
	default:
		return nil, fmt.Errorf("cannot construct a rewriter from type %s", t.Name())
	}
}

func parseRewriteTemplateBool(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	var v bool
	err := json.Unmarshal(j, &v)
	if !v || err != nil {
		return nil, err
	}
	return f.Bool(v), nil
}

func parseRewriteTemplateInt32(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	var v int32
	err := json.Unmarshal(j, &v)
	if v == 0 || err != nil {
		return nil, err
	}
	return f.Int32(v), nil
}

func parseRewriteTemplateInt64(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	var v int64
	err := json.Unmarshal(j, &v)
	if v == 0 || err != nil {
		return nil, err
	}
	return f.Int64(v), nil
}

func parseRewriteTemplateSint32(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	var v int32
	err := json.Unmarshal(j, &v)
	if v == 0 || err != nil {
		return nil, err
	}
	return f.Uint32(encodeZigZag32(v)), nil
}

func parseRewriteTemplateSint64(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	var v int64
	err := json.Unmarshal(j, &v)
	if v == 0 || err != nil {
		return nil, err
	}
	return f.Uint64(encodeZigZag64(v)), nil
}

func parseRewriteTemplateUint32(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	var v uint32
	err := json.Unmarshal(j, &v)
	if v == 0 || err != nil {
		return nil, err
	}
	return f.Uint32(v), nil
}

func parseRewriteTemplateUint64(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	var v uint64
	err := json.Unmarshal(j, &v)
	if v == 0 || err != nil {
		return nil, err
	}
	return f.Uint64(v), nil
}

func parseRewriteTemplateFix32(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	var v uint32
	err := json.Unmarshal(j, &v)
	if v == 0 || err != nil {
		return nil, err
	}
	return f.Fixed32(v), nil
}

func parseRewriteTemplateFix64(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	var v uint64
	err := json.Unmarshal(j, &v)
	if v == 0 || err != nil {
		return nil, err
	}
	return f.Fixed64(v), nil
}

func parseRewriteTemplateSfix32(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	var v int32
	err := json.Unmarshal(j, &v)
	if v == 0 || err != nil {
		return nil, err
	}
	return f.Fixed32(encodeZigZag32(v)), nil
}

func parseRewriteTemplateSfix64(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	var v int64
	err := json.Unmarshal(j, &v)
	if v == 0 || err != nil {
		return nil, err
	}
	return f.Fixed64(encodeZigZag64(v)), nil
}

func parseRewriteTemplateFloat(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	var v float32
	err := json.Unmarshal(j, &v)
	if v == 0 || err != nil {
		return nil, err
	}
	return f.Float32(v), nil
}

func parseRewriteTemplateDouble(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	var v float64
	err := json.Unmarshal(j, &v)
	if v == 0 || err != nil {
		return nil, err
	}
	return f.Float64(v), nil
}

func parseRewriteTemplateString(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	var v string
	err := json.Unmarshal(j, &v)
	if v == "" || err != nil {
		return nil, err
	}
	return f.String(v), nil
}

func parseRewriteTemplateBytes(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	var v string
	err := json.Unmarshal(j, &v)
	if v == "" || err != nil {
		return nil, err
	}
	return f.Bytes([]byte(v)), nil
}

func parseRewriteTemplateMap(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	st := &structType{
		name: t.Name(),
		fields: []Field{
			{Index: 0, Number: 1, Name: "key", Type: t.Key()},
			{Index: 1, Number: 2, Name: "value", Type: t.Elem()},
		},
		fieldsByName:   make(map[string]int),
		fieldsByNumber: make(map[FieldNumber]int),
	}

	for _, f := range st.fields {
		st.fieldsByName[f.Name] = f.Index
		st.fieldsByNumber[f.Number] = f.Index
	}

	template := map[string]json.RawMessage{}

	if err := json.Unmarshal(j, &template); err != nil {
		return nil, err
	}

	maplist := make([]json.RawMessage, 0, len(template))

	for key, value := range template {
		b, err := json.Marshal(struct {
			Key   string          `json:"key"`
			Value json.RawMessage `json:"value"`
		}{
			Key:   key,
			Value: value,
		})
		if err != nil {
			return nil, err
		}
		maplist = append(maplist, b)
	}

	rewriters := make([]Rewriter, len(maplist))

	for i, b := range maplist {
		r, err := parseRewriteTemplateStruct(st, f, b)
		if err != nil {
			return nil, err
		}
		rewriters[i] = r
	}

	return MultiRewriter(rewriters...), nil
}

func parseRewriteTemplateStruct(t Type, f FieldNumber, j json.RawMessage, rules ...RewriterRules) (Rewriter, error) {
	template := map[string]json.RawMessage{}

	if err := json.Unmarshal(j, &template); err != nil {
		return nil, err
	}

	fieldsByName := map[string]Field{}

	for i := range t.NumField() {
		f := t.Field(i)
		fieldsByName[f.Name] = f
	}

	message := MessageRewriter{}
	rewriters := []Rewriter{}

	for k, v := range template {
		f, ok := fieldsByName[k]
		if !ok {
			return nil, fmt.Errorf("rewrite template contained an invalid field named %q", k)
		}

		var fields []json.RawMessage
		if f.Repeated {
			if err := json.Unmarshal(v, &fields); err != nil {
				return nil, err
			}
		} else {
			fields = []json.RawMessage{v}
		}

		var rule any
		for i := range rules {
			if r, ok := rules[i][f.Name]; ok {
				rule = r
				break
			}
		}

		rewriters = rewriters[:0]

		for _, v := range fields {
			rw, err := parseRewriteTemplate(f.Type, f.Number, v, rule)
			if err != nil {
				return nil, fmt.Errorf("%s: %w", k, err)
			}
			if rw != nil {
				rewriters = append(rewriters, rw)
			}
		}

		if cap(message) <= int(f.Number) {
			m := make(MessageRewriter, f.Number+1)
			copy(m, message)
			message = m
		}

		message[f.Number] = MultiRewriter(rewriters...)
	}

	if f != 0 {
		return &embddedRewriter{number: f, message: message}, nil
	}

	return message, nil
}

type embddedRewriter struct {
	number  FieldNumber
	message MessageRewriter
}

func (f *embddedRewriter) Rewrite(out, in []byte) ([]byte, error) {
	prefix := len(out)

	out, err := f.message.Rewrite(out, in)
	if err != nil {
		return nil, err
	}
	if len(out) == prefix {
		return out, nil
	}

	b := [24]byte{}
	n1, _ := encodeVarint(b[:], EncodeTag(f.number, Varlen))
	n2, _ := encodeVarint(b[n1:], uint64(len(out)-prefix))
	tagAndLen := n1 + n2

	out = append(out, b[:tagAndLen]...)
	copy(out[prefix+tagAndLen:], out[prefix:])
	copy(out[prefix:], b[:tagAndLen])
	return out, nil
}

// RewriterRules defines a set of rules for overriding the Rewriter used for any
// particular field. These maps may be nested for defining rules for struct members.
//
// For example:
//
//	rules := proto.RewriterRules {
//		"flags": proto.BitOr[uint64]{},
//		"nested": proto.RewriterRules {
//			"name": myCustomRewriter,
//		},
//	}
type RewriterRules map[string]any

// Rewriterer is the interface for producing a Rewriter for a given Type, FieldNumber
// and json.RawMessage. The JSON value is the JSON-encoded payload that should be
// decoded to produce the appropriate Rewriter. Implementations of the Rewriterer
// interface are added to the RewriterRules to specify the rules for performing
// custom rewrite logic.
type Rewriterer interface {
	Rewriter(Type, FieldNumber, json.RawMessage) (Rewriter, error)
}

// BitOr implments the Rewriterer interface for providing a bitwise-or rewrite
// logic for integers rather than replacing them. Instances of this type are
// zero-size, carrying only the generic type for creating the appropriate
// Rewriter when requested.
//
// Adding these to a RewriterRules looks like:
//
//	rules := proto.RewriterRules {
//		"flags": proto.BitOr[uint64]{},
//	}
//
// When used as a rule when rewriting from a template, the BitOr expects a JSON-
// encoded integer passed into the Rewriter method. This parsed integer is then
// used to perform a bitwise-or against the protobuf message that is being rewritten.
//
// The above example can then be used like:
//
//	template := []byte(`{"flags": 8}`) // n |= 0b1000
//	rw, err := proto.ParseRewriteTemplate(typ, template, rules)
type BitOr[T integer] struct{}

// integer is the contraint used by the BitOr Rewriterer and the bitOrRW Rewriter.
// Because these perform bitwise-or operations, the types must be integer-like.
type integer interface {
	~int | ~int32 | ~int64 | ~uint | ~uint32 | ~uint64
}

// Rewriter implements the Rewriterer interface. The JSON value provided to this
// method comes from the template used for rewriting. The returned Rewriter will use
// this JSON-encoded integer to perform a bitwise-or against the protobuf message
// that is being rewritten.
func (BitOr[T]) Rewriter(t Type, f FieldNumber, j json.RawMessage) (Rewriter, error) {
	var v T
	err := json.Unmarshal(j, &v)
	if err != nil {
		return nil, err
	}
	return BitOrRewriter(t, f, v)
}

// BitOrRewriter creates a bitwise-or Rewriter for a given field type and number.
// The mask is the value or'ed with values in the target protobuf.
func BitOrRewriter[T integer](t Type, f FieldNumber, mask T) (Rewriter, error) {
	switch t.Kind() {
	case Int32, Int64, Sint32, Sint64, Uint32, Uint64, Fix32, Fix64, Sfix32, Sfix64:
	default:
		return nil, fmt.Errorf("cannot construct a rewriter from type %s", t.Name())
	}
	return bitOrRW[T]{mask: mask, t: t, f: f}, nil
}

// bitOrRW is the Rewriter returned by the BitOr Rewriter method.
type bitOrRW[T integer] struct {
	mask T
	t    Type
	f    FieldNumber
}

// Rewrite implements the Rewriter interface performing a bitwise-or between the
// template value and the input value.
func (r bitOrRW[T]) Rewrite(out, in []byte) ([]byte, error) {
	var v T
	if err := Unmarshal(in, &v); err != nil {
		return nil, err
	}

	v |= r.mask

	switch r.t.Kind() {
	case Int32:
		return r.f.Int32(int32(v)).Rewrite(out, in)
	case Int64:
		return r.f.Int64(int64(v)).Rewrite(out, in)
	case Sint32:
		return r.f.Uint32(encodeZigZag32(int32(v))).Rewrite(out, in)
	case Sint64:
		return r.f.Uint64(encodeZigZag64(int64(v))).Rewrite(out, in)
	case Uint32, Uint64:
		return r.f.Uint64(uint64(v)).Rewrite(out, in)
	case Fix32:
		return r.f.Fixed32(uint32(v)).Rewrite(out, in)
	case Fix64:
		return r.f.Fixed64(uint64(v)).Rewrite(out, in)
	case Sfix32:
		return r.f.Fixed32(encodeZigZag32(int32(v))).Rewrite(out, in)
	case Sfix64:
		return r.f.Fixed64(encodeZigZag64(int64(v))).Rewrite(out, in)
	}

	panic("unreachable") // Kind is validated when creating instances
}