package sim

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net"
	"time"
)

// Rendezvous makes a sim-engine behave like a multi-node one.
//
// It exists because the multi-node launch path — leader first, workers joining,
// only rank 0 serving, the whole replica failing when one shard dies — can only
// be tested against something that actually rendezvouses. Stub processes
// exercise the command construction and nothing else, which is how that path
// stayed unverified through four phases.
//
// The protocol is the smallest thing that has the same shape as vLLM's Ray head
// or SGLang's dist-init: rank 0 listens, every other rank dials, and rank 0
// does not serve until all of them have.
type Rendezvous struct {
	// Addr is the leader's dist-init endpoint. Empty means single node.
	Addr string
	// Rank 0 is the leader. NNodes counts every rank including it.
	Rank   int
	NNodes int
	// Timeout bounds the wait. A rendezvous that never completes must fail the
	// shard rather than hang it, which is the difference between a replica that
	// gets rebuilt and one that occupies its GPUs forever.
	Timeout time.Duration
}

// Multi reports whether this shard is part of a multi-node replica.
func (r Rendezvous) Multi() bool { return r.NNodes > 1 && r.Addr != "" }

// Leader reports whether this rank serves the API.
func (r Rendezvous) Leader() bool { return r.Rank == 0 }

// Join blocks until the replica is assembled.
//
// The leader waits for every worker to dial in; a worker dials the leader and
// holds the connection open, so the leader learns immediately when one dies.
func (r Rendezvous) Join(ctx context.Context, log *slog.Logger) error {
	if !r.Multi() {
		return nil
	}
	if r.Timeout <= 0 {
		r.Timeout = 5 * time.Minute
	}
	ctx, cancel := context.WithTimeout(ctx, r.Timeout)
	defer cancel()

	if r.Leader() {
		return r.host(ctx, log)
	}
	return r.dial(ctx, log)
}

func (r Rendezvous) host(ctx context.Context, log *slog.Logger) error {
	ln, err := net.Listen("tcp", r.Addr)
	if err != nil {
		return fmt.Errorf("binding the rendezvous address %s: %w", r.Addr, err)
	}
	// Deliberately not closed: holding it open is how a worker that restarts
	// can reconnect, and how the leader notices one going away.
	go func() {
		<-ctx.Done()
		_ = ln.Close()
	}()

	want := r.NNodes - 1
	log.Info("waiting for workers", "rank", r.Rank, "want", want, "addr", r.Addr)

	joined := 0
	for joined < want {
		conn, err := ln.Accept()
		if err != nil {
			if ctx.Err() != nil {
				return fmt.Errorf("only %d of %d workers joined before the deadline", joined, want)
			}
			return err
		}
		joined++
		log.Info("a worker joined", "joined", joined, "want", want, "from", conn.RemoteAddr())
		// Held open rather than closed: the connection is the worker's
		// liveness, exactly as a real distributed runtime would treat it.
		go func() { <-ctx.Done(); _ = conn.Close() }()
	}
	log.Info("every worker joined; serving", "nnodes", r.NNodes)
	return nil
}

func (r Rendezvous) dial(ctx context.Context, log *slog.Logger) error {
	log.Info("joining the leader", "rank", r.Rank, "addr", r.Addr)

	// Retried, because the leader is started first but "started" and "listening"
	// are not the same instant, and a worker that gives up on the first refusal
	// would make the whole replica a race.
	var lastErr error
	for ctx.Err() == nil {
		conn, err := (&net.Dialer{Timeout: 5 * time.Second}).DialContext(ctx, "tcp", r.Addr)
		if err == nil {
			log.Info("joined", "rank", r.Rank)
			go func() { <-ctx.Done(); _ = conn.Close() }()
			return nil
		}
		lastErr = err
		select {
		case <-ctx.Done():
		case <-time.After(250 * time.Millisecond):
		}
	}
	return fmt.Errorf("could not reach the leader at %s: %w", r.Addr, errors.Join(lastErr, ctx.Err()))
}