// waste-anchor: WebSocket signaling server for YAW/2.
// Deploy on your Hetzner VPS alongside a STUN server (coturn in STUN-only mode).
// The anchor never reads the content of "box" fields — it only routes sealed blobs.
package main

import (
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"flag"
	"log"
	"net/http"
	"sync"
	"time"

	"nhooyr.io/websocket"
	"nhooyr.io/websocket/wsjson"

	"github.com/waste-go/internal/proto"
)

func main() {
	bind := flag.String("bind", "0.0.0.0:17339", "address to listen on")
	flag.Parse()

	a := newAnchor()
	http.HandleFunc("/ws", a.handleWS)
	log.Printf("anchor: listening on %s", *bind)
	if err := http.ListenAndServe(*bind, nil); err != nil {
		log.Fatalf("anchor: %v", err)
	}
}

// ── Anchor ────────────────────────────────────────────────────────────────────

type client struct {
	id   string // hex peer id, set after join
	net  string // hashed network name, set after join
	send chan proto.AnchorMessage
	conn *websocket.Conn
}

type anchor struct {
	mu      sync.RWMutex
	clients map[string]*client // keyed by hex peer id
}

func newAnchor() *anchor {
	return &anchor{clients: make(map[string]*client)}
}

func (a *anchor) register(c *client) {
	a.mu.Lock()
	a.clients[c.id] = c
	a.mu.Unlock()
}

func (a *anchor) unregister(c *client) {
	if c.id == "" {
		return
	}
	a.mu.Lock()
	delete(a.clients, c.id)
	a.mu.Unlock()

	// Notify everyone in the same network.
	leave := proto.AnchorMessage{Type: proto.AnchorPeerLeave, ID: c.id}
	a.mu.RLock()
	for _, peer := range a.clients {
		if peer.net == c.net {
			select {
			case peer.send <- leave:
			default:
			}
		}
	}
	a.mu.RUnlock()
	log.Printf("anchor: peer left: %s", c.id[:min(8, len(c.id))])
}

// networkPeerIDs returns the hex ids of all peers in the same network as netHash.
func (a *anchor) networkPeerIDs(netHash, excludeID string) []string {
	a.mu.RLock()
	defer a.mu.RUnlock()
	var ids []string
	for id, c := range a.clients {
		if c.net == netHash && id != excludeID {
			ids = append(ids, id)
		}
	}
	return ids
}

func (a *anchor) forwardTo(toID string, msg proto.AnchorMessage) bool {
	a.mu.RLock()
	c, ok := a.clients[toID]
	a.mu.RUnlock()
	if !ok {
		return false
	}
	select {
	case c.send <- msg:
		return true
	default:
		return false
	}
}

// ── WebSocket handler ─────────────────────────────────────────────────────────

func (a *anchor) handleWS(w http.ResponseWriter, r *http.Request) {
	conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{InsecureSkipVerify: true})
	if err != nil {
		log.Printf("anchor: ws accept: %v", err)
		return
	}
	defer conn.Close(websocket.StatusNormalClosure, "bye")

	c := &client{
		send: make(chan proto.AnchorMessage, 64),
		conn: conn,
	}

	ctx, cancel := context.WithCancel(r.Context())
	defer cancel()

	// Send a challenge nonce immediately.
	nonce := make([]byte, 16)
	rand.Read(nonce)
	nonceHex := hex.EncodeToString(nonce)
	if err := wsjson.Write(ctx, conn, proto.AnchorMessage{
		Type:  proto.AnchorChallenge,
		Nonce: nonceHex,
	}); err != nil {
		return
	}

	// Writer goroutine.
	go func() {
		for msg := range c.send {
			if err := wsjson.Write(ctx, conn, msg); err != nil {
				cancel()
				return
			}
		}
	}()

	// Read loop.
	for {
		var msg proto.AnchorMessage
		if err := wsjson.Read(ctx, conn, &msg); err != nil {
			break
		}
		switch msg.Type {

		case proto.AnchorJoin:
			if msg.ID == "" || msg.Net == "" || msg.Sig == "" {
				log.Printf("anchor: join missing fields from %s", r.RemoteAddr)
				continue
			}
			pubBytes, err := hex.DecodeString(msg.ID)
			if err != nil || len(pubBytes) != ed25519.PublicKeySize {
				log.Printf("anchor: join: bad peer id from %s", r.RemoteAddr)
				continue
			}
			sigBytes, err := hex.DecodeString(msg.Sig)
			if err != nil {
				log.Printf("anchor: join: bad sig from %s", r.RemoteAddr)
				continue
			}
			// §5.1: sig covers nonce_raw || net_ascii (net as 64-char hex UTF-8 string)
			signed := append(nonce, []byte(msg.Net)...)
			if !ed25519.Verify(ed25519.PublicKey(pubBytes), signed, sigBytes) {
				log.Printf("anchor: join: sig verification failed for %s", msg.ID[:min(8, len(msg.ID))])
				continue
			}

			c.id = msg.ID
			c.net = msg.Net
			a.register(c)

			peers := a.networkPeerIDs(msg.Net, c.id)
			if err := wsjson.Write(ctx, conn, proto.AnchorMessage{
				Type:  proto.AnchorJoined,
				Peers: peers,
			}); err != nil {
				break
			}

			// Notify existing peers that someone new joined.
			join := proto.AnchorMessage{Type: proto.AnchorPeerJoin, ID: c.id}
			for _, pid := range peers {
				a.forwardTo(pid, join)
			}
			log.Printf("anchor: peer joined: %s net=%s peers=%d", c.id[:min(8, len(c.id))], msg.Net[:8], len(peers))

		case proto.AnchorTo:
			if c.id == "" {
				continue // not joined yet
			}
			if msg.To == "" || msg.Box == "" {
				continue
			}
			fwd := proto.AnchorMessage{
				Type: proto.AnchorFrom,
				From: c.id,
				Box:  msg.Box,
			}
			if !a.forwardTo(msg.To, fwd) {
				select {
				case c.send <- proto.AnchorMessage{Type: proto.AnchorNoPeer, ID: msg.To}:
				default:
				}
			}

		default:
			log.Printf("anchor: unknown message type %q from %s", msg.Type, r.RemoteAddr)
		}
	}

	a.unregister(c)
	close(c.send)
}

// hashNetName returns the SHA-256 hash of a plaintext network name.
// The anchor stores only hashed names — it never sees the plaintext.
func hashNetName(name string) string {
	h := sha256.Sum256([]byte("yaw2-net:" + name))
	return hex.EncodeToString(h[:])
}

// Compile-time check: hashNetName is used by daemons, kept here for reference.
var _ = hashNetName

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// Ensure wsjson import is used (avoids accidental drop).
var _ = time.Now