// Package testnode is a deterministic Bitcoin P2P peer that serves a short
// regtest chain with BIP157 compact filters. It exists so the tunnel, bridge
// and light-client code can be exercised end to end without a full node.
package testnode

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"sync"

	"git.smesh.lol/nostrino/pkg/btcwire"
	"git.smesh.lol/nostrino/pkg/chain"
	"git.smesh.lol/nostrino/pkg/crypto/chainhash"
	"git.smesh.lol/nostrino/pkg/filter"
)

// Config builds a fake chain.
type Config struct {
	// Blocks is the number of blocks after genesis. Minimum 1.
	Blocks int
	// Bits is the compact target. Zero means the regtest target 0x207fffff.
	Bits uint32
}

// Node is a fake peer with a generated chain.
type Node struct {
	profile       *chain.Profile
	headers       []btcwire.BlockHeader
	filters       [][]byte
	filterHashes  []chainhash.Hash
	filterHeaders []chainhash.Hash
	scripts       [][]byte

	mu     sync.Mutex
	ln     net.Listener
	closed bool
	conns  map[net.Conn]struct{}
	wg     sync.WaitGroup
}

// New generates a regtest chain of cfg.Blocks headers. The genesis header is
// the profile's.
func New(profile *chain.Profile, cfg Config) (*Node, error) {
	if profile == nil {
		return nil, errors.New("testnode: nil profile")
	}
	if profile.Name != "regtest" {
		return nil, fmt.Errorf("testnode: profile %q is not regtest", profile.Name)
	}
	if cfg.Blocks < 1 {
		cfg.Blocks = 1
	}
	if cfg.Bits == 0 {
		cfg.Bits = 0x207fffff
	}
	n := &Node{
		profile: profile,
		conns:   make(map[net.Conn]struct{}),
	}
	n.headers = append(n.headers, profile.GenesisHeader)
	n.filters = append(n.filters, filter.BuildFromScripts([][]byte{[]byte("block-0")}, [32]byte(profile.GenesisHash)))
	n.scripts = append(n.scripts, []byte("block-0"))
	for i := 1; i <= cfg.Blocks; i++ {
		prev := n.headers[i-1]
		prevHash := prev.BlockHash()
		merkle := chainhash.DoubleHashH([]byte(fmt.Sprintf("merkle-%d", i)))
		h := btcwire.BlockHeader{
			Version:    1,
			PrevBlock:  prevHash,
			MerkleRoot: merkle,
			Timestamp:  uint32(1296688602 + i*600),
			Bits:       cfg.Bits,
		}
		for nonce := uint32(0); nonce < 1_000_000; nonce++ {
			h.Nonce = nonce
			if h.CheckProofOfWork() {
				break
			}
		}
		if !h.CheckProofOfWork() {
			return nil, fmt.Errorf("testnode: could not mine block %d", i)
		}
		n.headers = append(n.headers, h)
		script := []byte(fmt.Sprintf("block-%d", i))
		n.scripts = append(n.scripts, script)
		f := filter.BuildFromScripts([][]byte{script}, [32]byte(h.BlockHash()))
		n.filters = append(n.filters, f)
	}
	for i, f := range n.filters {
		fh := chainhash.DoubleHashH(f)
		n.filterHashes = append(n.filterHashes, fh)
		var prev chainhash.Hash
		if i > 0 {
			prev = n.filterHeaders[i-1]
		}
		n.filterHeaders = append(n.filterHeaders, btcwire.FilterHeader(fh, prev))
	}
	return n, nil
}

// Height is the chain tip height.
func (n *Node) Height() uint32 { return uint32(len(n.headers) - 1) }

// Header returns the header at height.
func (n *Node) Header(height uint32) (btcwire.BlockHeader, bool) {
	if int(height) >= len(n.headers) {
		return btcwire.BlockHeader{}, false
	}
	return n.headers[height], true
}

// Filter returns the compact filter at height.
func (n *Node) Filter(height uint32) ([]byte, bool) {
	if int(height) >= len(n.filters) {
		return nil, false
	}
	return n.filters[height], true
}

// Script returns the script embedded in the filter at height.
func (n *Node) Script(height uint32) []byte {
	if int(height) >= len(n.scripts) {
		return nil
	}
	return n.scripts[height]
}

// Listen starts the node on addr (127.0.0.1:0 for any port).
func (n *Node) Listen(addr string) (string, error) {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return "", err
	}
	n.mu.Lock()
	n.ln = ln
	n.mu.Unlock()
	n.wg.Add(1)
	go n.acceptLoop(ln)
	return ln.Addr().String(), nil
}

func (n *Node) acceptLoop(ln net.Listener) {
	defer n.wg.Done()
	for {
		conn, err := ln.Accept()
		if err != nil {
			return
		}
		n.mu.Lock()
		if n.closed {
			n.mu.Unlock()
			conn.Close()
			return
		}
		n.conns[conn] = struct{}{}
		n.mu.Unlock()
		n.wg.Add(1)
		go func() {
			defer n.wg.Done()
			defer func() {
				n.mu.Lock()
				delete(n.conns, conn)
				n.mu.Unlock()
				conn.Close()
			}()
			n.serve(conn)
		}()
	}
}

// Close stops the listener and all connections.
func (n *Node) Close() error {
	n.mu.Lock()
	if n.closed {
		n.mu.Unlock()
		return nil
	}
	n.closed = true
	ln := n.ln
	conns := make([]net.Conn, 0, len(n.conns))
	for c := range n.conns {
		conns = append(conns, c)
	}
	n.mu.Unlock()
	if ln != nil {
		ln.Close()
	}
	for _, c := range conns {
		c.Close()
	}
	n.wg.Wait()
	return nil
}

func (n *Node) serve(conn net.Conn) {
	reader := btcwire.NewReader(conn, n.profile.Magic)
	write := func(cmd 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(n.profile.Magic, cmd, raw)
		if err != nil {
			return err
		}
		_, err = conn.Write(frame)
		return err
	}
	for {
		msg, err := reader.Next()
		if err != nil {
			if errors.Is(err, io.EOF) {
				return
			}
			return
		}
		switch msg.Command {
		case btcwire.CmdVersion:
			v := &btcwire.Version{
				ProtocolVersion: btcwire.ProtocolVersion,
				Services:        btcwire.ServiceNodeNetwork | btcwire.ServiceNodeWitness | btcwire.ServiceNodeCompactFilters,
				Timestamp:       1700000000,
				Nonce:           42,
				UserAgent:       "/nostrino-testnode:0.1.0/",
				LastBlock:       int32(n.Height()),
				Relay:           false,
			}
			if err := write(btcwire.CmdVersion, v); err != nil {
				return
			}
			if err := write(btcwire.CmdVerAck, nil); err != nil {
				return
			}
		case btcwire.CmdVerAck:
			// handshake complete
		case btcwire.CmdPing:
			var p btcwire.Ping
			if p.DecodePayload(msg.Payload) == nil {
				if write(btcwire.CmdPong, &btcwire.Pong{Nonce: p.Nonce}) != nil {
					return
				}
			}
		case btcwire.CmdGetHeaders:
			var req btcwire.GetHeaders
			if err := req.DecodePayload(msg.Payload); err != nil {
				return
			}
			if err := n.sendHeaders(write, &req); err != nil {
				return
			}
		case btcwire.CmdGetCFHeaders:
			var req btcwire.GetCFHeaders
			if err := req.DecodePayload(msg.Payload); err != nil {
				return
			}
			if err := n.sendCFHeaders(write, &req); err != nil {
				return
			}
		case btcwire.CmdGetCFilters:
			var req btcwire.GetCFilters
			if err := req.DecodePayload(msg.Payload); err != nil {
				return
			}
			if err := n.sendCFilters(write, &req); err != nil {
				return
			}
		case btcwire.CmdGetCFCheckpt:
			var req btcwire.GetCFCheckpt
			if err := req.DecodePayload(msg.Payload); err != nil {
				return
			}
			if err := n.sendCFCheckpt(write, &req); err != nil {
				return
			}
		case btcwire.CmdSendHeaders, btcwire.CmdWtxidRelay, btcwire.CmdSendAddrV2, btcwire.CmdFeeFilter:
			// accepted, no response
		default:
			// ignore unknown commands
		}
	}
}

func (n *Node) heightOf(h chainhash.Hash) (uint32, bool) {
	for i := range n.headers {
		if n.headers[i].BlockHash() == h {
			return uint32(i), true
		}
	}
	return 0, false
}

func (n *Node) startAfterLocator(locator []chainhash.Hash) uint32 {
	for _, h := range locator {
		if height, ok := n.heightOf(h); ok {
			return height + 1
		}
	}
	return 0
}

func (n *Node) sendHeaders(write func(string, btcwire.Payload) error, req *btcwire.GetHeaders) error {
	start := n.startAfterLocator(req.LocatorHashes)
	end := uint32(len(n.headers))
	if req.StopHash != (chainhash.Hash{}) {
		if h, ok := n.heightOf(req.StopHash); ok {
			end = h + 1
		}
	}
	if start >= uint32(len(n.headers)) || start >= end {
		return write(btcwire.CmdHeaders, &btcwire.Headers{})
	}
	if end > start+btcwire.MaxBlockHeadersPerMsg {
		end = start + btcwire.MaxBlockHeadersPerMsg
	}
	return write(btcwire.CmdHeaders, &btcwire.Headers{Headers: n.headers[start:end]})
}

func (n *Node) sendCFHeaders(write func(string, btcwire.Payload) error, req *btcwire.GetCFHeaders) error {
	if req.FilterType != btcwire.FilterTypeBasic {
		return write(btcwire.CmdCFHeaders, &btcwire.CFHeaders{FilterType: req.FilterType, StopHash: req.StopHash})
	}
	start := req.StartHeight
	if start > n.Height() {
		return write(btcwire.CmdCFHeaders, &btcwire.CFHeaders{FilterType: req.FilterType, StopHash: req.StopHash})
	}
	end := uint32(len(n.headers) - 1)
	if req.StopHash != (chainhash.Hash{}) {
		if h, ok := n.heightOf(req.StopHash); ok {
			end = h
		}
	}
	if end > start+btcwire.MaxFilterHashesPerMsg-1 {
		end = start + btcwire.MaxFilterHashesPerMsg - 1
	}
	var prev chainhash.Hash
	if start > 0 {
		prev = n.filterHeaders[start-1]
	}
	hashes := append([]chainhash.Hash(nil), n.filterHashes[start:end+1]...)
	stop, ok := n.Header(end)
	if !ok {
		return errors.New("testnode: no stop header")
	}
	return write(btcwire.CmdCFHeaders, &btcwire.CFHeaders{
		FilterType:       btcwire.FilterTypeBasic,
		StopHash:         stop.BlockHash(),
		PrevFilterHeader: prev,
		FilterHashes:     hashes,
	})
}

func (n *Node) sendCFilters(write func(string, btcwire.Payload) error, req *btcwire.GetCFilters) error {
	if req.FilterType != btcwire.FilterTypeBasic {
		return errors.New("testnode: unsupported filter type")
	}
	start := req.StartHeight
	end := uint32(len(n.headers) - 1)
	if req.StopHash != (chainhash.Hash{}) {
		h, ok := n.heightOf(req.StopHash)
		if !ok {
			return errors.New("testnode: unknown stop hash")
		}
		end = h
	}
	if start > end {
		return nil
	}
	for h := start; h <= end; h++ {
		hdr := n.headers[h]
		if err := write(btcwire.CmdCFilter, &btcwire.CFilter{
			FilterType: btcwire.FilterTypeBasic,
			BlockHash:  hdr.BlockHash(),
			Filter:     n.filters[h],
		}); err != nil {
			return err
		}
	}
	return nil
}

func (n *Node) sendCFCheckpt(write func(string, btcwire.Payload) error, req *btcwire.GetCFCheckpt) error {
	end := uint32(len(n.headers) - 1)
	if req.StopHash != (chainhash.Hash{}) {
		h, ok := n.heightOf(req.StopHash)
		if !ok {
			h = end
		}
		end = h
	}
	// One checkpoint per 1000 filters, as BIP157 recommends for serving.
	step := uint32(1000)
	var pts []chainhash.Hash
	for h := step - 1; h <= end; h += step {
		pts = append(pts, n.filterHeaders[h])
	}
	stop, _ := n.Header(end)
	return write(btcwire.CmdCFCheckpt, &btcwire.CFCheckpt{
		FilterType:    btcwire.FilterTypeBasic,
		StopHash:      stop.BlockHash(),
		FilterHeaders: pts,
	})
}

// MerkleOf is a helper for callers that build their own headers.
func MerkleOf(height uint32) chainhash.Hash {
	var b [4]byte
	binary.LittleEndian.PutUint32(b[:], height)
	return chainhash.DoubleHashH(b[:])
}

// FilterHeader returns the rolling filter header at height.
func (n *Node) FilterHeader(height uint32) (chainhash.Hash, bool) {
	if int(height) >= len(n.filterHeaders) {
		return chainhash.Hash{}, false
	}
	return n.filterHeaders[height], true
}

// FilterHash returns the hash of the compact filter at height.
func (n *Node) FilterHash(height uint32) (chainhash.Hash, bool) {
	if int(height) >= len(n.filterHashes) {
		return chainhash.Hash{}, false
	}
	return n.filterHashes[height], true
}

// GenesisHash returns the profile's genesis hash.
func (n *Node) GenesisHash() chainhash.Hash { return n.profile.GenesisHash }

// ForkAt replaces every block above height with blocks new ones built on the
// common ancestor, producing a chain the same length but a different branch.
// It is used to test reorg handling.
func (n *Node) ForkAt(height uint32, blocks int, tag string) error {
	if int(height) >= len(n.headers) {
		return fmt.Errorf("testnode: fork height %d above tip %d", height, len(n.headers)-1)
	}
	n.headers = n.headers[:height+1]
	n.filters = n.filters[:height+1]
	n.filterHashes = n.filterHashes[:height+1]
	n.filterHeaders = n.filterHeaders[:height+1]
	n.scripts = n.scripts[:height+1]
	for i := 1; i <= blocks; i++ {
		h := height + uint32(i)
		prev := n.headers[h-1]
		prevHash := prev.BlockHash()
		merkle := chainhash.DoubleHashH([]byte(fmt.Sprintf("merkle-%d-%s", h, tag)))
		hdr := btcwire.BlockHeader{
			Version:    1,
			PrevBlock:  prevHash,
			MerkleRoot: merkle,
			Timestamp:  prev.Timestamp + 600,
			Bits:       prev.Bits,
		}
		for nonce := uint32(0); nonce < 1_000_000; nonce++ {
			hdr.Nonce = nonce
			if hdr.CheckProofOfWork() {
				break
			}
		}
		if !hdr.CheckProofOfWork() {
			return fmt.Errorf("testnode: could not mine fork block %d", h)
		}
		n.headers = append(n.headers, hdr)
		script := []byte(fmt.Sprintf("block-%d-%s", h, tag))
		n.scripts = append(n.scripts, script)
		f := filter.BuildFromScripts([][]byte{script}, [32]byte(hdr.BlockHash()))
		n.filters = append(n.filters, f)
		fh := chainhash.DoubleHashH(f)
		n.filterHashes = append(n.filterHashes, fh)
		n.filterHeaders = append(n.filterHeaders, btcwire.FilterHeader(fh, n.filterHeaders[h-1]))
	}
	return nil
}