// Package p2p is a minimal Bitcoin P2P peer: the version handshake, frame
// reading and writing, ping/pong, and the message loop the compact-filter
// client drives. It knows nothing about wallets or chain state.
package p2p

import (
	"context"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"sync"
	"time"

	"git.smesh.lol/nostrino/pkg/btcwire"
)

// DefaultUserAgent identifies this client to a chain node.
const DefaultUserAgent = "/nostrino:0.1.0/"

// Config configures one peer connection.
type Config struct {
	// Magic is the chain's P2P network magic.
	Magic uint32
	// UserAgent is the BIP14 user agent. Empty means DefaultUserAgent.
	UserAgent string
	// Services is the service bits this client advertises. A light client
	// advertises none by default.
	Services uint64
	// LastBlock is the caller's best height, sent in version.
	LastBlock int32
	// ReadTimeout bounds one frame read. Zero means 60s.
	ReadTimeout time.Duration
	// RequireCompactFilters rejects a peer that does not advertise NODE
	// COMPACT_FILTERS.
	RequireCompactFilters bool
	// Logf receives diagnostics; nil discards them.
	Logf func(string, ...any)
}

// Conn is one P2P connection.
type Conn struct {
	conn   net.Conn
	reader *btcwire.Reader
	cfg    Config

	writeMu sync.Mutex
	closeMu sync.Mutex
	closed  bool

	peer *btcwire.Version
}

// New wraps an established byte stream (a TCP connection or a tunnel Conn).
func New(conn net.Conn, cfg Config) *Conn {
	if cfg.Magic == 0 {
		panic("p2p: zero network magic")
	}
	if cfg.ReadTimeout == 0 {
		cfg.ReadTimeout = 60 * time.Second
	}
	if cfg.UserAgent == "" {
		cfg.UserAgent = DefaultUserAgent
	}
	if cfg.Logf == nil {
		cfg.Logf = func(string, ...any) {}
	}
	return &Conn{
		conn:   conn,
		reader: btcwire.NewReader(conn, cfg.Magic),
		cfg:    cfg,
	}
}

// DialTCP opens a TCP connection and wraps it.
func DialTCP(ctx context.Context, addr string, cfg Config) (*Conn, error) {
	var d net.Dialer
	conn, err := d.DialContext(ctx, "tcp", addr)
	if err != nil {
		return nil, fmt.Errorf("p2p: dial %s: %w", addr, err)
	}
	return New(conn, cfg), nil
}

// PeerVersion is the version message received from the remote peer.
func (c *Conn) PeerVersion() *btcwire.Version { return c.peer }

// Handshake performs version/verack. It returns the peer's version message.
func (c *Conn) Handshake(ctx context.Context) (*btcwire.Version, error) {
	our := &btcwire.Version{
		ProtocolVersion: btcwire.ProtocolVersion,
		Services:        c.cfg.Services,
		Timestamp:       time.Now().Unix(),
		Nonce:           randomNonce(),
		UserAgent:       c.cfg.UserAgent,
		LastBlock:       c.cfg.LastBlock,
		Relay:           false,
	}
	if err := c.Send(btcwire.CmdVersion, our); err != nil {
		return nil, err
	}

	var peer *btcwire.Version
	for peer == nil {
		msg, err := c.Next(ctx)
		if err != nil {
			return nil, fmt.Errorf("p2p: waiting for version: %w", err)
		}
		if msg.Command != btcwire.CmdVersion {
			continue
		}
		var v btcwire.Version
		if err := v.DecodePayload(msg.Payload); err != nil {
			return nil, fmt.Errorf("p2p: decode version: %w", err)
		}
		peer = &v
	}
	if peer.ProtocolVersion < 70015 {
		return nil, fmt.Errorf("p2p: peer protocol version %d is too old for compact filters", peer.ProtocolVersion)
	}
	if c.cfg.RequireCompactFilters && peer.Services&btcwire.ServiceNodeCompactFilters == 0 {
		return nil, fmt.Errorf("p2p: peer does not serve compact filters (services %#x)", peer.Services)
	}
	if err := c.SendEmpty(btcwire.CmdVerAck); err != nil {
		return nil, err
	}
	// A node may announce sendheaders/wtxidrelay before verack.
	for {
		msg, err := c.Next(ctx)
		if err != nil {
			return nil, fmt.Errorf("p2p: waiting for verack: %w", err)
		}
		if msg.Command == btcwire.CmdVerAck {
			break
		}
	}
	c.peer = peer
	if err := c.SendEmpty(btcwire.CmdSendHeaders); err != nil {
		return nil, err
	}
	return peer, nil
}

func randomNonce() uint64 {
	var b [8]byte
	if _, err := rand.Read(b[:]); err != nil {
		return uint64(time.Now().UnixNano())
	}
	return binary.LittleEndian.Uint64(b[:])
}

// Send encodes and writes one message.
func (c *Conn) Send(command string, payload btcwire.Payload) error {
	var raw []byte
	if payload != nil {
		b, err := payload.EncodePayload()
		if err != nil {
			return err
		}
		raw = b
	}
	frame, err := btcwire.EncodeMessage(c.cfg.Magic, command, raw)
	if err != nil {
		return err
	}
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	c.closeMu.Lock()
	closed := c.closed
	c.closeMu.Unlock()
	if closed {
		return errors.New("p2p: closed")
	}
	if _, err := c.conn.Write(frame); err != nil {
		return fmt.Errorf("p2p: write %s: %w", command, err)
	}
	return nil
}

// SendEmpty writes a command with an empty payload.
func (c *Conn) SendEmpty(command string) error { return c.Send(command, nil) }

// Next reads the next frame. Ping is answered automatically and not returned.
func (c *Conn) Next(ctx context.Context) (*btcwire.Message, error) {
	for {
		if err := c.setReadDeadline(ctx); err != nil {
			return nil, err
		}
		msg, err := c.reader.Next()
		if err != nil {
			return nil, err
		}
		if msg.Command == btcwire.CmdPing {
			var p btcwire.Ping
			if err := p.DecodePayload(msg.Payload); err != nil {
				return nil, fmt.Errorf("p2p: decode ping: %w", err)
			}
			if err := c.Send(btcwire.CmdPong, &btcwire.Pong{Nonce: p.Nonce}); err != nil {
				return nil, err
			}
			continue
		}
		return msg, nil
	}
}

func (c *Conn) setReadDeadline(ctx context.Context) error {
	deadline := time.Now().Add(c.cfg.ReadTimeout)
	if d, ok := ctx.Deadline(); ok && d.Before(deadline) {
		deadline = d
	}
	if err := c.conn.SetReadDeadline(deadline); err != nil {
		return err
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
		return nil
	}
}

// ReadUntil reads frames until one with the given command arrives, returning
// it. Other messages are discarded.
func (c *Conn) ReadUntil(ctx context.Context, command string) (*btcwire.Message, error) {
	for {
		msg, err := c.Next(ctx)
		if err != nil {
			return nil, err
		}
		if msg.Command == command {
			return msg, nil
		}
	}
}

// ReadUntilMax reads frames until one with the given command arrives, refusing any frame whose
// payload exceeds max before the payload is allocated.
//
// It exists for one reason: a peer answering getdata can answer with anything, and a client that
// allocates whatever the header claims has handed a stranger its memory. The limit is per call
// and restored afterwards, so a block read cannot loosen the cap every other read lives under.
func (c *Conn) ReadUntilMax(ctx context.Context, command string, max int) (*btcwire.Message, error) {
	c.reader.SetLimit(max)

	defer c.reader.SetLimit(btcwire.MaxPayload)

	for {
		msg, err := c.Next(ctx)
		if err != nil {
			return nil, err
		}

		if msg.Command == command {
			return msg, nil
		}
	}
}

// Close closes the underlying connection.
func (c *Conn) Close() error {
	c.closeMu.Lock()
	if c.closed {
		c.closeMu.Unlock()
		return nil
	}
	c.closed = true
	c.closeMu.Unlock()
	return c.conn.Close()
}

// Conn exposes the underlying stream for callers that need deadlines.
func (c *Conn) Conn() net.Conn { return c.conn }