package watch_aggregator //nolint:staticcheck

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"sync"

	"buf.build/go/protovalidate"
	"github.com/coder/websocket"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modshared"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/logz"
	"google.golang.org/protobuf/encoding/protojson"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	k8swatch "k8s.io/apimachinery/pkg/watch"
)

const (
	watchTypeWatch              = "watch"
	watchTypeUnwatch            = "unwatch"
	watchEventTypeWatchEvent    = "event"
	watchEventTypeStopEvent     = "stop"
	watchEventTypeError         = "error"
	watchFailedErrorType        = "WATCH_FAILED"
	watchRequestFailedErrorType = "WATCH_REQUEST_FAILED"
	watchRequestLimit           = 100
)

type watchFactory interface {
	new(watchRequest *WatchRequest, sender watchResponseSender) watchInterface
}

type watchResponseSender interface {
	sendEvent(watchID string, e *k8sWatchEvent)
	sendStop(watchID string)
	sendError(watchID string, errorType string, err error, obj runtime.Object)
}

type watchInterface interface {
	handle(context.Context)
}

type WatchAggregator struct {
	log       *slog.Logger
	api       modshared.API
	validator protovalidate.Validator

	factory watchFactory

	ws         WebSocketInterface
	mu         sync.Mutex
	watches    map[string]context.CancelFunc
	sendCh     chan *watchEvent
	readBuffer bytes.Buffer
}

type watchEvent struct {
	Type    string         `json:"type"`
	WatchID string         `json:"watchId"`
	Event   *k8sWatchEvent `json:"watchEvent,omitempty"`
	Error   *watchError    `json:"error,omitempty"`
}

// NOTE: k8sWatchEvent contains a Kubernetes watch.Event
// The reason for not using Event from k8s.io/apimachinery/pkg/watch
// is that the import aliasing functionality in gomock is not working properly
// and with that it always imports the aforementioned package as "watch"
// conflicting with the watch_aggregator.watch type.
// We don't want to rename our types because of that and simply
// cloning and owning this type isn't a big deal.
// We can get rid of this once at least
// https://github.com/uber-go/mock/issues/166 is fixed.
type k8sWatchEvent struct {
	Type   k8swatch.EventType `json:"type"`
	Object runtime.Object     `json:"object"`
}

type watchError struct {
	Type    string         `json:"type"`
	Message string         `json:"message"`
	Object  runtime.Object `json:"object,omitempty"`
}

func NewWatchAggregator(log *slog.Logger, api modshared.API, v protovalidate.Validator, accept WebSocketAcceptFunc, rt http.RoundTripper) *WatchAggregator {
	factory := &basicWatchFactory{
		rt: rt,
	}
	return newWatchAggregator(factory, log, api, v, accept)
}

func newWatchAggregator(factory watchFactory, log *slog.Logger, api modshared.API, v protovalidate.Validator, accept WebSocketAcceptFunc) *WatchAggregator {
	c, err := accept()
	if err != nil {
		log.Debug("Unable to accept aggregated watch WebSocket request", logz.Error(err))
		return nil
	}

	return &WatchAggregator{
		log:       log,
		api:       api,
		validator: v,
		factory:   factory,
		ws:        c,
		watches:   make(map[string]context.CancelFunc),
		sendCh:    make(chan *watchEvent),
	}
}

func (a *WatchAggregator) Handle(ctx context.Context) {
	var wg wait.Group
	defer wg.Wait()

	defer close(a.sendCh)

	var watchesWg wait.Group
	defer watchesWg.Wait()

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	// Start the write loop
	wg.StartWithContext(ctx, a.writer)

	// Read and command loop
readLoop:
	for {
		wr, err := a.read(ctx)
		if err != nil {
			if !isShutdown(err, ctx) {
				a.sendError("", watchRequestFailedErrorType, fmt.Errorf("unable to read watch request: %w", err), nil)
			}
			break
		}

		if err := a.validator.Validate(wr); err != nil {
			a.sendError(wr.WatchId, watchRequestFailedErrorType, fmt.Errorf("validation of watch request failed: %w", err), nil)
			break
		}

		switch wr.Type {
		case watchTypeWatch:
			watchesWg.StartWithContext(ctx, func(ctx context.Context) { a.watch(ctx, wr) })
		case watchTypeUnwatch:
			a.stopWatch(wr)
		default:
			a.sendError(wr.WatchId, watchRequestFailedErrorType, fmt.Errorf("unknown watch type %q", wr.Type), nil)
			break readLoop
		}
	}
}

func (a *WatchAggregator) writer(ctx context.Context) {
	defer func() {
		// draining the channel
		for event := range a.sendCh {
			a.log.Debug("Discarding event because of ongoing shutdown", logz.WatchID(event.WatchID))
		}
	}()

	for event := range a.sendCh {
		if !a.write(ctx, event) {
			break
		}
	}

	// The connection may have already been closed by the writer or reader,
	// we don't really care.
	// TODO: improve once the WebSocket package supports canceling with a cause.
	_ = a.ws.Close(websocket.StatusNormalClosure, "")
}

func (a *WatchAggregator) write(ctx context.Context, event *watchEvent) (success bool) {
	success = true
	w, err := a.ws.Writer(ctx, websocket.MessageText)
	if err != nil {
		// When KAS is shutting down the ctx is canceled leading to
		// the WebSocket package automatically closing the connection
		// and hence we should immediately stop to write anything to it.
		// We don't need to log anything as those would be false positives.
		if !isShutdown(err, ctx) {
			a.log.Debug("Unable to get WebSocket writer", logz.Error(err), logz.WatchID(event.WatchID))
		}
		return false
	}
	defer func() {
		closeErr := w.Close()
		if closeErr != nil && !isShutdown(closeErr, ctx) {
			a.log.Debug("Unable to close WebSocket writer", logz.Error(closeErr), logz.WatchID(event.WatchID))
			success = false
		}
	}()

	b, err := json.Marshal(event)
	if err != nil {
		success = false
		err = fmt.Errorf("unable to marshal %s event: %w", event.Type, err)
		a.api.HandleProcessingError(ctx, a.log, "", err)
		b, err = json.Marshal(&watchEvent{
			Type:    watchEventTypeError,
			WatchID: event.WatchID,
			Error: &watchError{
				Type:    watchFailedErrorType,
				Message: err.Error(),
			},
		})
		if err != nil {
			a.api.HandleProcessingError(ctx, a.log, "Unable to marshal error message for event marshaling error", err)
			return false
		}
	}

	_, err = w.Write(b)
	if err != nil {
		// When KAS is shutting down the ctx is canceled leading to
		// the WebSocket package automatically closing the connection
		// and hence we should immediately stop to write anything to it.
		// We don't need to log anything as those would be false positives.
		if !isShutdown(err, ctx) {
			a.log.Debug("Unable to write event on WebSocket connection", logz.Error(err), logz.WatchID(event.WatchID))
		}
		return false
	}

	return success
}

func (a *WatchAggregator) read(ctx context.Context) (*WatchRequest, error) {
	mt, r, err := a.ws.Reader(ctx)
	if err != nil {
		return nil, err
	}

	if mt != websocket.MessageText {
		return nil, fmt.Errorf("expected message type 'text', got %d", mt)
	}

	defer a.readBuffer.Reset()
	_, err = a.readBuffer.ReadFrom(r)
	if err != nil {
		return nil, err
	}

	var wr WatchRequest
	err = protojson.Unmarshal(a.readBuffer.Bytes(), &wr)
	if err != nil {
		return nil, fmt.Errorf("failed to unmarshal JSON as a watch request: %w", err)
	}

	return &wr, nil
}

func (a *WatchAggregator) watch(ctx context.Context, wr *WatchRequest) {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	// Creating a registering a new watch
	w, err := func() (watchInterface, error) {
		a.mu.Lock()
		defer a.mu.Unlock()

		if _, ok := a.watches[wr.WatchId]; ok {
			return nil, errors.New("a watch with this id is already registered, stop it before starting a new one with the same id")
		}

		if len(a.watches) >= watchRequestLimit {
			return nil, fmt.Errorf("reached watch request limit of %d watches. Either stop some unused watches or start a new aggregated watch", watchRequestLimit)
		}

		w := a.factory.new(wr, a)

		a.watches[wr.WatchId] = cancel
		return w, nil
	}()

	if err != nil {
		a.sendError(wr.WatchId, watchRequestFailedErrorType, err, nil)
		return
	}

	defer func() {
		a.mu.Lock()
		defer a.mu.Unlock()
		delete(a.watches, wr.WatchId)
	}()

	// Start handling the actual watch
	w.handle(ctx)
}

func (a *WatchAggregator) stopWatch(wr *WatchRequest) {
	a.mu.Lock()
	defer a.mu.Unlock()
	if w, ok := a.watches[wr.WatchId]; ok {
		w()
	}
}

func (a *WatchAggregator) sendEvent(watchID string, e *k8sWatchEvent) {
	a.sendCh <- &watchEvent{
		Type:    watchEventTypeWatchEvent,
		WatchID: watchID,
		Event:   e,
	}
}

func (a *WatchAggregator) sendStop(watchID string) {
	a.sendCh <- &watchEvent{
		Type:    watchEventTypeStopEvent,
		WatchID: watchID,
	}
}

func (a *WatchAggregator) sendError(watchID, errorType string, err error, obj runtime.Object) {
	a.sendCh <- &watchEvent{
		Type:    watchEventTypeError,
		WatchID: watchID,
		Error: &watchError{
			Type:    errorType,
			Message: err.Error(),
			Object:  obj,
		},
	}
}

func isShutdown(err error, ctx context.Context) bool {
	return websocket.CloseStatus(err) != -1 || ctx.Err() != nil
}
