package sqls

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/storage"
)

func TestStatementsV2DeltaLogKey(t *testing.T) {
	assert.Equal(t,
		"deltalake/org=org1/cluster=cluster1/type=sqlstatement/component=sqlstatement/_delta_log/_last_checkpoint",
		statementsV2DeltaLogKey("org1", "cluster1", statementsV2DeltaLastCheckpoint),
	)
	assert.Equal(t,
		"deltalake/org=org1/cluster=cluster1/type=sqlstatement/component=sqlstatement/_delta_log/00000000000000000000.json",
		statementsV2DeltaLogKey("org1", "cluster1", statementsV2DeltaInitialCommit),
	)
}

func TestProbeStatementsDataSourcesUsesObjectExists(t *testing.T) {
	checkpointKey := "deltalake/org=tenant/cluster=cid/type=sqlstatement/component=sqlstatement/_delta_log/_last_checkpoint"
	initialKey := "deltalake/org=tenant/cluster=cid/type=sqlstatement/component=sqlstatement/_delta_log/00000000000000000000.json"
	missingCheckpoint := "deltalake/org=tenant/cluster=missing/type=sqlstatement/component=sqlstatement/_delta_log/_last_checkpoint"
	missingInitial := "deltalake/org=tenant/cluster=missing/type=sqlstatement/component=sqlstatement/_delta_log/00000000000000000000.json"
	onlyInitialCluster := "early"
	onlyInitialKey := "deltalake/org=tenant/cluster=early/type=sqlstatement/component=sqlstatement/_delta_log/00000000000000000000.json"

	store := &probeMockStorage{
		exists: map[string]bool{
			checkpointKey:  true,
			onlyInitialKey: true,
		},
	}
	adapter := &DuckDBAdapter{store: store}

	// Prefer _last_checkpoint; stop after first hit.
	resp, err := adapter.ProbeStatementsDataSources(context.Background(), &StatementsDataSourceProbeParams{
		Env:       "prod",
		Region:    "us-west-2",
		TenantID:  "tenant",
		ClusterID: "cid",
	})
	require.NoError(t, err)
	assert.True(t, resp.V2Available)
	assert.Equal(t, 1, store.existsCalls)

	// Before first checkpoint: fall back to 00000000000000000000.json.
	resp, err = adapter.ProbeStatementsDataSources(context.Background(), &StatementsDataSourceProbeParams{
		Env:       "prod",
		Region:    "us-west-2",
		TenantID:  "tenant",
		ClusterID: onlyInitialCluster,
	})
	require.NoError(t, err)
	assert.True(t, resp.V2Available)
	assert.Equal(t, 3, store.existsCalls) // checkpoint miss + initial hit

	// Neither object exists.
	resp, err = adapter.ProbeStatementsDataSources(context.Background(), &StatementsDataSourceProbeParams{
		Env:       "prod",
		Region:    "us-west-2",
		TenantID:  "tenant",
		ClusterID: "missing",
	})
	require.NoError(t, err)
	assert.False(t, resp.V2Available)
	assert.Equal(t, 5, store.existsCalls) // + checkpoint + initial
	assert.False(t, store.exists[missingCheckpoint])
	assert.False(t, store.exists[missingInitial])
	assert.False(t, store.exists[initialKey]) // not present for cid in this map beyond checkpoint
	assert.Zero(t, store.listCalls)
}

func TestStatementsDataSourceHandler(t *testing.T) {
	gin.SetMode(gin.TestMode)
	mockAdapter := &mockDuckDBAdapter{
		probeStatementsDataSourcesFunc: func(ctx context.Context, params *StatementsDataSourceProbeParams) (*StatementsDataSourcesResponse, error) {
			assert.Equal(t, "tenant", params.TenantID)
			assert.Equal(t, "cluster-1", params.ClusterID)
			assert.Equal(t, "aws", params.Provider)
			assert.False(t, params.SharedTier)
			assert.False(t, params.NextGen)
			return &StatementsDataSourcesResponse{V2Available: true}, nil
		},
	}
	svc := &Service{SQLsDuckDBAdapter: mockAdapter}
	r := gin.New()
	r.GET("/statements/datasource", svc.StatementsDataSourceHandler)

	req := httptest.NewRequest(http.MethodGet, "/statements/datasource", nil)
	req.Header.Set("x-provider", "aws")
	req.Header.Set("x-region", "us-west-2")
	req.Header.Set("x-org-id", "tenant")
	req.Header.Set("x-project-id", "proj")
	req.Header.Set("x-cluster-id", "cluster-1")
	req.Header.Set("x-deploy-type", "dedicated")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	require.Equal(t, http.StatusOK, w.Code)
	var resp StatementsDataSourcesResponse
	require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
	assert.True(t, resp.V2Available)
}

type probeMockStorage struct {
	exists      map[string]bool
	existsCalls int
	listCalls   int
}

func (m *probeMockStorage) Object(ctx context.Context, bucket, key string) (*storage.Object, error) {
	return nil, storage.ErrNotFound
}
func (m *probeMockStorage) ObjectExists(ctx context.Context, bucket, key string) (bool, error) {
	m.existsCalls++
	return m.exists[key], nil
}
func (m *probeMockStorage) ObjectSize(ctx context.Context, bucket, key string) (int64, error) {
	return 0, nil
}
func (m *probeMockStorage) Objects(ctx context.Context, bucket, prefix string) ([]storage.ObjectMeta, error) {
	m.listCalls++
	return nil, nil
}
func (m *probeMockStorage) ObjectsDelimiter(ctx context.Context, bucket, prefix string) ([]string, error) {
	m.listCalls++
	return nil, nil
}
func (m *probeMockStorage) ObjectsCallback(ctx context.Context, bucket, prefix string, callback func(storage.ObjectMeta) bool) error {
	m.listCalls++
	return nil
}
func (m *probeMockStorage) ObjectsCallbackWithStartKey(ctx context.Context, bucket, prefix, startKey string, callback func(storage.ObjectMeta) bool) error {
	m.listCalls++
	return nil
}
func (m *probeMockStorage) SaveObject(ctx context.Context, bucket, key string, body io.Reader) error {
	return nil
}

func (m *probeMockStorage) DeleteObject(ctx context.Context, bucket, key string) error {
	return nil
}

func (m *probeMockStorage) DeleteObjects(ctx context.Context, bucket string, keys []string) error {
	return nil
}
