package api

import (
	"context"
	"errors"
	"io"
	"sync"
	"time"
)

// defaultStreamBodyIdleTimeout bounds how long a streaming response
// body may go without producing bytes before the gateway cancels
// the upstream request and gives up. It is intentionally larger
// than the control-plane total timeout so SSE streams that emit
// keepalives every ~30s still pass, while a wedged upstream does
// not pin a Desktop handler forever.
const defaultStreamBodyIdleTimeout = 60 * time.Second

// errStreamBodyIdleTimeout is surfaced by streamBodyReader.Read
// after the watchdog cancels the upstream request context. The
// underlying source Read returns a context.Canceled error first;
// the wrapper translates the post-cancel Read result into this
// sentinel so callers can distinguish idle-timeout from other
// errors if needed.
var errStreamBodyIdleTimeout = errors.New("stream body idle timeout")

// errStreamBodyClosed is returned by streamBodyReader.Read when the
// reader has been closed externally (gateway teardown, generation
// swap, etc.) before any bytes were observed.
var errStreamBodyClosed = errors.New("stream body closed")

// streamBodyConfig tunes a streamBodyReader instance.
type streamBodyConfig struct {
	// IdleTimeout is the rolling per-Read deadline. Zero falls back
	// to defaultStreamBodyIdleTimeout so callers cannot accidentally
	// disable the watchdog.
	IdleTimeout time.Duration
	// Cancel, if non-nil, is invoked exactly once when the idle
	// timeout fires OR when Close is called. It must cancel the
	// upstream HTTP request context so the standard *http.Transport
	// tears down the underlying connection and any in-flight
	// source.Read returns with a context error.
	//
	// Custom RoundTrippers that do NOT honor request context
	// cancellation cannot be supported: source.Read would block
	// forever and the watchdog has no way to interrupt it. Such
	// transports must opt out of streaming proxying entirely.
	Cancel context.CancelFunc
	// OnTimeout, if non-nil, is invoked exactly once when the idle
	// timeout fires (NOT on plain Close). Use it for metrics / log
	// emission; it must not block.
	OnTimeout func()
	// NewTimer is a test seam for driving the watchdog without wall-clock
	// sleeps. Production callers leave it nil and use time.NewTimer.
	NewTimer func(time.Duration) streamBodyTimer
}

type streamBodyTimer interface {
	C() <-chan time.Time
	Stop() bool
	Reset(time.Duration) bool
}

type realStreamBodyTimer struct {
	timer *time.Timer
}

func (t *realStreamBodyTimer) C() <-chan time.Time        { return t.timer.C }
func (t *realStreamBodyTimer) Stop() bool                 { return t.timer.Stop() }
func (t *realStreamBodyTimer) Reset(d time.Duration) bool { return t.timer.Reset(d) }

// streamBodyReader wraps an upstream response body with a rolling
// idle watchdog. Each successful Read resets the deadline; if no
// bytes arrive within IdleTimeout, the watchdog cancels the
// upstream request context so the standard *http.Transport
// interrupts the blocked Read.
//
// Read is synchronous: no detached goroutine ever writes to the
// caller's buffer. The only background goroutine is the watchdog,
// whose sole escape hatch is the caller-supplied Cancel hook.
type streamBodyReader struct {
	source      io.ReadCloser
	idleTimeout time.Duration
	onTimeout   func()
	cancel      context.CancelFunc

	timer     streamBodyTimer
	closeOnce sync.Once
	closed    chan struct{}
}

func newStreamBodyReader(source io.ReadCloser, cfg streamBodyConfig) *streamBodyReader {
	return newStreamBodyReaderWithContext(context.Background(), source, cfg)
}

func newStreamBodyReaderWithContext(ctx context.Context, source io.ReadCloser, cfg streamBodyConfig) *streamBodyReader {
	if cfg.IdleTimeout <= 0 {
		cfg.IdleTimeout = defaultStreamBodyIdleTimeout
	}
	if cfg.NewTimer == nil {
		cfg.NewTimer = func(d time.Duration) streamBodyTimer {
			return &realStreamBodyTimer{timer: time.NewTimer(d)}
		}
	}
	r := &streamBodyReader{
		source:      source,
		idleTimeout: cfg.IdleTimeout,
		onTimeout:   cfg.OnTimeout,
		cancel:      cfg.Cancel,
		timer:       cfg.NewTimer(cfg.IdleTimeout),
		closed:      make(chan struct{}),
	}
	go r.watchdog()
	return r
}

// watchdog runs in a single goroutine per reader. It waits for the
// idle timer to fire and, on fire, triggers the close path with the
// idle-timeout sentinel. It does NOT touch the caller's buffer; the
// only side effect is invoking Cancel and OnTimeout.
func (r *streamBodyReader) watchdog() {
	select {
	case <-r.timer.C():
		r.fire(true)
	case <-r.closed:
	}
}

// fire performs the close-once sequence. idleTimeout=true records
// the timeout sentinel so a subsequent Read returns
// errStreamBodyIdleTimeout instead of errStreamBodyClosed.
func (r *streamBodyReader) fire(idleTimeout bool) {
	r.closeOnce.Do(func() {
		close(r.closed)
		if r.cancel != nil {
			r.cancel()
		}
		if idleTimeout && r.onTimeout != nil {
			r.onTimeout()
		}
		_ = r.source.Close()
	})
}

// Read performs a synchronous source.Read. The caller owns `p`
// throughout; no detached goroutine writes to it.
//
// Any non-nil source error (EOF, sentinel, etc.) terminates the
// reader: Close is invoked once to stop the watchdog and release the
// source, and (n, err) is returned verbatim. io.Copy / io.ReadAll
// honor the io.Reader contract by processing n>0 bytes before
// considering err, so a trailing-EOF-with-data still delivers the
// full payload to the client without losing the terminal error.
//
// The rolling idle deadline is only extended when source produces
// bytes (n > 0 && err == nil). A source that returns (0, nil) — which
// net/http can emit between chunks of a chunked body — must NOT
// extend the window, otherwise a wedged source that spins nil reads
// could starve the watchdog indefinitely.
func (r *streamBodyReader) Read(p []byte) (int, error) {
	select {
	case <-r.closed:
		return 0, errStreamBodyClosed
	default:
	}

	n, err := r.source.Read(p)
	if err != nil {
		r.Close()
		return n, err
	}
	if n > 0 {
		r.resetTimer()
	}
	return n, nil
}

func (r *streamBodyReader) Close() error {
	r.fire(false)
	return nil
}

// resetTimer restarts the rolling idle window after a byte-producing
// Read. The standard library Timer.Stop + drain pattern is safe here
// because resetTimer is only called from Read (which is logically
// serialized by io.Copy / the caller) and from the constructor.
func (r *streamBodyReader) resetTimer() {
	if !r.timer.Stop() {
		select {
		case <-r.timer.C():
		default:
		}
	}
	r.timer.Reset(r.idleTimeout)
}