package parquet

import (
	"os"
	"testing"

	"github.com/parquet-go/parquet-go"
	"github.com/parquet-go/parquet-go/format"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/observability/cmd/diagnosis-compute/pkg/types"
)

func TestOpenWriterUsesZstdCompression(t *testing.T) {
	dir := t.TempDir()
	m := &Manager{
		info: types.ClusterInfo{
			Env:       "prod",
			TenantID:  "t1",
			ClusterID: "c1",
			StorageID: "s1",
		},
		hourBegin:  1718708400, // 2024-06-18 11:00 UTC
		baseDir:    dir,
		maxBytes:   64 << 20,
		writeBatch: 64,
		writers:    map[string]*rollingWriter{},
	}

	w, err := m.openWriter("tidb", 1)
	require.NoError(t, err)
	_, err = w.writer.Write([]LogRow{{
		EventTime: "2024-06-18T11:00:00Z",
		Timestamp: 1718708400,
		Host:      "h1",
		Pod:       "p1",
		Level:     "INFO",
		Position:  "a.go:1",
		Message:   "hello world log line for compression check",
		Cluster:   "c1",
		Instance:  "db-tidb-0",
		TenantID:  "t1",
	}})
	require.NoError(t, err)
	require.NoError(t, m.closeWriter(w))

	f, err := os.Open(w.localPath)
	require.NoError(t, err)
	defer f.Close()
	st, err := f.Stat()
	require.NoError(t, err)

	pf, err := parquet.OpenFile(f, st.Size())
	require.NoError(t, err)
	require.NotEmpty(t, pf.Metadata().RowGroups)

	foundZstd := false
	for _, rg := range pf.Metadata().RowGroups {
		for _, col := range rg.Columns {
			if col.MetaData.Codec == format.Zstd {
				foundZstd = true
			} else {
				t.Fatalf("column %v codec = %v, want zstd", col.MetaData.PathInSchema, col.MetaData.Codec)
			}
		}
	}
	assert.True(t, foundZstd)
}