package watch_graph //nolint:staticcheck
import (
	"context"
)

// VID is vertex ID
// VD is vertex data
// AT is arc type
// AD is arc data

type ArcToID[VID, AT comparable] struct {
	To      VID
	ArcType AT
}

// ArcSetWithData is a set of arcs with data.
type ArcSetWithData[VID, AT comparable, AD any] map[ArcToID[VID, AT]]AD

// ArcSet is a set of arcs.
type ArcSet[VID, AT comparable] map[ArcToID[VID, AT]]struct{}

type SetVertexResult byte

const (
	// VertexNoop - no action was taken. Vertex exists and the data is the same.
	VertexNoop SetVertexResult = iota
	// VertexAdded - a new vertex was added.
	VertexAdded
	// VertexUpdated - vertex exists but data was updated.
	VertexUpdated
)

type GraphInspector[VID, AT comparable, VD, AD any] interface {
	InboundArcsFor(VID) ArcSet[VID, AT]
	OutboundArcsFor(VID) ArcSetWithData[VID, AT, AD]
	VertexData(VID) (VD, bool)
}

type ObjectGraphObserver[VID, AT comparable, VD, AD any] interface {
	OnSetVertex(context.Context, GraphInspector[VID, AT, VD, AD], VID, VD)
	OnDeleteVertex(context.Context, GraphInspector[VID, AT, VD, AD], VID)
	OnSetArc(ctx context.Context, insp GraphInspector[VID, AT, VD, AD], from VID, to ArcToID[VID, AT], data AD)
	OnDeleteArc(ctx context.Context, insp GraphInspector[VID, AT, VD, AD], from VID, to ArcToID[VID, AT])
}

type ObjectGraphOpts[VID, AT comparable, VD, AD any] struct {
	IsVertexDataEqual func(a, b VD) bool
	IsArcDataEqual    func(a, b AD) bool
	Observer          ObjectGraphObserver[VID, AT, VD, AD]
}

type ObjectGraph[VID, AT comparable, VD, AD any] struct {
	opts ObjectGraphOpts[VID, AT, VD, AD]

	id2data     map[VID]VD
	id2outbound map[VID]ArcSetWithData[VID, AT, AD]
	id2inbound  map[VID]ArcSet[VID, AT]
}

func NewObjectGraph[VID, AT comparable, VD, AD any](opts ObjectGraphOpts[VID, AT, VD, AD]) *ObjectGraph[VID, AT, VD, AD] {
	return &ObjectGraph[VID, AT, VD, AD]{
		opts:        opts,
		id2data:     map[VID]VD{},
		id2outbound: map[VID]ArcSetWithData[VID, AT, AD]{},
		id2inbound:  map[VID]ArcSet[VID, AT]{},
	}
}

// SetVertex sets a vertex in the graph declaratively:
// - replaces the existing vertex or adds a new one with the specified VID.
// - replaces all outbound arcs of that vertex with the provided arcs.
// arcs may be nil.
// NOTE: this function takes ownership of the arcs object. Do not use it after passing to this function.
func (g *ObjectGraph[VID, AT, VD, AD]) SetVertex(ctx context.Context, vid VID, vd VD, arcs ArcSetWithData[VID, AT, AD]) SetVertexResult {
	res := VertexNoop
	vOldData, isSet := g.id2data[vid]
	if isSet { // vertex exists
		// 1. Check if data is equal
		if g.opts.IsVertexDataEqual(vOldData, vd) {
			// Nothing to update, same data
		} else {
			// Update vertex data
			g.id2data[vid] = vd
			g.opts.Observer.OnSetVertex(ctx, g, vid, vd)
			res = VertexUpdated
		}
		// 2. Delete arcs that are no longer in the set
		arcSetForVertex := g.id2outbound[vid] // may be nil
		for aid := range arcSetForVertex {
			if _, ok := arcs[aid]; ok {
				// Arc is still in the set
			} else {
				// Arc is not in the new set, delete
				delete(arcSetForVertex, aid)
				g.deleteInboundArc(vid, aid)
				g.opts.Observer.OnDeleteArc(ctx, g, vid, aid)
			}
		}
		// 3. Add new arcs
		for aid, aData := range arcs {
			existingArcData, ok := arcSetForVertex[aid]
			if ok {
				// Arc exists, let's check the data
				if g.opts.IsArcDataEqual(existingArcData, aData) {
					// Data is equal, nothing to do.
				} else {
					// Data is not equal, arc must be updated.
					arcSetForVertex[aid] = aData
					// no need to call it since the arc already exists and inbound arcs don't store arc data (nothing to update).
					// g.addInboundArc(vid, aid)
					g.opts.Observer.OnSetArc(ctx, g, vid, aid, aData)
				}
			} else {
				// Arc does not exist, add
				if arcSetForVertex == nil { // ensure not nil.
					arcSetForVertex = ArcSetWithData[VID, AT, AD]{}
					g.id2outbound[vid] = arcSetForVertex
				}
				arcSetForVertex[aid] = aData
				g.addInboundArc(vid, aid)
				g.opts.Observer.OnSetArc(ctx, g, vid, aid, aData)
			}
		}
	} else {
		// Vertex does not exist
		// 1. Set vertex
		g.id2data[vid] = vd
		g.opts.Observer.OnSetVertex(ctx, g, vid, vd)
		// 2. Set arcs after vertex
		if len(arcs) > 0 { // Don't want to retain empty maps, let GC have them
			g.id2outbound[vid] = arcs
			for aid, aData := range arcs {
				g.addInboundArc(vid, aid)
				g.opts.Observer.OnSetArc(ctx, g, vid, aid, aData)
			}
		}
		res = VertexAdded
	}
	return res
}

// DeleteVertex deletes the vertex with the specified VID.
// All outbound arcs are also deleted.
func (g *ObjectGraph[VID, AT, VD, AD]) DeleteVertex(ctx context.Context, vid VID) bool {
	// 0. Check if vertex exists
	if _, ok := g.id2data[vid]; !ok {
		// Does not exist, exit early
		return false
	}

	// 1. Delete arcs first to avoid disconnected arcs
	arcSetForVertex := g.id2outbound[vid] // may be nil
	delete(g.id2outbound, vid)

	for aid := range arcSetForVertex {
		g.deleteInboundArc(vid, aid)
		g.opts.Observer.OnDeleteArc(ctx, g, vid, aid)
	}

	// 2. Delete vertex
	delete(g.id2data, vid)
	g.opts.Observer.OnDeleteVertex(ctx, g, vid)

	return true
}

// InboundArcsFor returns the set of arcs that have vid as their destination.
func (g *ObjectGraph[VID, AT, VD, AD]) InboundArcsFor(vid VID) ArcSet[VID, AT] {
	return g.id2inbound[vid]
}

func (g *ObjectGraph[VID, AT, VD, AD]) OutboundArcsFor(vid VID) ArcSetWithData[VID, AT, AD] {
	return g.id2outbound[vid]
}

func (g *ObjectGraph[VID, AT, VD, AD]) VertexData(vid VID) (VD, bool) {
	vd, isSet := g.id2data[vid]
	return vd, isSet
}

func (g *ObjectGraph[VID, AT, VD, AD]) addInboundArc(from VID, aid ArcToID[VID, AT]) {
	inbound := g.id2inbound[aid.To]
	if inbound == nil {
		inbound = ArcSet[VID, AT]{}
		g.id2inbound[aid.To] = inbound
	}
	inbound[ArcToID[VID, AT]{To: from, ArcType: aid.ArcType}] = struct{}{}
}

func (g *ObjectGraph[VID, AT, VD, AD]) deleteInboundArc(from VID, aid ArcToID[VID, AT]) {
	inbound := g.id2inbound[aid.To]
	delete(inbound, ArcToID[VID, AT]{To: from, ArcType: aid.ArcType})
	if len(inbound) == 0 { // it was the last inbound arc for the vertex. Delete empty map.
		delete(g.id2inbound, aid.To)
	}
}
