package btts

import (
	"context"
	"encoding/json"
	"fmt"
	"math"
	"net/url"
	"sort"
	"strings"
	"sync"
	"testing"
	"time"

	"cloud.google.com/go/bigtable"
	"github.com/stretchr/testify/assert"
	"go.skia.org/infra/go/paramtools"
	"go.skia.org/infra/go/query"
	"go.skia.org/infra/go/testutils/unittest"
	"go.skia.org/infra/go/vec32"
	"go.skia.org/infra/perf/go/btts_testutils"
	"go.skia.org/infra/perf/go/config"
	"go.skia.org/infra/perf/go/types"
)

var (
	cfg = &config.PerfBigTableConfig{
		TileSize: 256,
		Project:  "test",
		Instance: "test",
		Table:    "test",
		Topic:    "",
		GitUrl:   "",
		Shards:   8,
	}
)

func TestBasic(t *testing.T) {
	unittest.LargeTest(t)
	unittest.RequiresBigTableEmulator(t)

	ctx := context.Background()
	btts_testutils.CreateTestTable(t)
	defer btts_testutils.CleanUpTestTable(t)

	b, err := NewBigTableTraceStoreFromConfig(ctx, cfg, &btts_testutils.MockTS{}, true)
	assert.NoError(t, err)

	// Create an OPS in a fresh tile.
	tileKey := TileKeyFromOffset(1)
	op, err := b.updateOrderedParamSet(tileKey, paramtools.ParamSet{
		"cpu":    []string{"x86", "arm"},
		"config": []string{"8888", "565"},
	})
	assert.NoError(t, err)
	assert.Len(t, op.KeyOrder, 2)

	// Then update that OPS.
	op, err = b.updateOrderedParamSet(tileKey, paramtools.ParamSet{
		"os": []string{"linux", "win"},
	})
	assert.NoError(t, err)
	assert.Len(t, op.KeyOrder, 3)

	// Do we calculate LatestTile correctly?
	latest, err := b.GetLatestTile()
	assert.NoError(t, err)
	assert.Equal(t, int32(1), latest.Offset())

	// Add an OPS for a new tile.
	tileKey2 := TileKeyFromOffset(4)
	op, err = b.updateOrderedParamSet(tileKey2, paramtools.ParamSet{
		"os": []string{"win", "linux"},
	})

	// Do we calculate LatestTile correctly?
	latest, err = b.GetLatestTile()
	assert.NoError(t, err)
	assert.Equal(t, int32(4), latest.Offset())

	// Create another instance, so it has no cache.
	b2, err := NewBigTableTraceStoreFromConfig(ctx, cfg, &btts_testutils.MockTS{}, false)
	assert.NoError(t, err)

	// OPS for tile 4 should be a no-op since it's already in BT.
	op2, err := b2.updateOrderedParamSet(tileKey2, paramtools.ParamSet{
		// Note we reverse "linux", "win" order, but still get the same
		// result as op.
		"os": []string{"linux", "win"},
	})
	assert.Equal(t, op, op2)
}

func TestOPSThreaded(t *testing.T) {
	unittest.LargeTest(t)
	unittest.RequiresBigTableEmulator(t)

	ctx := context.Background()
	btts_testutils.CreateTestTable(t)
	defer btts_testutils.CleanUpTestTable(t)

	b, err := NewBigTableTraceStoreFromConfig(ctx, cfg, &btts_testutils.MockTS{}, true)
	assert.NoError(t, err)

	tileKey := TileKeyFromOffset(1)

	expected := paramtools.ParamSet{}
	// Add multiple params to the OPS in goroutines.
	wg := sync.WaitGroup{}
	for _, cpu := range []string{"x86", "arm"} {
		for _, config := range []string{"8888", "565"} {
			for _, os := range []string{"linux", "win"} {
				paramset := paramtools.ParamSet{
					"cpu":    []string{cpu},
					"config": []string{config},
					"os":     []string{os},
				}
				expected.AddParamSet(paramset)
				wg.Add(1)
				go func() {
					defer wg.Done()
					_, err := b.updateOrderedParamSet(tileKey, paramset)
					assert.NoError(t, err)
				}()
			}
		}
	}
	wg.Wait()

	// read current OPS
	entry, _, err := b.getOPS(tileKey)
	assert.NoError(t, err)
	expected.Normalize()
	entry.ops.ParamSet.Normalize()
	assert.Equal(t, expected, entry.ops.ParamSet)
}

// assertIndices asserts that the indices in the table match expectedKeys and expectedColumns.
//
// Returns the number of rows actually retrieved.
func assertIndices(t *testing.T, ops *paramtools.OrderedParamSet, b *BigTableTraceStore, expectedKeys []string, expectedColumns map[string][]string, msg string) int {
	var count int
	err := b.getTable().ReadRows(context.Background(), bigtable.PrefixRange(INDEX_PREFIX), func(row bigtable.Row) bool {
		count += 1
		assert.Contains(t, expectedKeys, row.Key())
		for _, col := range row[INDEX_FAMILY] {
			// Strip off the family name which is prefixed.
			rowKey := strings.Split(col.Column, ":")[3]
			p, err := ops.DecodeParamsFromString(rowKey)
			assert.NoError(t, err)
			decodedKey, err := query.MakeKeyFast(p)
			assert.NoError(t, err)
			assert.Contains(t, expectedColumns[col.Row], decodedKey, fmt.Sprintf("For row: %q", row.Key()))
		}
		return true
	}, bigtable.RowFilter(
		bigtable.ChainFilters(
			bigtable.LatestNFilter(1),
			bigtable.FamilyFilter(INDEX_FAMILY),
		),
	),
	)
	assert.NoError(t, err)
	assert.Equal(t, len(expectedKeys), b.indexed.Len())
	return count
}

func TestTraces(t *testing.T) {
	unittest.LargeTest(t)
	unittest.RequiresBigTableEmulator(t)

	ctx := context.Background()
	btts_testutils.CreateTestTable(t)
	defer btts_testutils.CleanUpTestTable(t)
	now := time.Now()

	b, err := NewBigTableTraceStoreFromConfig(ctx, cfg, &btts_testutils.MockTS{}, true)
	assert.NoError(t, err)

	tileKey := TileKeyFromOffset(1)
	ops, err := b.GetOrderedParamSet(ctx, tileKey)
	assert.NoError(t, err)
	assertIndices(t, ops, b, nil, nil, "Start empty")

	paramset := paramtools.ParamSet{
		"config": []string{"8888", "565"},
		"cpu":    []string{"x86", "arm"},
	}
	assert.NoError(t, err)
	params := []paramtools.Params{
		{"cpu": "x86", "config": "8888"},
		{"cpu": "x86", "config": "565"},
		{"cpu": "arm", "config": "8888"},
		{"cpu": "arm", "config": "565"},
	}
	values := []float32{
		1.0,
		1.1,
		1.2,
		1.3,
	}
	err = b.WriteTraces(257, params, values, paramset, "gs://some/test/location", now)
	assert.NoError(t, err)

	// Confirm that indices were written correctly.
	expectedKeys := []string{
		"i2147483646:config:565",
		"i2147483646:config:8888",
		"i2147483646:cpu:arm",
		"i2147483646:cpu:x86",
	}
	expectedColumns := map[string][]string{
		"i2147483646:config:565":  {",config=565,cpu=x86,", ",config=565,cpu=arm,"},
		"i2147483646:config:8888": {",config=8888,cpu=x86,", ",config=8888,cpu=arm,"},
		"i2147483646:cpu:arm":     {",config=565,cpu=arm,", ",config=8888,cpu=arm,"},
		"i2147483646:cpu:x86":     {",config=565,cpu=x86,", ",config=8888,cpu=x86,"},
	}
	ops, err = b.GetOrderedParamSet(ctx, tileKey)
	assert.NoError(t, err)
	count := assertIndices(t, ops, b, expectedKeys, expectedColumns, "First write")
	assert.Equal(t, 4, count)
	indexCount, err := b.CountIndices(context.Background(), tileKey)
	assert.NoError(t, err)
	assert.Equal(t, int64(4), indexCount)

	q, err := query.New(url.Values{"config": []string{"8888"}})
	assert.NoError(t, err)

	results, err := b.QueryTraces(ctx, tileKey, q)
	assert.NoError(t, err)
	vec1 := vec32.New(256)
	vec1[1] = 1.0
	vec2 := vec32.New(256)
	vec2[1] = 1.2
	expected := types.TraceSet{
		",config=8888,cpu=x86,": vec1,
		",config=8888,cpu=arm,": vec2,
	}
	assert.Equal(t, expected, results)
	results, err = b.QueryTracesByIndex(context.Background(), tileKey, q)
	assert.NoError(t, err)
	assert.Equal(t, expected, results)

	keys, err := b.TileKeys(ctx, tileKey)
	assert.NoError(t, err)
	sort.Strings(keys)
	assert.Equal(t, []string{",config=565,cpu=arm,", ",config=565,cpu=x86,", ",config=8888,cpu=arm,", ",config=8888,cpu=x86,"}, keys)

	// Now overwrite a value.
	params = []paramtools.Params{
		{"cpu": "x86", "config": "8888"},
	}
	values = []float32{
		2.0,
	}
	err = b.WriteTraces(257, params, values, paramset, "gs://some/other/test/location", now)
	assert.NoError(t, err)
	ops, err = b.GetOrderedParamSet(ctx, tileKey)
	assert.NoError(t, err)
	count = assertIndices(t, ops, b, expectedKeys, expectedColumns, "Overwrite")
	assert.Equal(t, 4, count)

	// Query again to get the updated value.
	results, err = b.QueryTraces(context.Background(), tileKey, q)
	assert.NoError(t, err)
	vec1 = vec32.New(256)
	vec1[1] = 2.0
	vec2 = vec32.New(256)
	vec2[1] = 1.2
	expected = types.TraceSet{
		",config=8888,cpu=x86,": vec1,
		",config=8888,cpu=arm,": vec2,
	}
	assert.Equal(t, expected, results)

	results, err = b.QueryTracesByIndex(context.Background(), tileKey, q)
	assert.NoError(t, err)
	assert.Equal(t, expected, results)

	// Write in the next column.
	params = []paramtools.Params{
		{"cpu": "x86", "config": "8888"},
	}
	values = []float32{
		3.0,
	}
	err = b.WriteTraces(258, params, values, paramset, "gs://some/other/test/location", now)
	assert.NoError(t, err)

	// Query again to get the updated value.
	results, err = b.QueryTraces(context.Background(), tileKey, q)
	assert.NoError(t, err)
	vec1 = vec32.New(256)
	vec1[1] = 2.0
	vec1[2] = 3.0
	vec2 = vec32.New(256)
	vec2[1] = 1.2
	expected = types.TraceSet{
		",config=8888,cpu=x86,": vec1,
		",config=8888,cpu=arm,": vec2,
	}
	assert.Equal(t, expected, results)
	count = assertIndices(t, ops, b, expectedKeys, expectedColumns, "Write new value.")
	assert.Equal(t, 4, count)

	results, err = b.QueryTracesByIndex(context.Background(), tileKey, q)
	assert.NoError(t, err)
	assert.Equal(t, expected, results)

	// Write to a new trace.
	params = []paramtools.Params{
		{"cpu": "risc-v", "config": "8888"},
	}
	values = []float32{
		2.0,
	}
	paramset.AddParams(params[0])
	err = b.WriteTraces(258, params, values, paramset, "gs://some/other/test/location", now)
	assert.NoError(t, err)

	// Add new trace to expectations.
	expectedKeys = append(expectedKeys, "i2147483646:cpu:risc-v")
	expectedColumns["i2147483646:cpu:risc-v"] = []string{",config=8888,cpu=risc-v,"}
	expectedColumns["i2147483646:config:8888"] = append(expectedColumns["i2147483646:config:8888"], ",config=8888,cpu=risc-v,")
	ops, err = b.GetOrderedParamSet(ctx, tileKey)
	assert.NoError(t, err)
	count = assertIndices(t, ops, b, expectedKeys, expectedColumns, "Add new trace")
	assert.Equal(t, 5, count)

	// Remove indices and repopulate with WriteIndices.
	muts := []*bigtable.Mutation{}
	for range expectedKeys {
		mut := bigtable.NewMutation()
		mut.DeleteRow()
		muts = append(muts, mut)
	}
	errs, err := b.getTable().ApplyBulk(ctx, expectedKeys, muts)
	assert.NoError(t, err)
	for _, e := range errs {
		assert.NoError(t, e)
	}
	count = assertIndices(t, ops, b, expectedKeys, expectedColumns, "Add new trace")
	assert.Equal(t, 0, count)
	err = b.WriteIndices(ctx, tileKey)
	assert.NoError(t, err)
	count = assertIndices(t, ops, b, expectedKeys, expectedColumns, "Add new trace")
	assert.Equal(t, 5, count)

	// Source
	traceId, err := query.MakeKey(paramtools.Params{"cpu": "x86", "config": "8888"})
	assert.NoError(t, err)
	s, err := b.GetSource(context.Background(), 258, traceId)
	assert.NoError(t, err)
	assert.Equal(t, "gs://some/other/test/location", s)

	s, err = b.GetSource(context.Background(), 259, traceId)
	assert.Error(t, err)
	assert.Equal(t, "", s)
}

func TestTileKey(t *testing.T) {
	unittest.SmallTest(t)

	numShards := int32(3)
	tileKey := TileKeyFromOffset(0)
	assert.Equal(t, int32(math.MaxInt32), int32(tileKey))
	assert.Equal(t, int32(0), tileKey.Offset())
	assert.Equal(t, "@2147483647", tileKey.OpsRowName())
	assert.Equal(t, "2:2147483647:", tileKey.TraceRowPrefix(2))
	assert.Equal(t, "1:2147483647:,0=1,", tileKey.TraceRowName(",0=1,", numShards))

	tileKey = TileKeyFromOffset(1)
	assert.Equal(t, int32(math.MaxInt32-1), int32(tileKey))
	assert.Equal(t, "@2147483646", tileKey.OpsRowName())
	assert.Equal(t, "3:2147483646:", tileKey.TraceRowPrefix(3))
	assert.Equal(t, "1:2147483646:,0=1,", tileKey.TraceRowName(",0=1,", numShards))

	tileKey = TileKeyFromOffset(-1)
	assert.Equal(t, BadTileKey, tileKey)

	var err error
	tileKey, err = TileKeyFromOpsRowName("2147483646")
	assert.Error(t, err)
	assert.Equal(t, BadTileKey, tileKey)

	tileKey, err = TileKeyFromOpsRowName("@2147483637")
	assert.NoError(t, err)
	assert.Equal(t, "@2147483637", tileKey.OpsRowName())
}

const rowJson = `{
  "D": [
    {
      "Row": "@2147483643",
      "Column": "D:H",
      "Timestamp": 1536145696388000,
      "Value": "NWY2MDQ5ZTk3ODdiMDcxMGFhY2U2MTYzNDU3NzRiNTk="
    },
    {
      "Row": "@2147483643",
      "Column": "D:OPS",
      "Timestamp": 1536145696388000,
      "Value": "Of+BAwEBD09yZGVyZWRQYXJhbVNldAH/ggABAgEIS2V5T3JkZXIB/4QAAQhQYXJhbVNldAH/hgAAABb/gwIBAQhbXXN0cmluZwH/hAABDAAAGf+FBAEBCFBhcmFtU2V0Af+GAAEMAf+EAAAY/4IBAQJvcwEBAm9zAgN3aW4FbGludXgA"
    }
  ]
}`

func TestOpsCacheEntry(t *testing.T) {
	unittest.SmallTest(t)
	// Entry for an empty OPS.
	o, err := NewOpsCacheEntry()
	assert.NoError(t, err)
	assert.Equal(t, 0, len(o.ops.KeyOrder))
	assert.Equal(t, "c011636276b346664a4d3a473ff07fc5", o.hash)

	// Entry for an OPS with just one key-value pair:
	ops := o.ops.Copy()
	ops.Update(paramtools.ParamSet{"config": []string{"8888"}})
	o2, err := opsCacheEntryFromOPS(ops)
	assert.NoError(t, err)
	assert.Equal(t, []string{"config"}, o2.ops.KeyOrder)
	assert.Equal(t, "7a59e4600a8f20900c933037d8e0011a", o2.hash)

	// From a BigTable row.
	goodRow := bigtable.Row{}
	err = json.Unmarshal([]byte(rowJson), &goodRow)
	assert.NoError(t, err)

	o3, err := NewOpsCacheEntryFromRow(goodRow)
	assert.NoError(t, err)
	assert.Equal(t, []string{"os"}, o3.ops.KeyOrder)
	assert.Equal(t, "5f6049e9787b0710aace616345774b59", o3.hash)

	// Empty BT row.
	row := bigtable.Row{}
	_, err = NewOpsCacheEntryFromRow(row)
	assert.Error(t, err)

	// Nothing missing.
	row = bigtable.Row{
		"D": []bigtable.ReadItem{
			goodRow["D"][0],
			goodRow["D"][1],
		},
	}
	_, err = NewOpsCacheEntryFromRow(row)
	assert.NoError(t, err)

	// Missing H.
	row = bigtable.Row{
		"D": []bigtable.ReadItem{
			goodRow["D"][0],
		},
	}
	_, err = NewOpsCacheEntryFromRow(row)
	assert.Error(t, err)

	// Missing OPS.
	row = bigtable.Row{
		"D": []bigtable.ReadItem{
			goodRow["D"][1],
		},
	}
	_, err = NewOpsCacheEntryFromRow(row)
	assert.Error(t, err)
}