package opt

import (
	"gitlink.org.cn/cloudream/common/pkgs/ioswitch/dag"
	"gitlink.org.cn/cloudream/common/utils/lo2"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch2"
	"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"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/types"
)

// 减少未使用的Multiply指令的输出流。如果减少到0，则删除该指令
func RemoveUnusedMultiplyOutput(ctx *state.GenerateState) bool {
	changed := false
	dag.WalkOnlyType[*ops2.ECMultiplyNode](ctx.DAG.Graph, func(node *ops2.ECMultiplyNode) bool {
		outArr := node.OutputStreams().Slots.RawArray()
		for i2, out := range outArr {
			if out.Dst.Len() > 0 {
				continue
			}

			outArr[i2] = nil
			node.OutputIndexes[i2] = -2
			changed = true
		}

		node.OutputStreams().Slots.SetRawArray(lo2.RemoveAllDefault(outArr))
		node.OutputIndexes = lo2.RemoveAll(node.OutputIndexes, -2)

		// 如果所有输出流都被删除，则删除该指令
		if node.OutputStreams().Len() == 0 {
			node.RemoveAllInputs()
			ctx.DAG.RemoveNode(node)
			changed = true
		}

		return true
	})
	return changed
}

// 替换ECMultiply指令
func UseECMultiplier(ctx *state.GenerateState) {
	dag.WalkOnlyType[*ops2.ECMultiplyNode](ctx.DAG.Graph, func(mulNode *ops2.ECMultiplyNode) bool {
		if mulNode.OutputStreams().Len() == 0 {
			return true
		}

		// 暂不支持编解码流的一部分
		if ctx.StreamRange.Offset != 0 || ctx.StreamRange.Length != nil {
			return true
		}

		var to *ioswitch2.ToShardStore
		var bwNodes []*ops2.BaseWriteNode
		// 所有的输出流必须有且只有一个相同的目的地
		// 暂时只支持分片存储
		for i := 0; i < mulNode.OutputStreams().Len(); i++ {
			if mulNode.OutputStreams().Get(i).Dst.Len() != 1 {
				return true
			}

			dstNode := mulNode.OutputStreams().Get(i).Dst.Get(0)
			swNode, ok := dstNode.(*ops2.BaseWriteNode)
			if !ok {
				return true
			}

			toShard, ok := swNode.To.(*ioswitch2.ToShardStore)
			if !ok {
				return true
			}

			if to == nil {
				to = toShard
			} else if to.UserSpace.UserSpace.UserSpaceID != swNode.UserSpace.UserSpace.UserSpaceID {
				return true
			}
			bwNodes = append(bwNodes, swNode)
		}
		_, err := factory.GetBuilder(&to.UserSpace).CreateECMultiplier(true)
		if err != nil {
			return true
		}

		// 每一个输入流都必须直接来自于存储服务，而且要支持通过HTTP读取文件
		var brNodes []*ops2.BaseReadDynNode
		var fromShards []*ioswitch2.FromShardStore
		for i := 0; i < mulNode.InputStreams().Len(); i++ {
			inNode := mulNode.InputStreams().Get(i).Src
			brNode, ok := inNode.(*ops2.BaseReadDynNode)
			if !ok {
				return true
			}
			fromShard, ok := brNode.From.(*ioswitch2.FromShardStore)
			if !ok {
				return true
			}

			store, err := factory.GetBuilder(&brNode.UserSpace).CreateShardStore(true)
			if err != nil {
				return true
			}
			_, ok = store.(types.HTTPShardRead)
			if !ok {
				return true
			}

			brNodes = append(brNodes, brNode)
			fromShards = append(fromShards, fromShard)
		}

		// 检查满足条件后，替换ECMultiply指令
		callMul := ctx.DAG.NewCallECMultiplier(to.UserSpace)
		if err := setEnvBySpace(callMul, &to.UserSpace); err != nil {
			return true
		}

		callMul.InitFrom(mulNode)
		for i, brNode := range brNodes {
			brNode.Output().Var().NotTo(mulNode)
			// 只有完全没有输出的ShardReadNode才可以被删除
			if brNode.Output().Var().Dst.Len() == 0 {
				ctx.DAG.RemoveNode(brNode)
			}

			hbr := ctx.DAG.NewGetShardHTTPRequest(brNode.UserSpace, fromShards[i].FileHash)
			hbr.Env().ToEnvDriver(true)
			hbr.HTTPRequestVar().ToSlot(callMul.HTTPRequestSlot(i))
		}

		for i, bwNode := range bwNodes {
			ctx.DAG.RemoveNode(bwNode)

			for _, dstSlot := range bwNode.FileInfoVar().ListDstSlots() {
				callMul.FileInfoVar(i).ToSlot(dstSlot)
			}
		}

		ctx.DAG.RemoveNode(mulNode)
		return true
	})
}