// waste-relay: bootstrap and relay server.
// Deploy this on your Hetzner VPS. Peers register here and can ask it
// to forward encrypted envelopes to other peers.
// The relay never sees plaintext — it only shuffles opaque blobs.
package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"log"
	"net"
	"sync"
	"time"

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

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

	ln, err := net.Listen("tcp", *bind)
	if err != nil {
		log.Fatalf("relay: listen on %s: %v", *bind, err)
	}
	log.Printf("relay: listening on %s", *bind)

	r := newRegistry()
	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Printf("relay: accept error: %v", err)
			continue
		}
		go r.handleClient(conn)
	}
}

// ── Registry ──────────────────────────────────────────────────────────────────

type registeredPeer struct {
	gossip proto.GossipEntry
	send   chan proto.RelayMessage // write to this to forward a message
}

type registry struct {
	mu    sync.RWMutex
	peers map[proto.PeerID]*registeredPeer
}

func newRegistry() *registry {
	return &registry{peers: make(map[proto.PeerID]*registeredPeer)}
}

func (r *registry) register(id proto.PeerID, entry proto.GossipEntry, send chan proto.RelayMessage) {
	r.mu.Lock()
	r.peers[id] = &registeredPeer{gossip: entry, send: send}
	r.mu.Unlock()
	log.Printf("relay: registered %s (%s)", id.Short(), entry.Peer.Alias)
}

func (r *registry) unregister(id proto.PeerID) {
	r.mu.Lock()
	delete(r.peers, id)
	r.mu.Unlock()
	log.Printf("relay: unregistered %s", id.Short())
}

func (r *registry) listPeers() []proto.GossipEntry {
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := make([]proto.GossipEntry, 0, len(r.peers))
	for _, p := range r.peers {
		out = append(out, p.gossip)
	}
	return out
}

func (r *registry) forward(to proto.PeerID, msg proto.RelayMessage) bool {
	r.mu.RLock()
	p, ok := r.peers[to]
	r.mu.RUnlock()
	if !ok {
		return false
	}
	select {
	case p.send <- msg:
		return true
	default:
		return false
	}
}

// ── Client handler ────────────────────────────────────────────────────────────

func (r *registry) handleClient(conn net.Conn) {
	defer conn.Close()
	addr := conn.RemoteAddr().String()

	sendCh := make(chan proto.RelayMessage, 64)
	var myID proto.PeerID

	// Writer goroutine
	go func() {
		enc := json.NewEncoder(conn)
		for msg := range sendCh {
			if err := enc.Encode(msg); err != nil {
				return
			}
		}
	}()

	send := func(msg proto.RelayMessage) {
		select {
		case sendCh <- msg:
		default:
		}
	}

	scanner := bufio.NewScanner(conn)
	for scanner.Scan() {
		var msg proto.RelayMessage
		if err := json.Unmarshal(scanner.Bytes(), &msg); err != nil {
			log.Printf("relay: bad message from %s: %v", addr, err)
			continue
		}

		switch msg.Type {

		case proto.RelayRegister:
			if msg.Peer == nil {
				send(proto.RelayMessage{Type: proto.RelayError, Message: "register: peer required"})
				continue
			}
			// TODO: verify msg.Signature against msg.Peer.PublicKey
			myID = msg.Peer.ID
			entry := proto.GossipEntry{
				Peer:     *msg.Peer,
				AddrHint: addr,
				LastSeen: time.Now(),
			}
			r.register(myID, entry, sendCh)

		case proto.RelayListPeers:
			send(proto.RelayMessage{
				Type:  proto.RelayPeerList,
				Peers: r.listPeers(),
			})

		case proto.RelayForward:
			if msg.To == nil || msg.Envelope == nil {
				send(proto.RelayMessage{Type: proto.RelayError, Message: "forward: to and envelope required"})
				continue
			}
			fwd := proto.RelayMessage{
				Type:     proto.RelayForwarded,
				From:     &myID,
				Envelope: msg.Envelope,
			}
			if !r.forward(*msg.To, fwd) {
				send(proto.RelayMessage{
					Type:    proto.RelayError,
					Message: "peer not connected to relay",
				})
			}

		default:
			log.Printf("relay: unknown message type %q from %s", msg.Type, addr)
		}
	}

	close(sendCh)
	if myID != "" {
		r.unregister(myID)
	}
}