// Package bridge is the gateway side of the relay tunnel: it accepts an
// encrypted stream from a light client and connects it to a chain node's P2P
// port, forwarding only the message types a compact-filter client needs.
//
// Bridging P2P rather than exposing JSON-RPC is deliberate. The P2P protocol
// has no wallet surface, so there is nothing to spend and nothing to list; the
// whitelist keeps the stream to chain data, and the client verifies the header
// and filter chains itself, so a bridge can withhold but not forge.
package bridge

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/url"
	"strings"
	"sync"
	"time"

	"git.smesh.lol/nostrino/pkg/btcwire"
	"git.smesh.lol/nostrino/pkg/chain"
	"git.smesh.lol/nostrino/pkg/nostr"
	"git.smesh.lol/nostrino/pkg/p2p"
	"git.smesh.lol/nostrino/pkg/registry"
	"git.smesh.lol/nostrino/pkg/relay"
	"git.smesh.lol/nostrino/pkg/tunnel"
	"git.smesh.lol/nostrino/pkg/wsserver"
)

// Config configures a bridge.
type Config struct {
	// Relays are the public relays the gateway announces on.
	Relays []string
	// NodeAddr is the chain node's P2P address, for example 127.0.0.1:18444.
	NodeAddr string
	// Key is the gateway's persistent transport key.
	Key [32]byte
	// Profile is the chain served. Magic is taken from it.
	Profile *chain.Profile
	// Chains is the chain name advertised to clients. Defaults to
	// Profile.Name.
	Chains []string
	// Allowlist is an optional set of client pubkeys. Empty serves anyone.
	Allowlist []string
	// MaxSessions bounds concurrent streams. Zero means 32.
	MaxSessions int
	// NodeDialTimeout bounds the node connection. Zero means 10s.
	NodeDialTimeout time.Duration
	// IdleTimeout closes a stream with no traffic. Zero means 2 minutes.
	IdleTimeout time.Duration
	// MaxBytesPerSec bounds each direction of a stream. Zero means unlimited;
	// the command defaults to 4 MiB/s.
	MaxBytesPerSec int64
	// MaxSessionsPerClient bounds one client pubkey's concurrent streams.
	// Zero means 2.
	MaxSessionsPerClient int
	// Presence publishes a discoverable gateway announcement. Zero value
	// disables it; the command enables it explicitly.
	Presence bool
	// PresenceInterval is the announcement period. Zero means 5 minutes.
	PresenceInterval time.Duration
	// Version is reported in presence and diagnostics.
	Version string
	// BridgeListen enables the WebSocket data plane on this TCP address.
	// Empty disables it.
	BridgeListen string
	// BridgePath is the WebSocket request path. Empty means /nostrino.
	BridgePath string
	// BridgeURL is the advertised bridge URL. Empty derives it from
	// BridgeListen and BridgePath; set it to a wss:// URL behind a proxy.
	BridgeURL string
	// BridgeToken, when set, is required as the token query parameter or as
	// a bearer Authorization header on the WebSocket upgrade.
	BridgeToken string
	// AllowCommands adds to the default BIP157 whitelist.
	AllowCommands []string
	// Logf receives diagnostics; nil discards them.
	Logf func(string, ...any)
	// TransportFactory overrides the relay WebSocket dialer, which is how
	// tests inject an in-process relay.
	TransportFactory func(url string) (relay.Transport, error)
}

// Server is a running bridge.
type Server struct {
	cfg      Config
	listener *tunnel.Listener
	ws       *wsserver.Server

	mu       sync.Mutex
	conns    map[net.Conn]struct{}
	wsActive int
	wsPerIP  map[string]int
	done     chan struct{}
	wg       sync.WaitGroup
}

// MaxBlockPayload is the largest block this gateway will forward.
//
// The client has the same cap, and it belongs on both ends: a cap that lives only in the client
// protects the client from the node, and this one protects the gateway from a node that answers
// a getdata with more than anybody asked for.
const MaxBlockPayload = 8 << 20

// DefaultCommands is the P2P whitelist. Everything a compact-filter client
// needs is here; address discovery, bloom filters, mempool queries and BIP152
// block relay are not.
var DefaultCommands = []string{
	btcwire.CmdVersion,
	btcwire.CmdVerAck,
	btcwire.CmdPing,
	btcwire.CmdPong,
	btcwire.CmdSendHeaders,
	btcwire.CmdWtxidRelay,
	btcwire.CmdSendAddrV2,
	btcwire.CmdFeeFilter,
	btcwire.CmdGetHeaders,
	btcwire.CmdHeaders,
	btcwire.CmdGetCFHeaders,
	btcwire.CmdCFHeaders,
	btcwire.CmdGetCFilters,
	btcwire.CmdCFilter,
	btcwire.CmdGetCFCheckpt,
	btcwire.CmdCFCheckpt,
	btcwire.CmdGetData,
	btcwire.CmdTx,
	btcwire.CmdInv,
	btcwire.CmdNotFound,
	// Core announces compact blocks straight after the handshake whether or not the
	// client asked for BIP152. Dropping it closes every session 100ms in.
	btcwire.CmdSendCmpct,
}

// Start connects the gateway to its relays and returns a running server.
func Start(ctx context.Context, cfg Config) (*Server, error) {
	if cfg.Profile == nil {
		return nil, errors.New("bridge: nil profile")
	}
	if err := cfg.Profile.Validate(); err != nil {
		return nil, err
	}
	if cfg.NodeAddr == "" {
		return nil, errors.New("bridge: no node address")
	}
	if cfg.Key == ([32]byte{}) {
		return nil, errors.New("bridge: no gateway key")
	}
	if len(cfg.Relays) == 0 {
		return nil, errors.New("bridge: no relay")
	}
	if cfg.MaxSessions == 0 {
		cfg.MaxSessions = 32
	}
	if cfg.NodeDialTimeout == 0 {
		cfg.NodeDialTimeout = 10 * time.Second
	}
	if cfg.IdleTimeout == 0 {
		cfg.IdleTimeout = 2 * time.Minute
	}
	if cfg.MaxSessionsPerClient == 0 {
		cfg.MaxSessionsPerClient = 2
	}
	if cfg.PresenceInterval == 0 {
		cfg.PresenceInterval = 5 * time.Minute
	}
	if cfg.Version == "" {
		cfg.Version = "nostrino"
	}
	if cfg.Logf == nil {
		cfg.Logf = func(string, ...any) {}
	}
	chains := cfg.Chains
	if len(chains) == 0 {
		chains = []string{cfg.Profile.Name}
	}
	l, err := tunnel.Listen(ctx, tunnel.ListenConfig{
		Relays:           cfg.Relays,
		Key:              cfg.Key,
		Chains:           chains,
		Allowlist:        cfg.Allowlist,
		MaxSessions:      cfg.MaxSessions,
		IdleTimeout:      cfg.IdleTimeout,
		Logf:             cfg.Logf,
		TransportFactory: cfg.TransportFactory,
	})
	if err != nil {
		return nil, err
	}
	cfg.Chains = chains
	s := &Server{
		cfg:      cfg,
		listener: l,
		conns:    make(map[net.Conn]struct{}),
		wsPerIP:  make(map[string]int),
		done:     make(chan struct{}),
	}
	if cfg.BridgeListen != "" {
		if cfg.BridgePath == "" {
			cfg.BridgePath = "/nostrino"
		}
		ws, err := wsserver.Start(wsserver.Config{
			Listen:  cfg.BridgeListen,
			Path:    cfg.BridgePath,
			Handler: s.handleWS,
			Logf:    cfg.Logf,
		})
		if err != nil {
			l.Close()
			return nil, err
		}
		s.ws = ws
		if cfg.BridgeURL == "" {
			cfg.BridgeURL = "ws://" + ws.Addr() + cfg.BridgePath
		}
		if cfg.BridgeToken != "" {
			sep := "?"
			if strings.Contains(cfg.BridgeURL, "?") {
				sep = "&"
			}
			cfg.BridgeURL += sep + "token=" + url.QueryEscape(cfg.BridgeToken)
		}
	}
	s.cfg = cfg
	if cfg.Presence {
		s.wg.Add(1)
		go s.presenceLoop(ctx)
	}
	return s, nil
}

// PubKey is the gateway transport public key.
func (s *Server) PubKey() [32]byte { return s.listener.PubKey() }

// URI returns the connection URI a client should use.
func (s *Server) URI() *tunnel.URI {
	u := &tunnel.URI{
		Gateway: s.listener.PubKey(),
		Relays:  s.cfg.Relays,
		Chains:  s.cfg.Chains,
		Version: 1,
	}
	if s.cfg.BridgeURL != "" {
		u.Bridges = []string{s.cfg.BridgeURL}
	}
	return u
}

// BridgeAddr is the bound WebSocket address, or empty when the data plane is
// disabled.
func (s *Server) BridgeAddr() string {
	if s.ws == nil {
		return ""
	}
	return s.ws.Addr()
}

func (s *Server) handleWS(c *wsserver.Conn, req *wsserver.Request) {
	if s.cfg.BridgeToken != "" {
		token := req.Query.Get("token")
		if token == "" {
			auth := req.Header["authorization"]
			if strings.HasPrefix(auth, "Bearer ") {
				token = strings.TrimPrefix(auth, "Bearer ")
			}
		}
		if token != s.cfg.BridgeToken {
			s.cfg.Logf("bridge: websocket client %s rejected: bad token", c.RemoteAddr())
			return
		}
	}
	chainName := req.Query.Get("chain")
	if chainName == "" && len(s.cfg.Chains) > 0 {
		chainName = s.cfg.Chains[0]
	}
	allowed := false
	for _, name := range s.cfg.Chains {
		if name == chainName {
			allowed = true
			break
		}
	}
	if !allowed {
		s.cfg.Logf("bridge: websocket client asked for chain %q", chainName)
		return
	}
	ip := c.RemoteAddr().String()
	if host, _, err := net.SplitHostPort(ip); err == nil {
		ip = host
	}
	s.mu.Lock()
	if s.wsActive >= s.cfg.MaxSessions {
		s.mu.Unlock()
		s.cfg.Logf("bridge: websocket session limit reached")
		return
	}
	if s.wsPerIP[ip] >= s.cfg.MaxSessionsPerClient {
		s.mu.Unlock()
		s.cfg.Logf("bridge: websocket per-client limit for %s", ip)
		return
	}
	s.wsActive++
	s.wsPerIP[ip]++
	s.mu.Unlock()
	defer func() {
		s.mu.Lock()
		s.wsActive--
		if s.wsPerIP[ip] > 0 {
			s.wsPerIP[ip]--
		}
		if s.wsPerIP[ip] == 0 {
			delete(s.wsPerIP, ip)
		}
		s.mu.Unlock()
	}()
	s.handleConn(context.Background(), c, chainName, "ws:"+c.RemoteAddr().String())
}

// Run accepts streams until ctx ends or the server is closed.
func (s *Server) Run(ctx context.Context) error {
	for {
		conn, err := s.listener.Accept(ctx)
		if err != nil {
			select {
			case <-s.done:
				return nil
			case <-ctx.Done():
				return ctx.Err()
			default:
				return err
			}
		}
		s.wg.Add(1)
		go func() {
			defer s.wg.Done()
			s.handle(ctx, conn)
		}()
	}
}

func (s *Server) handle(ctx context.Context, c *tunnel.Conn) {
	defer c.Close()
	s.handleConn(ctx, c, c.Chain(), "nostr:"+c.PeerHex())
}

func (s *Server) handleConn(ctx context.Context, c net.Conn, chainName, peer string) {
	allowedChain := false
	for _, name := range s.cfg.Chains {
		if name == chainName {
			allowedChain = true
			break
		}
	}
	if !allowedChain {
		s.cfg.Logf("bridge: client %s asked for chain %q", peer, chainName)
		return
	}

	dialCtx, cancel := context.WithTimeout(ctx, s.cfg.NodeDialTimeout)
	defer cancel()
	var d net.Dialer
	nodeConn, err := d.DialContext(dialCtx, "tcp", s.cfg.NodeAddr)
	if err != nil {
		s.cfg.Logf("bridge: node dial: %v", err)
		return
	}
	defer nodeConn.Close()
	s.track(nodeConn, true)
	defer s.track(nodeConn, false)

	s.cfg.Logf("bridge: session %s client=%s chain=%s node=%s", c.LocalAddr(), peer, chainName, s.cfg.NodeAddr)
	start := time.Now()
	var bytesUp, bytesDown int64
	upLimiter := newLimiter(s.cfg.MaxBytesPerSec)
	downLimiter := newLimiter(s.cfg.MaxBytesPerSec)

	errc := make(chan error, 2)
	go func() { errc <- s.pump(c, nodeConn, &bytesUp, upLimiter) }()
	go func() { errc <- s.pump2(c, nodeConn, &bytesDown, downLimiter) }()
	// The first error is the one that ended the session; the second is the refusal that
	// follows the close. Reporting only the counters left a dead tunnel with no reason.
	first := <-errc
	c.Close()
	nodeConn.Close()
	<-errc
	s.cfg.Logf("bridge: session %s closed after %s up=%d down=%d: %v",
		c.LocalAddr(), time.Since(start).Round(time.Millisecond), bytesUp, bytesDown, first)
}

func (s *Server) track(c net.Conn, add bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if add {
		s.conns[c] = struct{}{}
	} else {
		delete(s.conns, c)
	}
}

// pump forwards tunnel -> node.
func (s *Server) pump(tunnelConn net.Conn, node net.Conn, bytes *int64, limit *limiter) error {
	return s.copyFrames(tunnelConn, node, bytes, limit, "client")
}

// pump2 forwards node -> tunnel.
func (s *Server) pump2(tunnelConn net.Conn, node net.Conn, bytes *int64, limit *limiter) error {
	return s.copyFrames(node, tunnelConn, bytes, limit, "node")
}

func (s *Server) copyFrames(src, dst net.Conn, bytes *int64, limit *limiter, from string) error {
	reader := btcwire.NewReader(src, s.cfg.Profile.Magic)
	for {
		if s.cfg.IdleTimeout > 0 {
			if d, ok := src.(interface{ SetReadDeadline(time.Time) error }); ok {
				_ = d.SetReadDeadline(time.Now().Add(s.cfg.IdleTimeout))
			}
		}
		msg, err := reader.Next()
		if err != nil {
			return err
		}
		if !s.allowed(msg.Command) {
			return fmt.Errorf("bridge: %s sent disallowed command %q", from, msg.Command)
		}
		if msg.Command == btcwire.CmdBlock && len(msg.Payload) > MaxBlockPayload {
			return fmt.Errorf("bridge: %s sent a %d-byte block, the cap is %d", from, len(msg.Payload), MaxBlockPayload)
		}
		frame, err := btcwire.EncodeMessage(s.cfg.Profile.Magic, msg.Command, msg.Payload)
		if err != nil {
			return err
		}
		limit.wait(int64(len(frame)))
		if _, err := dst.Write(frame); err != nil {
			return err
		}
		*bytes += int64(len(frame))
	}
}

func (s *Server) allowed(command string) bool {
	for _, cmd := range DefaultCommands {
		if cmd == command {
			return true
		}
	}
	for _, cmd := range s.cfg.AllowCommands {
		if cmd == command {
			return true
		}
	}
	return false
}

func (s *Server) presenceLoop(ctx context.Context) {
	defer s.wg.Done()
	ticker := time.NewTicker(s.cfg.PresenceInterval)
	defer ticker.Stop()
	for {
		s.publishPresence(ctx)
		select {
		case <-s.done:
			return
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
	}
}

func (s *Server) publishPresence(ctx context.Context) {
	info := registry.Info{
		Version: s.cfg.Version,
		URI:     s.URI().String(),
		Chains:  s.cfg.Chains,
		Relays:  s.cfg.Relays,
	}
	tipCtx, cancel := context.WithTimeout(ctx, 6*time.Second)
	if node, err := CheckNode(tipCtx, s.cfg.NodeAddr, s.cfg.Profile, 5*time.Second); err == nil {
		info.Tip = int64(node.TipHeight)
	}
	cancel()
	ev, err := registry.Event(info, s.cfg.Key)
	if err != nil {
		s.cfg.Logf("bridge: presence: %v", err)
		return
	}
	pctx, pcancel := context.WithTimeout(ctx, 10*time.Second)
	defer pcancel()
	if err := s.listener.Publish(pctx, ev); err != nil {
		s.cfg.Logf("bridge: presence: %v", err)
		return
	}
	s.cfg.Logf("bridge: announced gateway as %s", nostr.Hex(ev.PubKey[:]))
}

// Close stops the listener and closes every node connection.
func (s *Server) Close() error {
	if s.ws != nil {
		s.ws.Close()
	}
	select {
	case <-s.done:
	default:
		close(s.done)
	}
	err := s.listener.Close()
	s.mu.Lock()
	for c := range s.conns {
		c.Close()
	}
	s.mu.Unlock()
	s.wg.Wait()
	return err
}

// NodeInfo is the result of a node compatibility check.
type NodeInfo struct {
	ProtocolVersion int32
	Services        uint64
	UserAgent       string
	TipHeight       int32
	CompactFilters  bool
}

// CheckNode connects to a node, handshakes, and reports whether it advertises
// compact filters. It is the bridge operator's preflight check.
func CheckNode(ctx context.Context, addr string, profile *chain.Profile, timeout time.Duration) (*NodeInfo, error) {
	if profile == nil {
		return nil, errors.New("bridge: nil profile")
	}
	if timeout == 0 {
		timeout = 10 * time.Second
	}
	dctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	peer, err := p2p.DialTCP(dctx, addr, p2p.Config{
		Magic:                 profile.Magic,
		UserAgent:             p2p.DefaultUserAgent,
		ReadTimeout:           timeout,
		RequireCompactFilters: false,
	})
	if err != nil {
		return nil, err
	}
	defer peer.Close()
	v, err := peer.Handshake(dctx)
	if err != nil {
		return nil, err
	}
	return &NodeInfo{
		ProtocolVersion: v.ProtocolVersion,
		Services:        v.Services,
		UserAgent:       v.UserAgent,
		TipHeight:       v.LastBlock,
		CompactFilters:  v.Services&btcwire.ServiceNodeCompactFilters != 0,
	}, nil
}