package main

import (
	"encoding/csv"
	"encoding/json"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestStatementStreamParser_ReadNext(t *testing.T) {
	input := strings.Join([]string{
		`{"begin":1784534264,"end":1784534325,"schema_name":"db1","serverless_tenant_id":"1372813089209260919"}`,
		`{"begin":1784534268,"end":1784534328,"schema_name":"db2","serverless_tenant_id":"other-tenant"}`,
	}, "\n")

	parser := NewStatementStreamParser(strings.NewReader(input))

	rec1, err := parser.ReadNext()
	require.NoError(t, err)
	require.NotNil(t, rec1)
	assert.Equal(t, int64(1784534325), rec1.End)
	assert.Equal(t, "1372813089209260919", rec1.ServerlessTenantID)
	assert.Equal(t, time.Unix(1784534325, 0).UTC(), rec1.eventTime())

	rec2, err := parser.ReadNext()
	require.NoError(t, err)
	require.NotNil(t, rec2)
	assert.Equal(t, "other-tenant", rec2.ServerlessTenantID)

	rec3, err := parser.ReadNext()
	require.NoError(t, err)
	assert.Nil(t, rec3)
}

func TestTruncateForLog(t *testing.T) {
	assert.Equal(t, "short", truncateForLog("short", 10))
	assert.Equal(t, "1234567890...(truncated)", truncateForLog("123456789012345", 10))
}

func TestStatementCSVHeaderOmitsRuV2ByDefault(t *testing.T) {
	header := statementCSVHeader(false)
	assert.NotContains(t, header, "sum_ruv2")
	assert.NotContains(t, header, "max_ruv2")
	assert.NotContains(t, header, "AVG_REQUEST_UNIT_V2")
	assert.Contains(t, header, "sum_rru")
	assert.Contains(t, header, "digest")

	headerWith := statementCSVHeader(true)
	assert.Contains(t, headerWith, "sum_ruv2")
	assert.Contains(t, headerWith, "max_ruv2")
	assert.Contains(t, headerWith, "AVG_REQUEST_UNIT_V2")
}

func TestStatementJSONToCSVRow(t *testing.T) {
	raw := json.RawMessage(`{
		"begin":1,"end":2,"digest":"abc","exec_count":4,
		"sum_ruv2":2.7,"max_ruv2":0.6,
		"auth_users":{"u1":{}},"backoff_types":{},
		"index_names":null,"is_internal":false,"sample_sql":"select 1"
	}`)

	row, err := statementJSONToCSVRow(raw, false)
	require.NoError(t, err)
	header := statementCSVHeader(false)
	require.Equal(t, len(header), len(row))

	col := map[string]string{}
	for i, name := range header {
		col[name] = row[i]
	}
	assert.Equal(t, "1", col["begin"])
	assert.Equal(t, "abc", col["digest"])
	assert.Equal(t, "4", col["exec_count"])
	assert.Equal(t, "false", col["is_internal"])
	assert.Equal(t, "select 1", col["sample_sql"])
	assert.Equal(t, `{"u1":{}}`, col["auth_users"])
	assert.Equal(t, `{}`, col["backoff_types"])
	assert.Equal(t, "", col["index_names"])
	_, hasRu := col["sum_ruv2"]
	assert.False(t, hasRu)

	rowWith, err := statementJSONToCSVRow(raw, true)
	require.NoError(t, err)
	headerWith := statementCSVHeader(true)
	require.Equal(t, len(headerWith), len(rowWith))
	colWith := map[string]string{}
	for i, name := range headerWith {
		colWith[name] = rowWith[i]
	}
	assert.Equal(t, "2.7", colWith["sum_ruv2"])
	assert.Equal(t, "0.6", colWith["max_ruv2"])
}

func TestStatementCSVRoundTripWritable(t *testing.T) {
	raw := json.RawMessage(`{"begin":1,"end":2,"digest":"d","sample_sql":"a,b\"c","sum_ruv2":1}`)
	row, err := statementJSONToCSVRow(raw, false)
	require.NoError(t, err)

	var buf strings.Builder
	w := csv.NewWriter(&buf)
	require.NoError(t, w.Write(statementCSVHeader(false)))
	require.NoError(t, w.Write(row))
	w.Flush()
	require.NoError(t, w.Error())

	r := csv.NewReader(strings.NewReader(buf.String()))
	records, err := r.ReadAll()
	require.NoError(t, err)
	require.Len(t, records, 2)
	assert.Equal(t, statementCSVHeader(false), records[0])
}