package cache

import (
	_ "embed"
	"fmt"
	"io"
	"log"
	"net/http/httptest"
	"net/url"
	"os"
	"testing"

	"github.com/google/go-containerregistry/pkg/name"
	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"edge-infra.dev/pkg/f8n/warehouse/cluster"
	"edge-infra.dev/pkg/f8n/warehouse/lift/unpack"
	"edge-infra.dev/pkg/f8n/warehouse/oci"
	"edge-infra.dev/pkg/f8n/warehouse/oci/layer"
	"edge-infra.dev/pkg/f8n/warehouse/oci/remote"
	"edge-infra.dev/pkg/f8n/warehouse/oci/walk"
	"edge-infra.dev/pkg/f8n/warehouse/pallet"
	"edge-infra.dev/pkg/f8n/warehouse/whtest/registry"
	"edge-infra.dev/pkg/lib/uuid"
	"edge-infra.dev/test/fixtures"
)

var (
	fpath    *fixtures.Path
	r        *httptest.Server
	rURL     *url.URL
	shoot    oci.Artifact
	redpanda oci.Artifact
	certmgr  oci.Artifact
	store    oci.Artifact
)

type simpleTestRecorder struct {
	gets int
	hits int
}

func (r *simpleTestRecorder) RecordGet(hit bool, _ string, _ string) {
	r.gets++
	if hit {
		r.hits++
	}
}

func (r *simpleTestRecorder) misses() int {
	return r.gets - r.hits
}

func TestMain(m *testing.M) {
	var err error

	rURL, r, err = registry.New(registry.Logger(log.New(io.Discard, "", 0)))
	if err != nil {
		panic(err)
	}
	defer func() {
		r.Close()
	}()

	fpath, err = fixtures.Layout()
	if err != nil {
		panic(err)
	}

	shoot = loadArtifact(fpath, "shoot:latest")
	redpanda = loadArtifact(fpath, "redpanda-system:latest")
	certmgr = loadArtifact(fpath, "cert-manager:latest")
	store = loadArtifact(fpath, "store:latest")

	os.Exit(m.Run())
}

func TestNew(t *testing.T) {
	warehouseCache, err := New(
		WithMemoryCacheSize(250),
	)
	assert.NoError(t, err)
	assert.NotEmpty(t, warehouseCache)
}

func TestCache(t *testing.T) {
	var (
		id            = "test-cache"
		cache         Cache
		recorder      = &simpleTestRecorder{gets: 0, hits: 0}
		expectedGets  = 0
		shootTag      = testTag(t, id, "shoot", uuid.New().Hash())
		shootDigest   = testDigest(t, id, "shoot", shoot)
		certmgrTag    = testTag(t, id, "cert-manager", uuid.New().Hash())
		certmgrDigest = testDigest(t, id, "cert-manager", certmgr)
	)

	cache, err := New(
		WithMemoryCacheSize(10),
		WithRecorder(recorder),
	)
	assert.NoError(t, err)

	t.Run("Add", func(t *testing.T) {
		_, err := cache.Add(shoot)
		assert.NoError(t, err)
	})

	// - Verifies we only call RecordGet() once per Get() operation
	// - Verifies expected cache hits
	t.Run("Get", func(t *testing.T) {
		// Ensure we fall back to remote by asking for ref we didn't add
		assert.NoError(t, remote.Write(certmgr, certmgrTag))
		assert.NoError(t, remote.Write(shoot, shootTag))

		refs := []name.Reference{
			shootTag, shootDigest, certmgrTag,
		}
		artifacts := []oci.Artifact{
			shoot, shoot, certmgr,
		}
		for i, ref := range refs {
			testGet(t, cache, ref, artifacts[i])
			expectedGets = expectedGets + 1
			assert.Equal(t, expectedGets, recorder.gets)
		}
		assert.Equal(t, 2, recorder.hits)
		assert.Equal(t, 1, recorder.misses())

		refs = []name.Reference{
			certmgrTag, certmgrDigest,
		}
		artifacts = []oci.Artifact{
			certmgr, certmgr,
		}
		for i, ref := range refs {
			testGet(t, cache, ref, artifacts[i])
			expectedGets = expectedGets + 1
			assert.Equal(t, expectedGets, recorder.gets)
		}
		assert.Equal(t, 4, recorder.hits)
		assert.Equal(t, 1, recorder.misses())

		// Ensure that a unfamiliar tag with a resolve digest doesn't
		// break
		newTag := testTag(t, id, "cert-manager", uuid.New().Hash())
		assert.NoError(t, remote.Write(certmgr, newTag))
		t.Log(certmgrDigest)
		testGet(t, cache, newTag, certmgr)
		expectedGets = expectedGets + 1
		assert.Equal(t, expectedGets, recorder.gets)

		assert.Equal(t, 5, recorder.hits)
		assert.Equal(t, 1, recorder.misses())
	})

	t.Run("Get_ResolveTag", func(t *testing.T) {
		_, err := cache.Add(redpanda)
		assert.NoError(t, err)
		assert.NoError(t, remote.Write(redpanda, certmgrTag))
		testGet(t, cache, certmgrTag, redpanda, ResolveTag())
		expectedGets = expectedGets + 1
		assert.Equal(t, expectedGets, recorder.gets)

		assert.Equal(t, 6, recorder.hits)
		assert.Equal(t, 1, recorder.misses())
	})
}

func TestLazyFetch(t *testing.T) {
	var (
		id          = "test-lazy-cache"
		cache       Cache
		recorder    = &simpleTestRecorder{0, 0}
		storeTag    = testTag(t, id, "store", uuid.New().Hash())
		storeDigest = testDigest(t, id, "store", store)
	)

	// compute what blobs exist in store artifact sourced from v1.Layout, so we
	// dont cause the cached artifact to be computed
	blobs, blobCount, err := analyzeArtifact(store)
	require.NoError(t, err)
	// same for creating pallet
	p, err := pallet.New(store)
	require.NoError(t, err)
	// write to registry, cache will lazily pull remote impl of store from here
	require.NoError(t, remote.Write(store, storeDigest))

	cache, err = New(
		WithMemoryCacheSize(500),
		WithRecorder(recorder),
	)
	require.NoError(t, err)

	cachedStore := testColdGet(t, cache, storeDigest, blobs)

	// We don't actually need to interact with result, just force all layers to be
	// visited.
	computeArtifact(t, cachedStore, p.Providers()...)

	// We should have visited all layers and all of our blobs should now be in the
	// cache.
	checkCacheForBlobs(t, cache, blobs, blobCount)
	assert.Equal(t, cache.Len(), blobCount, "cache has more blobs than expected")

	// Shouldn't trigger any misses going forward
	finalMissCount := recorder.misses()
	_, err = cache.Get(storeDigest)
	assert.NoError(t, err)
	assert.Equal(t, finalMissCount, recorder.misses())

	t.Run("Tags", func(t *testing.T) {
		// Ensure we resolve tag for digest we have already cached without performing
		// fetch.
		require.NoError(t, remote.Tag(storeTag.(name.Tag), store))
		cachedStore, err := cache.Get(storeTag)
		assert.NoError(t, err)
		sameDigest(t, store, cachedStore)
		assert.Equal(t, finalMissCount, recorder.misses())
	})
}

func TestCommonBlobs(t *testing.T) {
	// TODO: this would be best suited for a test recorder that keeps track of
	// what object types its seen. LazyCache implements [Recorder] including tracking
	// object type
	var (
		id          = "test-lazy-common-blobs"
		cache       Cache
		recorder    = &simpleTestRecorder{0, 0}
		shootRef    = testDigest(t, id, "store", shoot)
		redpandaRef = testDigest(t, id, "redpanda-system", redpanda)
	)

	cache, err := New(
		WithMemoryCacheSize(500),
		WithRecorder(recorder),
	)
	require.NoError(t, err)

	shootBlobs, shootBlobCount, err := analyzeArtifact(shoot)
	require.NoError(t, err)
	shootPallet, err := pallet.New(shoot)
	require.NoError(t, err)
	require.NoError(t, remote.Write(shoot, shootRef))

	redpandaBlobs, redpandaBlobCount, err := analyzeArtifact(redpanda)
	require.NoError(t, err)
	redpandaPallet, err := pallet.New(redpanda)
	require.NoError(t, err)
	require.NoError(t, remote.Write(redpanda, redpandaRef))

	// Get dependency first and force it to cache
	cachedRedpanda := testColdGet(t, cache, redpandaRef, redpandaBlobs)
	computeArtifact(t, cachedRedpanda, redpandaPallet.Providers()...)
	// Confirm it was cached
	checkCacheForBlobs(t, cache, redpandaBlobs, redpandaBlobCount)
	assert.Equal(t, cache.Len(), redpandaBlobCount, "cache has more blobs than expected")

	// Only verify the unique blobs for shoot are missing
	missingShootBlobs := map[string][]v1.Hash{} // should be missing
	for objType, hashes := range shootBlobs {
		for _, hash := range hashes {
			found := false
			for _, rpHash := range redpandaBlobs[objType] {
				if hash == rpHash {
					found = true
					break
				}
			}
			if !found {
				missingShootBlobs[objType] = append(missingShootBlobs[objType], hash)
			}
		}
	}
	cachedShoot := testColdGet(t, cache, shootRef, missingShootBlobs)
	// Capture misses at this point so that we can analyze the difference after
	// computing the artifact
	misses := recorder.misses()
	computeArtifact(t, cachedShoot, shootPallet.Providers()...)
	// Confirm it was cached
	checkCacheForBlobs(t, cache, shootBlobs, shootBlobCount)
	assert.Equal(t, cache.Len(), shootBlobCount, "cache has more blobs than expected")
	// We should have hit the cache for all common blobs. Because some blobs are used
	// for multiple providers, hits will grow faster than expected, but misses
	// should be equal to:
	// (shoots unique blobs - redpandas blobs) - 1 (because we fetched manifest earlier when we cached)
	// + the misses before we pulled (so we only count misses for pulling shoot through)
	assert.Equal(t, (shootBlobCount-redpandaBlobCount-1)+misses, recorder.misses())
}

func sameDigest(t *testing.T, exp, actual oci.Artifact) {
	ad, err := actual.Digest()
	assert.NoError(t, err)

	ed, err := exp.Digest()
	assert.NoError(t, err)

	assert.Equal(t, ed, ad)
}

func testTag(t *testing.T, r, n, tag string) name.Reference {
	t.Helper()
	ref, err := name.NewTag(fmt.Sprintf("%s/%s/%s:%s", rURL.Host, r, n, tag))
	require.NoError(t, err)
	return ref
}

func testDigest(t *testing.T, r, n string, a oci.Artifact) name.Reference {
	t.Helper()
	d, err := a.Digest()
	require.NoError(t, err)
	ref, err := name.NewDigest(fmt.Sprintf("%s/%s/%s@%s", rURL.Host, r, n, d))
	require.NoError(t, err)
	return ref
}

func testColdGet(t *testing.T, c Cache, ref name.Reference, blobs map[string][]v1.Hash) oci.Artifact {
	t.Helper()

	cached, err := c.Get(ref)
	require.NoError(t, err)

	d, err := cached.Digest()
	require.NoError(t, err)

	assert.True(t, c.Exists(d), "%s expected to be in cache", d)
	// Verify we are being as lazy as expected
	for objType, hashes := range blobs {
		for _, hash := range hashes {
			if hash != d {
				assert.False(t, c.Exists(hash),
					"%s %s unexpectedly in cache", objType, hash)
			}
		}
	}

	return cached
}

func checkCacheForBlobs(t *testing.T, c Cache, blobs map[string][]v1.Hash, expCount int) {
	t.Helper()

	var missing []v1.Hash
	for _, hashes := range blobs {
		for _, hash := range hashes {
			if !c.Exists(hash) {
				missing = append(missing, hash)
			}
		}
	}
	if len(missing) > 0 {
		t.Errorf("%d/%d, expected digests not in the cache: %v", len(missing), expCount, missing)
	}
}

func testGet(t *testing.T, c Cache, ref name.Reference, exp oci.Artifact, opts ...GetOption) {
	t.Helper()
	t.Log("ref", ref)
	a, err := c.Get(ref, opts...)
	assert.NoError(t, err)
	h, err := a.Digest()
	assert.NoError(t, err)
	assert.True(t, c.Exists(h))
	assert.NoError(t, err)
	sameDigest(t, exp, a)
}

// helper for loading artifacts in TestMain. panics if fetching the artifact
// from the layout fails, not appropriate for other Test functions
func loadArtifact(p *fixtures.Path, refstr string) oci.Artifact {
	ref, err := name.ParseReference(refstr)
	if err != nil {
		panic(err)
	}

	a, err := p.Get(ref)
	if err != nil {
		panic(err)
	}

	return a
}

// return number of blobs and map of type of blob to hashes present for that
// type. everything in the returned map should be present in the cache after
// computing the lazyImage/lazyIndex
//
// NOTE: this function will cause the cache to be at least partially computed,
// if you want to avoid that, use the uncached implementation (eg v1/layout or
// v1/remote)
func analyzeArtifact(a oci.Artifact) (map[string][]v1.Hash, int, error) {
	result := make(map[string][]v1.Hash)
	count := 0

	if err := walk.Walk(a, &walk.Fns{
		Index: func(ii v1.ImageIndex, _ v1.ImageIndex) error {
			h, err := ii.Digest()
			if err != nil {
				return err
			}
			for _, idxHash := range result["index"] {
				if idxHash == h {
					return nil
				}
			}
			result["index"] = append(result["index"], h)
			count++
			return nil
		},
		Image: func(img v1.Image, _ v1.ImageIndex) error {
			h, err := img.Digest()
			if err != nil {
				return err
			}
			for _, imgHash := range result["image"] {
				if imgHash == h {
					return nil
				}
			}
			result["image"] = append(result["image"], h)
			count++
			return nil
		},
		Layer: func(l layer.Layer, _ v1.Image) error {
			h, err := l.Digest()
			if err != nil {
				return err
			}
			for _, layerHash := range result["layer"] {
				if layerHash == h {
					return nil
				}
			}
			result["layer"] = append(result["layer"], h)
			count++
			return nil
		},
	}); err != nil {
		return nil, 0, err
	}

	return result, count, nil
}

// computeArtifact performs a no-op unpack of the input layers to force lazy
// implementations to resolve and cache themselves.
func computeArtifact(t *testing.T, a oci.Artifact, providers ...cluster.Provider) {
	t.Helper()

	fn := func(pallet.Pallet, []layer.Layer) error { return nil }
	for _, provider := range providers {
		t.Log("unpacked provider", provider)
		assert.NoError(t, unpack.Walk(a, fn, unpack.ForProvider(provider)))
	}
}