package btcwire

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

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

// Network magic values used to identify the network a frame belongs to.
//
// The values are the conventional four-byte strings in the order they appear
// on the wire, so MainNet is f9 be b4 d9 and RegTest is fa bf b5 da. They are
// therefore written and read big endian: the digits of the constant are the
// bytes on the wire, in the order every Bitcoin document prints them.
const (
	// MainNet is the magic for the Bitcoin main network.
	MainNet uint32 = 0xf9beb4d9

	// TestNet3 is the magic for the Bitcoin test network, version 3.
	TestNet3 uint32 = 0x0b110907

	// RegTest is the magic for a Bitcoin regression test network.
	RegTest uint32 = 0xfabfb5da

	// SigNet is the magic for a Bitcoin signet.
	SigNet uint32 = 0x0a03cf40
)

// Frame layout constants.
const (
	// CommandSize is the fixed size of the command field in a frame header.
	CommandSize = 12

	// MessageHeaderSize is the size of a frame header in bytes.
	MessageHeaderSize = 24

	// ChecksumSize is the number of payload SHA256d bytes carried by a frame
	// header.
	ChecksumSize = 4

	// MaxPayload is the largest payload this package will encode or accept.
	MaxPayload = 32 << 20
)

// Command strings used by the messages this package implements. A command is
// at most CommandSize bytes and consists of printable ASCII without NUL.
const (
	CmdVersion      = "version"
	CmdVerAck       = "verack"
	CmdPing         = "ping"
	CmdPong         = "pong"
	CmdGetHeaders   = "getheaders"
	CmdHeaders      = "headers"
	CmdGetCFilters  = "getcfilters"
	CmdCFilter      = "cfilter"
	CmdGetCFHeaders = "getcfheaders"
	CmdCFHeaders    = "cfheaders"
	CmdGetCFCheckpt = "getcfcheckpt"
	CmdCFCheckpt    = "cfcheckpt"
	CmdGetData      = "getdata"
	CmdTx           = "tx"
	CmdBlock        = "block"
	CmdInv          = "inv"
	CmdNotFound     = "notfound"
	CmdSendHeaders  = "sendheaders"
	CmdWtxidRelay   = "wtxidrelay"
	CmdSendAddrV2   = "sendaddrv2"
	CmdFeeFilter    = "feefilter"
	// CmdSendCmpct is how a peer announces compact-block support. A node sends it
	// unprompted right after the handshake even to a client that never asks for BIP152,
	// so a gateway that does not recognise it drops every session.
	CmdSendCmpct = "sendcmpct"
)

// Message is a decoded frame: a command name and its raw payload.
type Message struct {
	// Command is the command name without NUL padding.
	Command string

	// Payload is the raw message payload.
	Payload []byte
}

// Payload is implemented by the message payload types in this package, which
// convert themselves to and from the bytes carried inside a frame.
type Payload interface {
	// EncodePayload returns the wire encoding of the payload.
	EncodePayload() ([]byte, error)

	// DecodePayload decodes the payload from b.
	DecodePayload(b []byte) error
}

// Checksum returns the first four bytes of SHA256d(payload), the value carried
// by a frame header.
func Checksum(payload []byte) [ChecksumSize]byte {
	sum := chainhash.DoubleHashB(payload)
	var out [ChecksumSize]byte
	copy(out[:], sum[:ChecksumSize])
	return out
}

// encodeCommand validates command and returns it NUL padded to CommandSize
// bytes.
func encodeCommand(command string) ([CommandSize]byte, error) {
	var cmd [CommandSize]byte
	switch {
	case len(command) == 0:
		return cmd, fmt.Errorf("%w: empty command", ErrInvalidCommand)
	case len(command) > CommandSize:
		return cmd, fmt.Errorf(
			"%w: %q is %d bytes, maximum is %d",
			ErrInvalidCommand, command, len(command), CommandSize,
		)
	}
	for i := 0; i < len(command); i++ {
		c := command[i]
		if c < 0x20 || c > 0x7e {
			return cmd, fmt.Errorf(
				"%w: %q contains non-printable byte 0x%02x at index %d",
				ErrInvalidCommand, command, c, i,
			)
		}
		cmd[i] = c
	}
	return cmd, nil
}

// decodeCommand validates a NUL padded command field and returns it without
// the padding.
func decodeCommand(cmd []byte) (string, error) {
	if len(cmd) != CommandSize {
		return "", fmt.Errorf(
			"%w: command field is %d bytes, want %d",
			ErrInvalidCommand, len(cmd), CommandSize,
		)
	}
	n := CommandSize
	for n > 0 && cmd[n-1] == 0 {
		n--
	}
	if n == 0 {
		return "", fmt.Errorf("%w: command is all NUL padding", ErrInvalidCommand)
	}
	for i := 0; i < n; i++ {
		if cmd[i] < 0x20 || cmd[i] > 0x7e {
			return "", fmt.Errorf(
				"%w: command contains non-printable byte 0x%02x at index %d",
				ErrInvalidCommand, cmd[i], i,
			)
		}
	}
	return string(cmd[:n]), nil
}

// EncodeMessage returns a single complete frame for command and payload.
//
// The frame is the network magic, the NUL padded command, the payload length,
// the payload checksum, and the payload itself. The caller keeps ownership of
// payload; it is copied into the returned frame.
func EncodeMessage(magic uint32, command string, payload []byte) ([]byte, error) {
	cmd, err := encodeCommand(command)
	if err != nil {
		return nil, err
	}
	if len(payload) > MaxPayload {
		return nil, fmt.Errorf(
			"%w: %s payload is %d bytes, maximum is %d",
			ErrPayloadTooLarge, command, len(payload), MaxPayload,
		)
	}
	frame := make([]byte, MessageHeaderSize+len(payload))
	putMagic(frame[0:4], magic)
	copy(frame[4:4+CommandSize], cmd[:])
	binary.LittleEndian.PutUint32(frame[16:20], uint32(len(payload)))
	sum := Checksum(payload)
	copy(frame[20:24], sum[:])
	copy(frame[24:], payload)
	return frame, nil
}

// EncodeMessagePayload encodes p and returns a frame carrying it under
// command.
func EncodeMessagePayload(magic uint32, command string, p Payload) ([]byte, error) {
	payload, err := p.EncodePayload()
	if err != nil {
		return nil, fmt.Errorf("btcwire: encoding %s payload: %w", command, err)
	}
	return EncodeMessage(magic, command, payload)
}

// DecodeMessage reads and validates one frame from r. It validates the network
// magic, the command characters, the payload length bound, and the payload
// checksum.
//
// A clean end of input before any header byte is reported as io.EOF; a partial
// header or payload is reported as an error wrapping io.ErrUnexpectedEOF.
func DecodeMessage(magic uint32, r io.Reader) (*Message, error) {
	return decodeMessageLimit(magic, r, MaxPayload)
}

// decodeMessageLimit is DecodeMessage with a payload cap for this frame.
//
// The cap is applied before the payload is allocated, so a peer that answers a request for an
// 8 MiB block with a 30 MiB frame costs a header read and nothing else.
func decodeMessageLimit(magic uint32, r io.Reader, limit int) (*Message, error) {
	var header [MessageHeaderSize]byte
	if _, err := io.ReadFull(r, header[:]); err != nil {
		return nil, fmt.Errorf("btcwire: reading message header: %w", err)
	}
	gotMagic := readMagic(header[0:4])
	if gotMagic != magic {
		return nil, fmt.Errorf(
			"%w: frame has magic 0x%08x, want 0x%08x",
			ErrBadMagic, gotMagic, magic,
		)
	}
	command, err := decodeCommand(header[4 : 4+CommandSize])
	if err != nil {
		return nil, err
	}
	length := binary.LittleEndian.Uint32(header[16:20])
	if limit <= 0 || limit > MaxPayload {
		limit = MaxPayload
	}
	if length > uint32(limit) {
		return nil, fmt.Errorf(
			"%w: %s payload length %d, maximum is %d",
			ErrPayloadTooLarge, command, length, limit,
		)
	}
	payload := make([]byte, length)
	if _, err := io.ReadFull(r, payload); err != nil {
		// A zero-byte read at end of stream is reported as io.EOF by
		// io.ReadFull, but inside a frame it means truncation, not a clean
		// end of the message stream.
		if errors.Is(err, io.EOF) {
			err = io.ErrUnexpectedEOF
		}
		return nil, fmt.Errorf(
			"btcwire: reading %s payload of %d bytes: %w",
			command, length, err,
		)
	}
	want := Checksum(payload)
	if !bytes.Equal(header[20:24], want[:]) {
		return nil, fmt.Errorf(
			"%w: %s payload of %d bytes has checksum %x, want %x",
			ErrBadChecksum, command, length, header[20:24], want,
		)
	}
	return &Message{Command: command, Payload: payload}, nil
}

// Decode decodes the message payload into p.
func (m *Message) Decode(p Payload) error {
	if err := p.DecodePayload(m.Payload); err != nil {
		return fmt.Errorf("btcwire: decoding %s payload: %w", m.Command, err)
	}
	return nil
}

// putMagic writes the four network magic bytes into b, which must be four
// bytes long.
//
// The magic is the one field of the frame header that is big endian. Bitcoin
// prints it as the byte string that goes on the wire ("f9beb4d9" for mainnet),
// so the constant is that byte string read as a big-endian integer and the
// bytes come back out in the order they were printed.
func putMagic(b []byte, magic uint32) {
	binary.BigEndian.PutUint32(b, magic)
}

// readMagic reads the four network magic bytes in b.
func readMagic(b []byte) uint32 {
	return binary.BigEndian.Uint32(b)
}

// Reader reads a stream of frames that all belong to the same network.
type Reader struct {
	r     io.Reader
	magic uint32
	limit int
}

// SetLimit caps the payload of the next frames this reader decodes. Zero or a value above
// MaxPayload restores the package maximum.
func (r *Reader) SetLimit(max int) { r.limit = max }

// NewReader returns a Reader that reads frames carrying magic from r.
func NewReader(r io.Reader, magic uint32) *Reader {
	return &Reader{r: r, magic: magic}
}

// Next reads the next message. It returns io.EOF once the underlying reader is
// exhausted at a frame boundary.
func (r *Reader) Next() (*Message, error) {
	limit := r.limit
	if limit <= 0 || limit > MaxPayload {
		limit = MaxPayload
	}
	msg, err := decodeMessageLimit(r.magic, r.r, limit)
	if err != nil {
		if errors.Is(err, io.EOF) {
			return nil, io.EOF
		}
		return nil, err
	}
	return msg, nil
}

// Magic returns the network magic the Reader validates frames against.
func (r *Reader) Magic() uint32 {
	return r.magic
}