package dag

import (
	"context"
	"fmt"
	"sync"
)

type dag[T any] struct {
	nodes       []DAGNode[T]
	nodeToIndex map[DAGNode[T]]int

	// edges[to][] = from1, from2...
	edges [][]int
}

func NewDAG[T any]() DAG[T] {
	return &dag[T]{
		nodeToIndex: map[DAGNode[T]]int{},
	}
}

func (g *dag[T]) AddEdge(from, to DAGNode[T]) error {
	fromIndex, ok := g.addNode(from)
	toIndex, ok2 := g.addNode(to)
	// from node and to node have been added already
	if ok && ok2 {
		// detect cycle
		if g.isConnected(toIndex, fromIndex) {
			return fmt.Errorf("cycle is detected, from %s to %s", from, to)
		}

		// check whether edge has been added
		for _, addedFromIndex := range g.edges[toIndex] {
			if addedFromIndex == fromIndex {
				// edge has been added
				return nil
			}
		}
	}

	if fromIndex == -1 || toIndex == -1 {
		return nil
	}

	g.edges[toIndex] = append(g.edges[toIndex], fromIndex)

	return nil
}

func (g *dag[T]) addNode(n DAGNode[T]) (int, bool) {
	if n == nil {
		return -1, false
	}
	index, ok := g.nodeToIndex[n]
	if !ok {
		g.nodes = append(g.nodes, n)
		g.edges = append(g.edges, []int{})
		index = len(g.nodes) - 1
		g.nodeToIndex[n] = index
	}

	return index, ok
}

func (g *dag[T]) isConnected(from, to int) bool {
	for _, parent := range g.edges[from] {
		if g.isConnected(parent, to) {
			return true
		}
	}
	return false
}

func (g *dag[T]) Run(ctx context.Context) {
	wg := sync.WaitGroup{}

	readyChan := make([]chan struct{}, len(g.nodes))
	resArr := make([]T, len(g.nodes))
	for index := range g.nodes {
		readyChan[index] = make(chan struct{})
	}
	for index, node := range g.nodes {
		wg.Add(1)
		go func(index int, node DAGNode[T]) {
			defer func() {
				// notify all children
				close(readyChan[index])

				wg.Done()

			}()

			// parent result
			pres := []T{}
			nctx := ctx
			// wait all parents are ready
			for _, parent := range g.edges[index] {
				<-readyChan[parent]
				parentNode := g.nodes[parent]
				if err := parentNode.Error(); err != nil {
					node.Skip()
					return
				}
				if wc, ok := parentNode.(WithContext); ok {
					nctx = wc.WithContext(nctx)
				}
				pres = append(pres, resArr[parent])
			}
			nres := node.Run(nctx, pres...)
			resArr[index] = nres

		}(index, node)
	}
	wg.Wait()
}