package app

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"git.yuntuo.hk/k2safe/go-omni/lifecycle"
)

// Admission and completion share a lock: once stopped, no request can enter
// between observing zero active handlers and releasing their dependencies.
type requestDrain struct {
	mu      sync.Mutex
	stopped bool
	active  map[*requestWork]struct{}
	done    chan struct{}
}

type requestWork struct{ cancel context.CancelFunc }

func (d *requestDrain) begin(ctx context.Context) (context.Context, *requestWork) {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.stopped {
		return ctx, nil
	}
	ctx, cancel := context.WithCancel(ctx)
	w := &requestWork{cancel: cancel}
	if d.active == nil {
		d.active = make(map[*requestWork]struct{})
	}
	d.active[w] = struct{}{}
	return ctx, w
}

func (d *requestDrain) end(w *requestWork) {
	w.cancel()
	d.mu.Lock()
	defer d.mu.Unlock()
	delete(d.active, w)
	if d.stopped && len(d.active) == 0 {
		close(d.done)
	}
}

func (d *requestDrain) stop() {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.stopped {
		return
	}
	d.stopped = true
	if len(d.active) == 0 {
		close(d.done)
	}
}

func (d *requestDrain) cancelAndWait(timeout time.Duration) error {
	d.stop()
	d.mu.Lock()
	for w := range d.active {
		w.cancel()
	}
	d.mu.Unlock()
	timer := time.NewTimer(timeout)
	defer timer.Stop()
	select {
	case <-d.done:
		return nil
	case <-timer.C:
		select {
		case <-d.done:
			return nil
		default:
		}
		return fmt.Errorf("go-omni: HTTP handlers did not finish within cancellation grace %s: %w", timeout,
			errors.Join(context.DeadlineExceeded, lifecycle.ErrResourcesRetained))
	}
}

// watchShutdown stops admission on the early signal without canceling accepted
// work. The transport retains its graceful drain budget before Close cancels.
func (a *App) watchShutdown(ctx context.Context) func() bool {
	return context.AfterFunc(ctx, a.requests.stop)
}