package executor

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

	"github.com/google/uuid"

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

// shorter than the debug listener's own auth timeout, so a slow mill refuses instead of stalling the login
const debugKeyAskTimeout = 4 * time.Second

// asks the mill whether a key may open a debug shell. the mill resolves the repo from the lease, so
// the executor can't ask about a repo it wasn't given
type DebugKeyClient struct {
	mu      sync.Mutex
	pending map[string]chan *millv1.DebugKeyResponse
	sendFn  func(*millproto.Message) error
}

func NewDebugKeyClient() *DebugKeyClient {
	return &DebugKeyClient{pending: make(map[string]chan *millv1.DebugKeyResponse)}
}

func (c *DebugKeyClient) SetSendFn(sendFn func(*millproto.Message) error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.sendFn = sendFn
}

func (c *DebugKeyClient) OnDisconnect() {
	c.mu.Lock()
	old := c.pending
	c.pending = make(map[string]chan *millv1.DebugKeyResponse)
	c.sendFn = nil
	c.mu.Unlock()

	for _, ch := range old {
		select {
		case ch <- nil:
		default:
		}
	}
}

func (c *DebugKeyClient) HandleResponse(resp *millv1.DebugKeyResponse) {
	c.mu.Lock()
	ch, ok := c.pending[resp.GetRequestId()]
	c.mu.Unlock()
	if ok {
		select {
		case ch <- resp:
		default:
		}
	}
}

// fails closed, any error or silence from the mill refuses the key
func (c *DebugKeyClient) ForLease(leaseID string) func(ctx context.Context, key []byte) bool {
	return func(ctx context.Context, key []byte) bool {
		ctx, cancel := context.WithTimeout(ctx, debugKeyAskTimeout)
		defer cancel()
		resp, err := c.ask(ctx, &millv1.DebugKeyRequest{LeaseId: leaseID, PublicKey: key})
		return err == nil && resp.GetAllowed()
	}
}

func (c *DebugKeyClient) ask(ctx context.Context, req *millv1.DebugKeyRequest) (*millv1.DebugKeyResponse, error) {
	reqID := uuid.NewString()
	req.RequestId = reqID
	ch := make(chan *millv1.DebugKeyResponse, 1)

	c.mu.Lock()
	send := c.sendFn
	if send == nil {
		c.mu.Unlock()
		return nil, errors.New("debug key client not connected")
	}
	c.pending[reqID] = ch
	c.mu.Unlock()

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

	if err := send(&millproto.Message{DebugKeyReq: req}); err != nil {
		return nil, err
	}
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case resp := <-ch:
		if resp == nil {
			return nil, errors.New("debug key client disconnected")
		}
		if resp.GetError() != "" {
			return nil, errors.New(resp.GetError())
		}
		return resp, nil
	}
}