package main

import (
	"context"
	"sync"
)

type snapshotKey struct {
	NodeLimit      int
	EdgeLimit      int
	IncludeBNodes  bool
	GraphQueryID   string
}

type snapshotInflight struct {
	ready    chan struct{}
	snapshot GraphResponse
	err      error
}

type GraphSnapshotService struct {
	sparql *AnzoGraphClient
	cfg    Config

	mu       sync.Mutex
	cache    map[snapshotKey]GraphResponse
	inflight map[snapshotKey]*snapshotInflight
}

func NewGraphSnapshotService(sparql *AnzoGraphClient, cfg Config) *GraphSnapshotService {
	return &GraphSnapshotService{
		sparql:    sparql,
		cfg:      cfg,
		cache:    make(map[snapshotKey]GraphResponse),
		inflight: make(map[snapshotKey]*snapshotInflight),
	}
}

func (s *GraphSnapshotService) Get(ctx context.Context, nodeLimit int, edgeLimit int, graphQueryID string) (GraphResponse, error) {
	key := snapshotKey{NodeLimit: nodeLimit, EdgeLimit: edgeLimit, IncludeBNodes: s.cfg.IncludeBNodes, GraphQueryID: graphQueryID}

	s.mu.Lock()
	if snap, ok := s.cache[key]; ok {
		s.mu.Unlock()
		return snap, nil
	}

	if inf, ok := s.inflight[key]; ok {
		ready := inf.ready
		s.mu.Unlock()
		select {
		case <-ctx.Done():
			return GraphResponse{}, ctx.Err()
		case <-ready:
			return inf.snapshot, inf.err
		}
	}

	inf := &snapshotInflight{ready: make(chan struct{})}
	s.inflight[key] = inf
	s.mu.Unlock()

	snap, err := fetchGraphSnapshot(ctx, s.sparql, s.cfg, nodeLimit, edgeLimit, graphQueryID)

	s.mu.Lock()
	inf.snapshot = snap
	inf.err = err
	delete(s.inflight, key)
	if err == nil {
		s.cache[key] = snap
	}
	close(inf.ready)
	s.mu.Unlock()

	return snap, err
}