package engine

import (
	"context"
	"sync"

	"github.com/miekg/dns"
)

// flightKey identifies exchanges that would produce the same upstream answer:
// the cache key, the upstream set an ACL group pinned, and the client subnet
// that was forwarded (an ECS-aware upstream answers by subnet).
type flightKey struct {
	name, ups, ecs string
	qtype, qclass  uint16
	do             bool
}

// flight is one upstream exchange other queries may wait on.
type flight struct {
	done chan struct{}
	up   *dns.Msg
	addr string
	err  error
	// leaderGaveUp is true when the leader's own context ended the exchange.
	// That says nothing about the upstream, so waiters try again themselves.
	leaderGaveUp bool
}

// flightGroup coalesces identical concurrent upstream exchanges, so a burst of
// cache misses for one name (a restart, a cache flush, a viral link) costs the
// upstream one query rather than one per client.
//
// The leader runs the exchange on its own goroutine and its own context;
// there is no detached work, so a shutdown that waits for in-flight queries
// still waits for every exchange.
type flightGroup struct {
	mu sync.Mutex
	m  map[flightKey]*flight
}

// do runs fn once per key among concurrent callers. shared is true for a
// caller that received another caller's answer; its message is a private copy.
func (g *flightGroup) do(ctx context.Context, key flightKey, fn func() (*dns.Msg, string, error)) (up *dns.Msg, addr string, shared bool, err error) {
	for {
		g.mu.Lock()
		if g.m == nil {
			g.m = make(map[flightKey]*flight)
		}
		if f, ok := g.m[key]; ok {
			g.mu.Unlock()
			select {
			case <-f.done:
			case <-ctx.Done():
				return nil, "", true, ctx.Err()
			}
			if f.leaderGaveUp {
				continue // the leader's client left; take over
			}
			if f.err != nil {
				return nil, f.addr, true, f.err
			}
			return f.up.Copy(), f.addr, true, nil
		}
		f := &flight{done: make(chan struct{})}
		g.m[key] = f
		g.mu.Unlock()

		f.up, f.addr, f.err = fn()
		f.leaderGaveUp = f.err != nil && ctx.Err() != nil
		g.mu.Lock()
		delete(g.m, key)
		g.mu.Unlock()
		close(f.done)
		// The leader keeps the original; waiters get copies. f.up is read by
		// waiters only after done is closed, and the leader's later changes to
		// its message must not reach them, so hand the leader a copy instead.
		if f.err == nil && f.up != nil {
			return f.up.Copy(), f.addr, false, nil
		}
		return f.up, f.addr, false, f.err
	}
}