package executor

import (
	"context"
	"errors"
	"sync"

	"github.com/google/uuid"

	millproto "tangled.org/core/spindle/mill/proto"
)

type millAnswer interface {
	comparable
	GetRequestId() string
	GetError() string
}

// a disconnect fails every pending ask because answers never cross to a new session
type millAsker[R millAnswer] struct {
	mu      sync.Mutex
	pending map[string]chan R
	sendFn  func(*millproto.Message) error
}

func newMillAsker[R millAnswer]() *millAsker[R] {
	return &millAsker[R]{pending: make(map[string]chan R)}
}

func (a *millAsker[R]) SetSendFn(sendFn func(*millproto.Message) error) {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.sendFn = sendFn
}

func (a *millAsker[R]) OnDisconnect() {
	a.mu.Lock()
	old := a.pending
	a.pending = make(map[string]chan R)
	a.sendFn = nil
	a.mu.Unlock()

	var none R
	for _, ch := range old {
		select {
		case ch <- none:
		default:
		}
	}
}

func (a *millAsker[R]) HandleResponse(resp R) {
	a.mu.Lock()
	ch, ok := a.pending[resp.GetRequestId()]
	a.mu.Unlock()
	if ok {
		select {
		case ch <- resp:
		default:
		}
	}
}

func (a *millAsker[R]) ask(ctx context.Context, request func(requestID string) *millproto.Message) (R, error) {
	var none R
	reqID := uuid.NewString()
	ch := make(chan R, 1)

	a.mu.Lock()
	send := a.sendFn
	if send == nil {
		a.mu.Unlock()
		return none, errors.New("not connected to the mill")
	}
	a.pending[reqID] = ch
	a.mu.Unlock()

	defer func() {
		a.mu.Lock()
		delete(a.pending, reqID)
		a.mu.Unlock()
	}()

	if err := send(request(reqID)); err != nil {
		return none, err
	}
	select {
	case <-ctx.Done():
		return none, ctx.Err()
	case resp := <-ch:
		if resp == none {
			return none, errors.New("disconnected from the mill")
		}
		if resp.GetError() != "" {
			return none, errors.New(resp.GetError())
		}
		return resp, nil
	}
}