package agentk

import (
	"context"
	"maps"
	"sync"

	"k8s.io/apimachinery/pkg/util/wait"
)

type errorDetailsTracker struct {
	mx    sync.Mutex
	store map[errorTrackerKey]operationState
	wg    wait.Group
}

type errorTrackerKey struct {
	name      string
	namespace string
}

// operationState indicates the state of an async operation as it is being watched by the errorDetailsTracker.
type operationState struct {
	// version indicates the version of the operation for which the state is being tracked
	version uint64

	// errorDetails contains error details available at the end of an async operation. It will be nil if the operation has yet to finish execution
	errorDetails *ErrorDetails
}

func newErrorDetailsTracker() *errorDetailsTracker {
	return &errorDetailsTracker{store: make(map[errorTrackerKey]operationState)}
}

func (t *errorDetailsTracker) deleteErrorIfVersion(name, namespace string, version uint64) {
	key := errorTrackerKey{
		name:      name,
		namespace: namespace,
	}

	t.mx.Lock()
	defer t.mx.Unlock()

	existingState, exists := t.store[key]

	if exists && existingState.version == version {
		// only delete entries if they correspond to the version passed in the
		// function call. This check prevents cases where entries for a more
		// recent operation(with a higher version) may be overwritten with
		// error details of an older operation(with a lower version)
		delete(t.store, key)
	}
}

// saveErrorIfVersion will record error details only if the provided version
// matches the version in an existing record. In every other case, nothing
// will be written
func (t *errorDetailsTracker) saveErrorIfVersion(name, namespace string, errorDetails *ErrorDetails, version uint64) {
	key := errorTrackerKey{
		name:      name,
		namespace: namespace,
	}

	t.mx.Lock()
	defer t.mx.Unlock()

	existingState, exists := t.store[key]
	if !exists {
		// Do NOT write anything if nothing exists
		return
	}

	if existingState.version != version {
		// this check is added so to prevent overwrite of errors
		// for entries with mismatched versions
		return
	}
	t.store[key] = operationState{
		version:      existingState.version,
		errorDetails: errorDetails,
	}
}

// watchForLatestErrors will watch the provided non-nil channel for error details and record them asynchronously. If multiple watches are created for
// different versions for the same workspace & namespace, only the error details corresponding to the latest version are tracked while entries for earlier
// versions are discarded. If multiple values are published in the channel, only the first error detail will be recorded. The caller is responsible for ensuring that at-most one value is published on the passed channel
func (t *errorDetailsTracker) watchForLatestErrors(ctx context.Context, name, namespace string, version uint64, asyncErrDetails <-chan *ErrorDetails) {
	t.markEntryWithVersion(name, namespace, version)

	t.wg.StartWithContext(ctx, func(ctx context.Context) {
		errDetails := <-asyncErrDetails
		if errDetails != nil {
			// at least one error details was received on the channel, so it must be saved
			t.saveErrorIfVersion(name, namespace, errDetails, version)
		} else {
			// either nil value or nothing was received on the channel and so the
			// entry for provided version can be safely evicted
			t.deleteErrorIfVersion(name, namespace, version)
		}
	})
}

// waitForErrors waits on existing channels being watched to finish publishing error details(if any)
func (t *errorDetailsTracker) waitForErrors() {
	t.wg.Wait()
}

// markEntryWithVersion will create an entry in the store if and only if
// no entry exists for the particular key OR an entry exists with an older version.
// If the version of an existing entry is higher (not equal) than the version passed in the function
// call, then the writes are skipped and nothing is updated
func (t *errorDetailsTracker) markEntryWithVersion(workspace, namespace string, version uint64) {
	key := errorTrackerKey{
		name:      workspace,
		namespace: namespace,
	}

	t.mx.Lock()
	defer t.mx.Unlock()

	existingState, exists := t.store[key]

	if exists && existingState.version > version {
		// writes should be skipped here as an older version MUST not
		// overwrite the entries corresponding to a newer version
		return
	}

	t.store[key] = operationState{
		version:      version,
		errorDetails: nil,
	}
}

func (t *errorDetailsTracker) createSnapshot() map[errorTrackerKey]operationState {
	snapshot := make(map[errorTrackerKey]operationState)

	t.mx.Lock()
	defer t.mx.Unlock()

	maps.Copy(snapshot, t.store)

	return snapshot
}
