package edge

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	"sync"
	"syscall"
	"time"

	"github.com/carlosframework/platform/internal/unixdial"
)

// genproxy.go is the edge half of generations (zero-downtime spec §7.2): a
// reverse proxy for an activator-managed route whose app socket moves from
// one generation to the next. Every request asks the activator for the
// current generation and uses that generation's pool, whose connections all
// reach that generation's socket (dialGen), so the response is attributed to
// the generation that served it. A connection to generation n is never
// reused for n+1, and a generation's idle connections are closed when it
// stops being current.
//
// Used only when the activator's GenerationsOn is true. With it false the
// edge builds unixProxy(t.Socket) exactly as before.

// errNoGeneration is a dial's answer while the route has no current
// generation (between a stop and the next wake). It wraps ENOENT on purpose:
// to dialUnixRetrying it is the same "socket briefly absent" gap a restart
// always had, retried inside the same window, and each retry asks again.
var errNoGeneration = fmt.Errorf("edge: instance has no current generation: %w", syscall.ENOENT)

// dialUnixRetrying is dialUnixWithRetry with the socket resolved per attempt:
// resolve is asked before the first dial and again before every retry, so a
// dial that races a switch lands on the generation that is current by then.
func dialUnixRetrying(ctx context.Context, d *net.Dialer, resolve func() (string, bool)) (net.Conn, error) {
	const window = 13 * time.Second
	deadline := time.Now().Add(window)
	for {
		var (
			conn net.Conn
			err  error
		)
		if sock, ok := resolve(); ok {
			conn, err = unixdial.Checked(ctx, d, sock)
		} else {
			err = errNoGeneration
		}
		if err == nil {
			return conn, nil
		}
		if ctx.Err() != nil || !isRetryableDialErr(err) || time.Now().After(deadline) {
			return nil, err
		}
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-time.After(100 * time.Millisecond):
		}
	}
}

// genMovedError is a generation pool's dial giving up: its own socket is
// gone and another generation is current. Nothing was sent, so the request
// can go again on the current generation's pool.
type genMovedError struct{ to GenInfo }

func (e *genMovedError) Error() string {
	return fmt.Sprintf("edge: generation moved to %d while dialling the one before it", e.to.N)
}

// dialGen is a generation pool's dial: own's socket and nothing else, and
// a connection proven to be own's (a shared foreign path is checked
// against the current generation once connected), so every connection in
// the pool is own's and a response through it is attributed to own exactly
// (spec §7.2). The retiring generation still
// accepts until its SIGTERM, which comes after the switch, so a request
// that chose own's pool an instant before the switch is served there. A
// retryable failure is retried inside the usual window while own is still
// current or no generation is; once another one is current, own is gone for
// good and the dial returns genMovedError.
func dialGen(ctx context.Context, d *net.Dialer, src GenSource, host string, own GenInfo) (net.Conn, error) {
	const window = 13 * time.Second
	deadline := time.Now().Add(window)
	for {
		conn, err := unixdial.Checked(ctx, d, own.Socket)
		if err == nil {
			// Spec §7.2 (rev 17): if the generation stops being current
			// before the dial completes, the request is retried against the
			// new current generation. So every new connection is checked
			// once connected: own must still be current. This also covers a
			// foreign row, whose generations share one canonical path, so a
			// connection on it could have reached the replacement. A
			// connection already made (an idle one in own's pool) is not
			// re-checked: the request that chose own's pool sampled own.
			cur, ok := src.ServingGen(host)
			if ok && cur.N == own.N {
				return conn, nil
			}
			conn.Close()
			if ok {
				return nil, &genMovedError{to: cur}
			}
			err = errNoGeneration
		}
		if ctx.Err() != nil || !isRetryableDialErr(err) || time.Now().After(deadline) {
			return nil, err
		}
		if cur, ok := src.ServingGen(host); ok && (cur.N != own.N || cur.Socket != own.Socket) {
			return nil, &genMovedError{to: cur}
		}
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-time.After(100 * time.Millisecond):
		}
	}
}

// genPool is one generation's transport and the generation it reaches.
type genPool struct {
	tr  *http.Transport
	gen GenInfo
}

// GenPools is the per-socket-path set of transports, one idle pool per
// generation socket. One set per consumer (the edge has one, scheduled
// delivery has its own) and shared across route reloads, because the pool
// belongs to the generation, not to the table build that first reached it.
type GenPools struct {
	mu sync.Mutex
	m  map[string]*genPool
}

// NewGenPools returns an empty set.
func NewGenPools() *GenPools { return &GenPools{m: map[string]*genPool{}} }

// pool is generation g's pool, made on first use, or false when g is no
// longer the host's current generation. Its dial reaches g's own socket only
// (dialGen). A foreign row reuses its canonical path for every generation
// (plan 1's reading), so a pool found for the path under another generation
// number is replaced, and its idle connections closed.
//
// g was sampled by the caller some time ago. The check against src happens
// under p.mu, and Retire takes the same lock, so a pool is published only for
// a generation that was still current at that instant: if a switch has
// retired g since, either Retire has already run (and this refuses to publish
// a pool nothing would ever remove) or it runs after the publish and removes
// it. The caller samples again on false.
//
// Lock order: p.mu, then whatever lock src.ServingGen takes (the activator's
// own). That is safe because Retire runs from OnGenRetired hooks, which the
// activator calls outside every lock of its own, and ServingGen never calls
// into this set. A GenSource whose ServingGen reached Retire would deadlock.
func (p *GenPools) pool(host string, g GenInfo, src GenSource) (*genPool, bool) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if cur, ok := src.ServingGen(host); !ok || cur.N != g.N {
		return nil, false
	}
	if gp, ok := p.m[g.Socket]; ok {
		if gp.gen.N == g.N {
			return gp, true
		}
		gp.tr.CloseIdleConnections()
	}
	gp := &genPool{tr: genTransportFor(host, src, &g), gen: g}
	p.m[g.Socket] = gp
	return gp, true
}

// Retire closes socket's idle connections and forgets its pool. It has the
// shape of GenSource.OnGenRetired's hook. A request still in flight on a
// retired pool finishes; its connection is not reused by anyone after.
func (p *GenPools) Retire(_, socket string) {
	p.mu.Lock()
	gp, ok := p.m[socket]
	delete(p.m, socket)
	p.mu.Unlock()
	if ok {
		gp.tr.CloseIdleConnections()
	}
}

// Len is the number of live pools.
func (p *GenPools) Len() int {
	p.mu.Lock()
	defer p.mu.Unlock()
	return len(p.m)
}

// genTransportFor builds the transport behind one pool: unixProxy's idle
// bounds, and a dial that reaches own's socket only (dialGen). RoundTrip
// always passes own: a request that finds no current generation waits for
// one (waitServingGen) and then uses that generation's pool. With own nil
// it is a one-shot transport whose dial resolves the socket per attempt.
func genTransportFor(host string, src GenSource, own *GenInfo) *http.Transport {
	dialer := &net.Dialer{}
	dial := func(ctx context.Context, _, _ string) (net.Conn, error) {
		return dialUnixRetrying(ctx, dialer, func() (string, bool) { return src.AppSocket(host) })
	}
	if own != nil {
		g := *own
		dial = func(ctx context.Context, _, _ string) (net.Conn, error) {
			return dialGen(ctx, dialer, src, host, g)
		}
	}
	return &http.Transport{
		DialContext:       dial,
		MaxIdleConns:      32,
		IdleConnTimeout:   90 * time.Second,
		DisableKeepAlives: own == nil,
	}
}

// genRoundTripper picks the pool of host's current generation for each
// request, stamps the response with the build of the generation whose
// connection carried it, and records that generation on the request's
// outcome for the cache (spec §7.2). A pool's connections are all its own
// generation's (dialGen), so the pool's generation is that generation.
type genRoundTripper struct {
	host  string
	src   GenSource
	pools *GenPools
}

// genRetries bounds how often one request is sent again after its
// generation moved under it.
const genRetries = 2

// genStaleRetries bounds how often one request goes back to choosing a
// generation because the one it chose was no longer current when its pool was
// asked for. That is not a dial attempt, so it does not use genRetries, but a
// source that keeps switching must not hold the request forever.
const genStaleRetries = 3

func (t *genRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	// The generation is chosen, and any switch retried, before a byte of the
	// body is consumed (spec §7.2, rev 17). The transport gets the body
	// behind heldBody: it reads it only once it is writing the request on a
	// connection dialGen has proven to be the chosen generation's, and a
	// Close it makes after a failed dial is held. So a dial that finds the
	// generation moved leaves the body untouched, and the request goes again
	// on the current generation's pool, a streaming POST with no GetBody
	// included.
	var hb *heldBody
	if req.Body != nil && req.Body != http.NoBody {
		hb = &heldBody{rc: req.Body}
		req.Body = hb
	}
	for attempt, stale := 0, 0; ; {
		g, ok := t.src.ServingGen(t.host)
		if !ok {
			// No current generation (between a stop and the next wake): wait
			// for one inside the dial's retry window, then take the normal
			// path, so the response is attributed to the generation that
			// serves it and its connection passes dialGen's check (spec §7.2).
			g, ok = waitServingGen(req.Context(), t.src, t.host)
			if !ok {
				if hb != nil {
					hb.release()
				}
				return nil, errNoGeneration
			}
		}
		gp, ok := t.pools.pool(t.host, g, t.src)
		if !ok {
			// g was retired between the choice and the pool: choose again.
			if stale++; stale > genStaleRetries {
				if hb != nil {
					hb.release()
				}
				if cur, ok := t.src.ServingGen(t.host); ok {
					return nil, &genMovedError{to: cur}
				}
				return nil, errNoGeneration
			}
			continue
		}
		resp, err := gp.tr.RoundTrip(req)
		var moved *genMovedError
		if errors.As(err, &moved) && attempt < genRetries && (hb == nil || !hb.wasRead()) {
			// The generation this request chose stopped being current before
			// its dial completed. Nothing was sent and nothing of the body
			// was read: send it again, on the current one.
			attempt++
			continue
		}
		if err != nil {
			if hb != nil {
				hb.release()
			}
			return nil, err
		}
		// X-Carlos-Version names the build of the generation that served the
		// request, not the registry's Version (spec §7.2). Set, not Add: the
		// platform's voice on this header is the only one.
		if gp.gen.Build != "" {
			resp.Header.Set("X-Carlos-Version", gp.gen.Build)
		}
		noteServedGen(req, t.host, gp.gen.N)
		return resp, nil
	}
}

// waitServingGen waits, within the dial's 13 s retry window and ctx, for
// host to have a current generation, asking every 100 ms.
func waitServingGen(ctx context.Context, src GenSource, host string) (GenInfo, bool) {
	deadline := time.Now().Add(13 * time.Second)
	for {
		if g, ok := src.ServingGen(host); ok {
			return g, true
		}
		if time.Now().After(deadline) {
			return GenInfo{}, false
		}
		select {
		case <-ctx.Done():
			return GenInfo{}, false
		case <-time.After(100 * time.Millisecond):
		}
	}
}

// heldBody is a request body the transport cannot consume or close before
// the request is on a connection. The transport reads a body only while it
// writes the request on a dialled connection, and closes it on any error,
// including a dial that never connected. Until the first Read, that Close
// is held, so the same body can go again on another generation's pool.
// After it, Close passes through as usual.
type heldBody struct {
	rc   io.ReadCloser
	mu   sync.Mutex
	read bool
}

func (b *heldBody) Read(p []byte) (int, error) {
	b.mu.Lock()
	b.read = true
	b.mu.Unlock()
	return b.rc.Read(p)
}

func (b *heldBody) Close() error {
	b.mu.Lock()
	read := b.read
	b.mu.Unlock()
	if !read {
		return nil // a failed dial's close: the body may still be sent
	}
	return b.rc.Close()
}

func (b *heldBody) wasRead() bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.read
}

// release closes the body for good: the request has finished with an error.
func (b *heldBody) release() { _ = b.rc.Close() }

// appSocketProxy is unixProxy for an activator-managed route with
// generations on: the same Rewrite, Host preservation, upgrade pass-through
// and error handler, over genRoundTripper instead of one fixed socket.
func appSocketProxy(host string, src GenSource, pools *GenPools) *httputil.ReverseProxy {
	return &httputil.ReverseProxy{
		Rewrite: func(pr *httputil.ProxyRequest) {
			setForwarded(pr)
			pr.Out.URL.Scheme = "http"
			pr.Out.URL.Host = "carlos" // placeholder; the unix dialer ignores it
			pr.Out.Host = pr.In.Host   // backend sees the real hostname
			scrubInboundPurgeHeader(pr)
		},
		Transport:    &genRoundTripper{host: host, src: src, pools: pools},
		ErrorHandler: defaultProxyError(proxyCause),
	}
}

// AppSocketProxy is appSocketProxy for callers outside this package: the
// scheduled-work deliverer in internal/host, which must reach the current
// generation the same way the edge does (spec §7.1) without the edge's
// cache, compression or usage layers.
func AppSocketProxy(host string, src GenSource, pools *GenPools) http.Handler {
	return appSocketProxy(host, src, pools)
}