package wire

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

	"github.com/gorilla/websocket"
)

const elasticRequestLinkWriteWait = 10 * time.Second

type elasticRequestLink struct {
	linkID           string
	conn             *websocket.Conn
	allowedEndpoints map[string]struct{}

	cancel context.CancelFunc
	done   chan struct{}

	writeMu       sync.Mutex
	mu            sync.Mutex
	closed        bool
	busy          bool
	draining      bool
	currentID     string
	currentToken  string
	currentCancel context.CancelFunc
}

func (link *elasticRequestLink) tryStart(requestID string) bool {
	link.mu.Lock()
	defer link.mu.Unlock()
	if link.closed || link.busy || link.draining {
		return false
	}
	link.busy = true
	link.currentID = strings.TrimSpace(requestID)
	return true
}

func (link *elasticRequestLink) allowsEndpoint(endpointID string, version int) bool {
	if link == nil || len(link.allowedEndpoints) == 0 {
		return false
	}
	_, ok := link.allowedEndpoints[desktopRemoteEndpointKey(endpointID, version)]
	return ok
}

func (link *elasticRequestLink) setCurrentCancel(cancelToken string, cancel context.CancelFunc) {
	link.mu.Lock()
	link.currentToken = strings.TrimSpace(cancelToken)
	link.currentCancel = cancel
	link.mu.Unlock()
}

func (link *elasticRequestLink) finish() {
	link.mu.Lock()
	shouldClose := link.draining
	link.busy = false
	link.currentID = ""
	link.currentToken = ""
	link.currentCancel = nil
	link.mu.Unlock()
	if shouldClose {
		link.close()
	}
}

func (link *elasticRequestLink) isBusy() bool {
	link.mu.Lock()
	defer link.mu.Unlock()
	return link.busy
}

func (link *elasticRequestLink) requestDrain() bool {
	link.mu.Lock()
	defer link.mu.Unlock()
	link.draining = true
	return !link.busy
}

// cancelAssignment requires the per-request cancel token from the assignment;
// matching only the link ID or request ID is not enough to cancel local work.
func (link *elasticRequestLink) cancelAssignment(requestID string, cancelToken string) {
	link.mu.Lock()
	cancel := link.currentCancel
	currentID := strings.TrimSpace(link.currentID)
	currentToken := strings.TrimSpace(link.currentToken)
	link.mu.Unlock()
	requestID = strings.TrimSpace(requestID)
	cancelToken = strings.TrimSpace(cancelToken)
	if cancel != nil && (requestID == "" || requestID == currentID) && cancelToken != "" && cancelToken == currentToken {
		cancel()
	}
}

func (link *elasticRequestLink) sendJSON(message any) error {
	link.mu.Lock()
	if link.closed {
		link.mu.Unlock()
		return errors.New("elastic request link closed")
	}
	conn := link.conn
	link.mu.Unlock()
	if conn == nil {
		return errors.New("elastic request link has no carrier")
	}
	link.writeMu.Lock()
	defer link.writeMu.Unlock()
	_ = conn.SetWriteDeadline(time.Now().Add(elasticRequestLinkWriteWait))
	return conn.WriteJSON(message)
}

func (link *elasticRequestLink) close() {
	link.mu.Lock()
	if link.closed {
		link.mu.Unlock()
		return
	}
	link.closed = true
	cancel := link.cancel
	conn := link.conn
	link.mu.Unlock()
	if cancel != nil {
		cancel()
	}
	if conn != nil {
		link.writeMu.Lock()
		_ = conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(elasticRequestLinkWriteWait))
		link.writeMu.Unlock()
		_ = conn.Close()
	}
}