package inlinedtree

import (
	"container/heap"

	"bonanza.build/pkg/ds"
	model_core "bonanza.build/pkg/model/core"
	"bonanza.build/pkg/model/encoding"
	model_filesystem_pb "bonanza.build/pkg/proto/model/filesystem"
	"bonanza.build/pkg/storage/object"

	"github.com/buildbarn/bb-storage/pkg/util"

	"google.golang.org/protobuf/proto"
)

// ParentAppender is invoked to append the data to be stored in the
// parent message. If externalObject is set, the data is about to
// be stored in an external message. If externalObject is not set, the
// data may be inlined.
//
// This function may be invoked multiple times to determine how the size
// of the parent object is affected by either inlining the data or
// storing it externally.
type ParentAppender[TParentMessage any, TMetadata model_core.ReferenceMetadata] func(
	parent model_core.PatchedMessage[TParentMessage, TMetadata],
	externalObject *model_core.Decodable[model_core.CreatedObject[TMetadata]],
) error

// Candidate of data that needs to be stored in a message, but can
// optionally be stored in a child object if no space is present.
type Candidate[TParentMessage any, TMetadata model_core.ReferenceMetadata] struct {
	// Message to store in a child object if no space is present to
	// inline it into the parent object.
	ExternalMessage model_core.PatchedMessage[model_core.Marshalable, TMetadata]
	// Encoder to use when storing ExternalMessage in a child object.
	Encoder encoding.BinaryEncoder
	// Function to invoke to either inline the message into the
	// output, or create a reference to the child object.
	ParentAppender ParentAppender[TParentMessage, TMetadata]
}

// AlwaysInline creates a simple Candidate that always inlines data into
// the parent message. This can be used to set message fields that have
// no way of getting stored in an external object.
func AlwaysInline[TParentMessage any, TMetadata model_core.ReferenceMetadata](
	patcher *model_core.ReferenceMessagePatcher[TMetadata],
	parentAppender func(parent model_core.PatchedMessage[TParentMessage, TMetadata]),
) Candidate[TParentMessage, TMetadata] {
	return Candidate[TParentMessage, TMetadata]{
		ExternalMessage: model_core.NewPatchedMessage((model_core.Marshalable)(nil), patcher),
		ParentAppender: func(
			out model_core.PatchedMessage[TParentMessage, TMetadata],
			externalObject *model_core.Decodable[model_core.CreatedObject[TMetadata]],
		) error {
			parentAppender(out)
			return nil
		},
	}
}

var marshalOptions = proto.MarshalOptions{UseCachedSize: true}

type queuedCandidate struct {
	candidateIndex           int
	inlinedSizeIncreaseBytes int
}

type queuedCandidates struct {
	ds.Slice[queuedCandidate]
}

func (l queuedCandidates) Less(i, j int) bool {
	return l.Slice[i].inlinedSizeIncreaseBytes < l.Slice[j].inlinedSizeIncreaseBytes
}

// CandidateList is a list of candidates that may be inlined into message
// generated by Build, or stored separately.
type CandidateList[TParentMessage any, TMetadata model_core.ReferenceMetadata] []Candidate[TParentMessage, TMetadata]

// Discard any state associated with all candidates. This ensures that
// if a call to Build() fails, any files/resources owned by the metadata
// gets released.
func (cl *CandidateList[TParentMessage, TMetadata]) Discard() {
	for i := range *cl {
		(*cl)[i].ExternalMessage.Discard()
	}
	*cl = nil
}

// Options used by the Build function.
type Options struct {
	ReferenceFormat  object.ReferenceFormat
	MaximumSizeBytes int
}

// Build a Protobuf message that contains data, either inlined into the
// message directly or stored in the form of a reference to a separate
// object.
//
// Data to be inlined needs to be provided in the form of candidates.
// For each candidate, a function needs to be provided that can either
// inline the data into the message or emit a reference to it. These
// functions are eventually applied in the same order in which they are
// provided, meaning that the order in which items are appended to lists
// is preserved.
//
// This function selects candidates to inline based on their height or
// size. It initially attempts to keep the height of the resulting
// message minimal, followed by inlining as many candidates as possible
// to fill up the remaining space.
func Build[
	TParentMessage any,
	TMetadata model_core.ReferenceMetadata,
	TParentMessagePtr interface {
		*TParentMessage
		proto.Message
	},
](
	candidates CandidateList[TParentMessagePtr, TMetadata],
	options *Options,
) (model_core.PatchedMessage[TParentMessagePtr, TMetadata], error) {
	defer candidates.Discard()

	// It is often the case that everything can be inlined. First
	// determine the size of the resulting message if all candidates
	// were to be inlined.
	candidatesInlinedSizeBytes := make([]int, 0, len(candidates))
	everythingInlinedSizeBytes := 0
	for _, candidate := range candidates {
		parentInlined := model_core.PatchedMessage[TParentMessagePtr, TMetadata]{
			Message: TParentMessagePtr(new(TParentMessage)),
		}
		if err := candidate.ParentAppender(parentInlined, nil); err != nil {
			return model_core.PatchedMessage[TParentMessagePtr, TMetadata]{}, err
		}
		inlinedSizeBytes := candidate.ExternalMessage.Patcher.GetReferencesSizeBytes() + marshalOptions.Size(parentInlined.Message)
		candidatesInlinedSizeBytes = append(candidatesInlinedSizeBytes, inlinedSizeBytes)
		everythingInlinedSizeBytes += inlinedSizeBytes
	}

	output := model_core.NewSimplePatchedMessage[TMetadata](TParentMessagePtr(new(TParentMessage)))
	defer output.Discard()
	candidatesToInline := make([]bool, len(candidates))
	if everythingInlinedSizeBytes <= options.MaximumSizeBytes {
		// All candidates can be inlined. This means we don't
		// need to do any additional processing.
		for i, candidate := range candidates {
			candidatesToInline[i] = true
			output.Patcher.Merge(candidate.ExternalMessage.Patcher)
		}
	} else {
		// We can't fit all candidates. For each candidate,
		// compute how much the size of the output message
		// increases when inlined, compared to storing it
		// externally.
		//
		// TODO: Prevent these from getting recomputed every time.
		bogusData := []byte("A")
		bogusContents, err := options.ReferenceFormat.NewContents(nil, bogusData)
		if err != nil {
			panic(err)
		}
		bogusCreatedObject := model_core.CreatedObject[TMetadata]{
			Contents: bogusContents,
		}

		outputSizeBytes := 0
		queuedCandidates := queuedCandidates{
			Slice: make(ds.Slice[queuedCandidate], 0, len(candidates)),
		}
		for i, candidate := range candidates {
			// Determine how much space this candidate uses when
			// stored externally. The caller can set
			// ExternalMessage.Message to nil to indicate that
			// inlining should always be performed, meaning this can
			// be skipped.
			inlinedSizeBytes := candidatesInlinedSizeBytes[i]
			if candidate.ExternalMessage.Message != nil {
				parentExternal := model_core.NewSimplePatchedMessage[TMetadata](TParentMessagePtr(new(TParentMessage)))
				decodableBogusCreatedObject, err := model_core.NewDecodable(
					bogusCreatedObject,
					make([]byte, candidate.Encoder.GetDecodingParametersSizeBytes()),
				)
				if err != nil {
					return output, err
				}
				if err := candidate.ParentAppender(parentExternal, &decodableBogusCreatedObject); err != nil {
					parentExternal.Discard()
					return model_core.PatchedMessage[TParentMessagePtr, TMetadata]{}, err
				}
				externalSizeBytes := parentExternal.Patcher.GetReferencesSizeBytes() + marshalOptions.Size(parentExternal.Message)
				parentExternal.Discard()

				if inlinedSizeBytes > externalSizeBytes {
					// Inlining the data takes up more space.
					// Queue it, so that we can inline it only
					// if enough space is available.
					outputSizeBytes += externalSizeBytes
					queuedCandidates.Push(queuedCandidate{
						candidateIndex:           i,
						inlinedSizeIncreaseBytes: inlinedSizeBytes - externalSizeBytes,
					})
					continue
				}
			}

			// Inlining the data is smaller than storing it
			// externally. Inline it immediately.
			candidatesToInline[i] = true
			output.Patcher.Merge(candidate.ExternalMessage.Patcher)
			outputSizeBytes += inlinedSizeBytes
		}

		// Determine how much space is needed to inline all candidates
		// that would otherwise cause the height of the object tree to
		// increase. If space is available, inline them all. This
		// ensures that the height is minimized.
		heightAllInlined := output.Patcher.GetHeight()
		highestInlinedSizeIncreaseBytes := 0
		for _, queuedCandidate := range queuedCandidates.Slice {
			candidate := &candidates[queuedCandidate.candidateIndex]
			if h := candidate.ExternalMessage.Patcher.GetHeight(); heightAllInlined == h {
				highestInlinedSizeIncreaseBytes += queuedCandidate.inlinedSizeIncreaseBytes
			} else if heightAllInlined < h {
				heightAllInlined = h
				highestInlinedSizeIncreaseBytes = queuedCandidate.inlinedSizeIncreaseBytes
			}
		}

		if newOutputSizeBytes := outputSizeBytes + highestInlinedSizeIncreaseBytes; newOutputSizeBytes <= options.MaximumSizeBytes {
			for i := 0; i < queuedCandidates.Len(); {
				queuedCandidate := &queuedCandidates.Slice[i]
				candidate := &candidates[queuedCandidate.candidateIndex]
				if candidate.ExternalMessage.Patcher.GetHeight() == heightAllInlined {
					candidatesToInline[queuedCandidate.candidateIndex] = true
					output.Patcher.Merge(candidate.ExternalMessage.Patcher)

					queuedCandidates.Swap(i, queuedCandidates.Len()-1)
					queuedCandidates.Pop()
				} else {
					i++
				}
			}
			outputSizeBytes = newOutputSizeBytes
		}

		// Fill the remaining space with as many candidates as possible,
		// inlining the smallest ones first. Even though this is
		// detrimental to the overall height, it reduces the total
		// number of objects significantly.
		//
		// It may be that multiple candidates share the same size. In
		// that case we could inline the candidate with the lowest
		// index. However, this has the disadvantage that the order in
		// which candidates are provided to this function matters even
		// if they pertain to different fields, making refactoring hard.
		//
		// Work around this by inlining all equally sized candidates
		// collectively. If there is no space to fit all of them, we
		// store them externally and inline less preferential candidates
		// instead.
		heap.Init(&queuedCandidates)
		var queuedCandidatesWithSameSize []queuedCandidate
		for queuedCandidates.Len() > 0 && outputSizeBytes+queuedCandidates.Slice[0].inlinedSizeIncreaseBytes <= options.MaximumSizeBytes {
			queuedCandidatesWithSameSize = append(queuedCandidatesWithSameSize[:0], heap.Pop(&queuedCandidates).(queuedCandidate))
			inlinedSizeIncreaseBytes := queuedCandidatesWithSameSize[0].inlinedSizeIncreaseBytes
			for queuedCandidates.Len() > 0 && queuedCandidates.Slice[0].inlinedSizeIncreaseBytes == inlinedSizeIncreaseBytes {
				queuedCandidatesWithSameSize = append(queuedCandidatesWithSameSize, heap.Pop(&queuedCandidates).(queuedCandidate))
			}

			if newOutputSizeBytes := outputSizeBytes + inlinedSizeIncreaseBytes*len(queuedCandidatesWithSameSize); newOutputSizeBytes <= options.MaximumSizeBytes {
				for _, queuedCandidate := range queuedCandidatesWithSameSize {
					candidate := &candidates[queuedCandidate.candidateIndex]
					candidatesToInline[queuedCandidate.candidateIndex] = true
					output.Patcher.Merge(candidate.ExternalMessage.Patcher)
				}
				outputSizeBytes = newOutputSizeBytes
			}
		}
	}

	for i := range candidates {
		candidate := &candidates[i]
		if candidatesToInline[i] {
			// Inline the message into the parent.
			if err := candidate.ParentAppender(output, nil); err != nil {
				return model_core.PatchedMessage[TParentMessagePtr, TMetadata]{}, err
			}
		} else {
			// Store the message separately, and store a
			// reference in the parent.
			createdObject, err := model_core.MarshalAndEncode(candidate.ExternalMessage.Move(), options.ReferenceFormat, candidate.Encoder)
			if err != nil {
				return model_core.PatchedMessage[TParentMessagePtr, TMetadata]{}, util.StatusWrapf(err, "Failed to create object contents for candidate at index %d", i)
			}
			if err := candidate.ParentAppender(output, &createdObject); err != nil {
				return model_core.PatchedMessage[TParentMessagePtr, TMetadata]{}, err
			}
		}
	}
	return output.Move(), nil
}

// ParentAppenderForTesting is an instantiation of ParentAppender for
// generating mocks to be used by tests.
type ParentAppenderForTesting ParentAppender[*model_filesystem_pb.DirectoryContents, model_core.ReferenceMetadata]
