package opt

import (
	"gitlink.org.cn/cloudream/common/utils/math2"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch/dag"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch2/ops2"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch2/parser/state"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/factory"
)

// 将SegmentJoin指令替换成分片上传指令
func UseMultipartUploadToShardStore(ctx *state.GenerateState) {
	dag.WalkOnlyType[*ops2.SegmentJoinNode](ctx.DAG.Graph, func(joinNode *ops2.SegmentJoinNode) bool {
		if joinNode.Joined().Dst.Len() != 1 {
			return true
		}

		joinDst := joinNode.Joined().Dst.Get(0)
		bwNode, ok := joinDst.(*ops2.BaseWriteNode)
		if !ok {
			return true
		}

		// SegmentJoin的输出流的范围必须与ToShardStore的输入流的范围相同，
		// 虽然可以通过调整SegmentJoin的输入流来调整范围，但太复杂，暂不支持
		toStrIdx := bwNode.GetTo().GetStreamIndex()
		toStrRng := bwNode.GetTo().GetRange()
		if toStrIdx.IsRaw() {
			if !toStrRng.Equals(ctx.StreamRange) {
				return true
			}
		} else {
			return true
		}

		// Join的目的地必须支持MultipartUpload功能才能替换成分片上传
		multiUpload, err := factory.GetBuilder(&bwNode.UserSpace).CreateMultiparter(true)
		if err != nil {
			return true
		}

		// Join的每一个段的大小必须超过最小分片大小。
		// 目前只支持拆分超过最大分片的流，不支持合并多个小段流以达到最小分片大小。
		for _, size := range joinNode.Segments {
			if size < multiUpload.MinPartSize() {
				return true
			}
		}

		initNode := ctx.DAG.NewMultipartInitiator(bwNode.UserSpace)
		initNode.Env().CopyFrom(bwNode.Env())

		partNumber := 1
		for i, size := range joinNode.Segments {
			joinInput := joinNode.InputSlot(i)

			if size > multiUpload.MaxPartSize() {
				// 如果一个分段的大小大于最大分片大小，则需要拆分为多个小段上传
				// 拆分以及上传指令直接在流的产生节点执行
				splits := math2.SplitLessThan(size, multiUpload.MaxPartSize())
				splitNode := ctx.DAG.NewSegmentSplit(splits)
				splitNode.Env().CopyFrom(joinInput.Var().Src.Env())

				joinInput.Var().ToSlot(splitNode.InputSlot())

				for i2 := 0; i2 < len(splits); i2++ {
					uploadNode := ctx.DAG.NewMultipartUpload(bwNode.UserSpace, partNumber, splits[i2])
					uploadNode.Env().CopyFrom(joinInput.Var().Src.Env())

					initNode.UploadArgsVar().ToSlot(uploadNode.UploadArgsSlot())
					splitNode.SegmentVar(i2).ToSlot(uploadNode.PartStreamSlot())
					uploadNode.UploadResultVar().ToSlot(initNode.PartInfoSlot())

					partNumber++
				}
			} else {
				// 否则直接上传整个分段
				uploadNode := ctx.DAG.NewMultipartUpload(bwNode.UserSpace, partNumber, size)
				// 上传指令直接在流的产生节点执行
				uploadNode.Env().CopyFrom(joinInput.Var().Src.Env())

				initNode.UploadArgsVar().ToSlot(uploadNode.UploadArgsSlot())
				joinInput.Var().ToSlot(uploadNode.PartStreamSlot())
				uploadNode.UploadResultVar().ToSlot(initNode.PartInfoSlot())

				partNumber++
			}

			joinInput.Var().NotTo(joinNode)
		}

		// BaseWriteNode的FileInfoVar替换为MultipartInitiatorNode的FileInfoVar
		for _, dstSlot := range bwNode.FileInfoVar().ListDstSlots() {
			initNode.FileInfoVar().ToSlot(dstSlot)
		}

		// 最后删除Join指令和ToShardStore指令
		ctx.DAG.RemoveNode(joinNode)
		ctx.DAG.RemoveNode(bwNode)
		return true
	})
}