package btcwire

import (
	"encoding/binary"
	"fmt"
	"io"

	"git.smesh.lol/nostrino/pkg/crypto/chainhash"
)

// Compact size prefixes, as defined by Bitcoin's variable length integer
// encoding.
const (
	compactSizeUint16 = 0xfd
	compactSizeUint32 = 0xfe
	compactSizeUint64 = 0xff
)

// Encoder appends primitive Bitcoin wire encodings to an internal buffer. The
// zero value is ready for use. Encoding never fails except through the
// functions that return an error, such as WriteVarBytes.
type Encoder struct {
	buf []byte
}

// Uint8 appends v as a single byte.
func (e *Encoder) Uint8(v uint8) {
	e.buf = append(e.buf, v)
}

// Uint16 appends v as a little-endian uint16.
func (e *Encoder) Uint16(v uint16) {
	e.buf = append(e.buf, byte(v), byte(v>>8))
}

// Uint16BE appends v as a big-endian uint16. Network ports use this encoding
// in the version message.
func (e *Encoder) Uint16BE(v uint16) {
	e.buf = append(e.buf, byte(v>>8), byte(v))
}

// Uint32 appends v as a little-endian uint32.
func (e *Encoder) Uint32(v uint32) {
	e.buf = append(e.buf, byte(v), byte(v>>8), byte(v>>16), byte(v>>24))
}

// Uint64 appends v as a little-endian uint64.
func (e *Encoder) Uint64(v uint64) {
	e.buf = append(e.buf,
		byte(v), byte(v>>8), byte(v>>16), byte(v>>24),
		byte(v>>32), byte(v>>40), byte(v>>48), byte(v>>56),
	)
}

// Int32 appends v as a little-endian signed 32-bit integer.
func (e *Encoder) Int32(v int32) {
	e.Uint32(uint32(v))
}

// Int64 appends v as a little-endian signed 64-bit integer.
func (e *Encoder) Int64(v int64) {
	e.Uint64(uint64(v))
}

// CompactSize appends v in Bitcoin's compact size encoding, using the
// shortest representation.
func (e *Encoder) CompactSize(v uint64) {
	switch {
	case v < compactSizeUint16:
		e.Uint8(uint8(v))
	case v <= 0xffff:
		e.Uint8(compactSizeUint16)
		e.Uint16(uint16(v))
	case v <= 0xffffffff:
		e.Uint8(compactSizeUint32)
		e.Uint32(uint32(v))
	default:
		e.Uint8(compactSizeUint64)
		e.Uint64(v)
	}
}

// Hash appends h in wire byte order.
func (e *Encoder) Hash(h chainhash.Hash) {
	e.buf = append(e.buf, h[:]...)
}

// Write appends b verbatim.
func (e *Encoder) Write(b []byte) {
	e.buf = append(e.buf, b...)
}

// Bytes returns the encoded bytes. The result aliases the encoder's storage
// and is only valid until the next append.
func (e *Encoder) Bytes() []byte {
	return e.buf
}

// Len returns the number of bytes encoded so far.
func (e *Encoder) Len() int {
	return len(e.buf)
}

// Decoder decodes primitive Bitcoin wire types from a byte slice. All methods
// return an error rather than panicking when the input is truncated.
type Decoder struct {
	buf []byte
	off int
}

// NewDecoder returns a Decoder over buf. The Decoder does not copy buf, so
// the caller must not modify it while decoding.
func NewDecoder(buf []byte) *Decoder {
	return &Decoder{buf: buf}
}

// Remaining returns the number of undecoded bytes.
func (d *Decoder) Remaining() int {
	return len(d.buf) - d.off
}

// Offset returns the number of bytes decoded so far.
func (d *Decoder) Offset() int {
	return d.off
}

// take returns the next n bytes without copying them.
func (d *Decoder) take(n int) ([]byte, error) {
	if n < 0 {
		return nil, fmt.Errorf("%w: cannot read %d bytes", ErrInvalidLength, n)
	}
	if n > d.Remaining() {
		return nil, fmt.Errorf(
			"btcwire: reading %d bytes at offset %d, only %d available: %w",
			n, d.off, d.Remaining(), io.ErrUnexpectedEOF,
		)
	}
	b := d.buf[d.off : d.off+n]
	d.off += n
	return b, nil
}

// Bytes returns the next n bytes. The result aliases the decoder's input.
func (d *Decoder) Bytes(n int) ([]byte, error) {
	return d.take(n)
}

// Uint8 decodes a single byte.
func (d *Decoder) Uint8() (uint8, error) {
	b, err := d.take(1)
	if err != nil {
		return 0, err
	}
	return b[0], nil
}

// Uint16 decodes a little-endian uint16.
func (d *Decoder) Uint16() (uint16, error) {
	b, err := d.take(2)
	if err != nil {
		return 0, err
	}
	return binary.LittleEndian.Uint16(b), nil
}

// Uint16BE decodes a big-endian uint16.
func (d *Decoder) Uint16BE() (uint16, error) {
	b, err := d.take(2)
	if err != nil {
		return 0, err
	}
	return binary.BigEndian.Uint16(b), nil
}

// Uint32 decodes a little-endian uint32.
func (d *Decoder) Uint32() (uint32, error) {
	b, err := d.take(4)
	if err != nil {
		return 0, err
	}
	return binary.LittleEndian.Uint32(b), nil
}

// Uint64 decodes a little-endian uint64.
func (d *Decoder) Uint64() (uint64, error) {
	b, err := d.take(8)
	if err != nil {
		return 0, err
	}
	return binary.LittleEndian.Uint64(b), nil
}

// Int32 decodes a little-endian signed 32-bit integer.
func (d *Decoder) Int32() (int32, error) {
	v, err := d.Uint32()
	if err != nil {
		return 0, err
	}
	return int32(v), nil
}

// Int64 decodes a little-endian signed 64-bit integer.
func (d *Decoder) Int64() (int64, error) {
	v, err := d.Uint64()
	if err != nil {
		return 0, err
	}
	return int64(v), nil
}

// Hash decodes a chainhash.Hash in wire byte order.
func (d *Decoder) Hash() (chainhash.Hash, error) {
	var h chainhash.Hash
	b, err := d.take(chainhash.HashSize)
	if err != nil {
		return h, err
	}
	copy(h[:], b)
	return h, nil
}

// CompactSize decodes a compact size integer, rejecting non-minimal
// encodings.
func (d *Decoder) CompactSize() (uint64, error) {
	prefix, err := d.Uint8()
	if err != nil {
		return 0, err
	}
	switch prefix {
	case compactSizeUint16:
		v, err := d.Uint16()
		if err != nil {
			return 0, err
		}
		if v < compactSizeUint16 {
			return 0, fmt.Errorf(
				"btcwire: compact size %d encoded with 0x%02x prefix: %w",
				v, compactSizeUint16, ErrNonCanonicalCompactSize,
			)
		}
		return uint64(v), nil
	case compactSizeUint32:
		v, err := d.Uint32()
		if err != nil {
			return 0, err
		}
		if v <= 0xffff {
			return 0, fmt.Errorf(
				"btcwire: compact size %d encoded with 0x%02x prefix: %w",
				v, compactSizeUint32, ErrNonCanonicalCompactSize,
			)
		}
		return uint64(v), nil
	case compactSizeUint64:
		v, err := d.Uint64()
		if err != nil {
			return 0, err
		}
		if v <= 0xffffffff {
			return 0, fmt.Errorf(
				"btcwire: compact size %d encoded with 0x%02x prefix: %w",
				v, compactSizeUint64, ErrNonCanonicalCompactSize,
			)
		}
		return v, nil
	default:
		return uint64(prefix), nil
	}
}

// maxCompactSizeCount is the largest element count that can be safely
// converted to an int on every supported platform.
const maxCompactSizeCount = 1 << 31

// count decodes a compact size used as an element count and verifies that it
// does not exceed limit. It also rejects counts that cannot fit in memory
// given the remaining input, where minItemSize is the smallest encoded size
// of one element.
func (d *Decoder) count(limit int, minItemSize int, what string) (int, error) {
	n, err := d.CompactSize()
	if err != nil {
		return 0, fmt.Errorf("btcwire: reading %s count: %w", what, err)
	}
	if n > maxCompactSizeCount || n > uint64(limit) {
		return 0, fmt.Errorf(
			"btcwire: %s count %d exceeds limit %d: %w",
			what, n, limit, ErrTooManyItems,
		)
	}
	if minItemSize > 0 && n > 0 {
		need := n * uint64(minItemSize)
		if need > uint64(d.Remaining()) {
			return 0, fmt.Errorf(
				"btcwire: %s count %d needs at least %d bytes, only %d available: %w",
				what, n, need, d.Remaining(), io.ErrUnexpectedEOF,
			)
		}
	}
	return int(n), nil
}

// WriteVarBytes appends b to e prefixed by its compact size length. It fails
// when b is longer than MaxPayload, since such a field cannot be carried by a
// valid frame.
func WriteVarBytes(e *Encoder, b []byte) error {
	if len(b) > MaxPayload {
		return fmt.Errorf(
			"btcwire: variable length field of %d bytes exceeds maximum %d: %w",
			len(b), MaxPayload, ErrVarBytesTooLarge,
		)
	}
	e.CompactSize(uint64(len(b)))
	e.Write(b)
	return nil
}

// ReadVarBytes decodes a variable length byte string whose length must not
// exceed maxLen. maxLen must be between zero and MaxPayload. The returned
// slice aliases the decoder's input.
func ReadVarBytes(d *Decoder, maxLen int) ([]byte, error) {
	if maxLen < 0 || maxLen > MaxPayload {
		return nil, fmt.Errorf("%w: maximum %d", ErrInvalidLength, maxLen)
	}
	n, err := d.CompactSize()
	if err != nil {
		return nil, fmt.Errorf("btcwire: reading varbytes length: %w", err)
	}
	if n > uint64(maxLen) {
		return nil, fmt.Errorf(
			"btcwire: varbytes length %d exceeds maximum %d: %w",
			n, maxLen, ErrVarBytesTooLarge,
		)
	}
	b, err := d.Bytes(int(n))
	if err != nil {
		return nil, fmt.Errorf("btcwire: reading %d byte varbytes: %w", n, err)
	}
	return b, nil
}