package query

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestResolveWaitSeconds(t *testing.T) {
	s := &Service{cfg: Config{DefaultWaitSeconds: 30, MaxWaitSeconds: 60}}
	assert.Equal(t, 30, s.resolveWaitSeconds(nil))
	zero := 0
	assert.Equal(t, 0, s.resolveWaitSeconds(&zero))
	neg := -1
	assert.Equal(t, 0, s.resolveWaitSeconds(&neg))
	five := 5
	assert.Equal(t, 5, s.resolveWaitSeconds(&five))
	tooBig := 120
	assert.Equal(t, 60, s.resolveWaitSeconds(&tooBig))

	s.cfg.DefaultWaitSeconds = 45
	s.cfg.MaxWaitSeconds = 90
	assert.Equal(t, 45, s.resolveWaitSeconds(nil))
	assert.Equal(t, 90, s.resolveWaitSeconds(&tooBig))
}

func TestDefaultConfigFromEnvWaitSeconds(t *testing.T) {
	t.Setenv(envOlapDefaultWaitSeconds, "")
	t.Setenv(envOlapMaxWaitSeconds, "")
	cfg := DefaultConfigFromEnv("aws", "us-east-1")
	assert.Equal(t, 30, cfg.DefaultWaitSeconds)
	assert.Equal(t, 60, cfg.MaxWaitSeconds)

	t.Setenv(envOlapDefaultWaitSeconds, "45")
	t.Setenv(envOlapMaxWaitSeconds, "120")
	cfg = DefaultConfigFromEnv("aws", "us-east-1")
	assert.Equal(t, 45, cfg.DefaultWaitSeconds)
	assert.Equal(t, 120, cfg.MaxWaitSeconds)

	t.Setenv(envOlapDefaultWaitSeconds, "200")
	t.Setenv(envOlapMaxWaitSeconds, "90")
	cfg = DefaultConfigFromEnv("aws", "us-east-1")
	assert.Equal(t, 90, cfg.DefaultWaitSeconds) // clamped to max
	assert.Equal(t, 90, cfg.MaxWaitSeconds)
}

func TestResolveMaxRows(t *testing.T) {
	n, err := resolveMaxRows(nil)
	require.NoError(t, err)
	assert.Equal(t, defaultResultRows, n)

	ok := int64(100)
	n, err = resolveMaxRows(&ok)
	require.NoError(t, err)
	assert.Equal(t, int64(100), n)

	tooBig := int64(1001)
	_, err = resolveMaxRows(&tooBig)
	require.Error(t, err)
}

type fakeEngine struct {
	mu            sync.Mutex
	statusSeq     []string
	statusCalls   int
	page          *Page
	submitCalled  int
	fetchCalled   int
	presignCalled int
	lastMaxRows   int64
	statusDelay   time.Duration
	presign       *PresignedResult
	presignErr    error
}

func (f *fakeEngine) Name() string { return "fake" }

func (f *fakeEngine) Submit(ctx context.Context, req *EngineSubmitRequest) (*Handle, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.submitCalled++
	return &Handle{EngineName: f.Name(), ExecutionID: "exec-1"}, nil
}

func (f *fakeEngine) Status(ctx context.Context, handle *Handle) (*ExecutionStatus, error) {
	if f.statusDelay > 0 {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-time.After(f.statusDelay):
		}
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	idx := f.statusCalls
	if idx >= len(f.statusSeq) {
		idx = len(f.statusSeq) - 1
	}
	f.statusCalls++
	status := f.statusSeq[idx]
	st := &ExecutionStatus{
		Status:                    status,
		ScannedBytes:              42,
		EngineExecutionTimeMillis: 7,
		TotalExecutionTimeMillis:  11,
	}
	if isTerminalStatus(status) {
		now := time.Now().UTC()
		st.CompletedAt = &now
	}
	return st, nil
}

func (f *fakeEngine) FetchPage(ctx context.Context, handle *Handle, nextToken string, maxRows int64) (*Page, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.fetchCalled++
	f.lastMaxRows = maxRows
	if f.page == nil {
		return &Page{}, nil
	}
	return f.page, nil
}

func (f *fakeEngine) PresignResult(ctx context.Context, handle *Handle, expires time.Duration) (*PresignedResult, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.presignCalled++
	if f.presignErr != nil {
		return nil, f.presignErr
	}
	if f.presign == nil {
		n := int64(128)
		return &PresignedResult{
			Format:    "csv",
			S3URI:     "s3://bucket/prefix/exec-1.csv",
			URL:       "https://example.com/exec-1.csv?X-Amz-Signature=fake",
			ExpiresAt: time.Now().UTC().Add(expires),
			Bytes:     &n,
		}, nil
	}
	return f.presign, nil
}

func (f *fakeEngine) Cancel(ctx context.Context, handle *Handle) error { return nil }

func newTestService(engine Engine, region string) *Service {
	cfg := Config{
		Vendor:             "aws",
		Region:             region,
		LogsBucket:         "test-bucket",
		DefaultWaitSeconds: fallbackDefaultWaitSeconds,
		MaxWaitSeconds:     fallbackMaxWaitSeconds,
	}
	return &Service{
		cfg:     cfg,
		engine:  engine,
		jobs:    newJobStore(),
		catalog: &Catalog{cfg: cfg, tables: resolveRegistryTables(cfg)},
	}
}

func TestHandleSubmitWaitIncludesResult(t *testing.T) {
	gin.SetMode(gin.TestMode)
	engine := &fakeEngine{
		statusSeq: []string{StatusRunning, StatusSucceeded},
		page: &Page{
			Columns: []Column{{Name: "serverless_cluster_id", Type: "varchar"}},
			Rows:    [][]any{{"c1"}, {"c2"}},
		},
	}
	s := newTestService(engine, "us-east-1")
	r := gin.New()
	InitOlapApi(r.Group("/api/v1"), s)

	body := `{"sql":"SELECT 1"}`
	req := httptest.NewRequest(http.MethodPost, "/api/v1/olap", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("x-provider", "aws")
	req.Header.Set("x-region", "us-east-1")
	req.Header.Set("x-cluster-id", "cluster-1")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	require.Equal(t, http.StatusOK, w.Code)
	assert.Contains(t, w.Body.String(), `"status":"SUCCEEDED"`)
	assert.Contains(t, w.Body.String(), `"result"`)
	assert.Contains(t, w.Body.String(), `"serverless_cluster_id"`)
	assert.Contains(t, w.Body.String(), `"c1"`)
	assert.Equal(t, 1, engine.fetchCalled)
}

func TestHandleSubmitWaitSecondsZeroKeepsAsync(t *testing.T) {
	gin.SetMode(gin.TestMode)
	engine := &fakeEngine{
		statusSeq: []string{StatusSucceeded},
		page: &Page{
			Columns: []Column{{Name: "id", Type: "varchar"}},
			Rows:    [][]any{{"1"}},
		},
	}
	s := newTestService(engine, "us-east-1")
	r := gin.New()
	InitOlapApi(r.Group("/api/v1"), s)

	body := `{"sql":"SELECT 1","wait_seconds":0}`
	req := httptest.NewRequest(http.MethodPost, "/api/v1/olap", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("x-provider", "aws")
	req.Header.Set("x-region", "us-east-1")
	req.Header.Set("x-cluster-id", "cluster-1")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	require.Equal(t, http.StatusAccepted, w.Code)
	assert.Contains(t, w.Body.String(), `"status":"RUNNING"`)
	assert.NotContains(t, w.Body.String(), `"result"`)
	assert.Equal(t, 0, engine.fetchCalled)
	assert.Equal(t, 0, engine.statusCalls)
}

func TestWaitForTerminalTimeoutReturnsLastStatus(t *testing.T) {
	engine := &fakeEngine{statusSeq: []string{StatusRunning, StatusRunning, StatusRunning}}
	s := newTestService(engine, "us-east-1")
	st, err := s.waitForTerminal(context.Background(), &Handle{EngineName: "fake", ExecutionID: "e1"}, 250*time.Millisecond)
	require.NoError(t, err)
	require.NotNil(t, st)
	assert.Equal(t, StatusRunning, st.Status)
	assert.GreaterOrEqual(t, engine.statusCalls, 1)
}

func TestHandleSubmitWaitRespectsMaxRows(t *testing.T) {
	gin.SetMode(gin.TestMode)
	engine := &fakeEngine{
		statusSeq: []string{StatusSucceeded},
		page: &Page{
			Columns: []Column{{Name: "id", Type: "varchar"}},
			Rows:    [][]any{{"1"}},
		},
	}
	s := newTestService(engine, "us-east-1")
	r := gin.New()
	InitOlapApi(r.Group("/api/v1"), s)

	body := `{"sql":"SELECT 1","max_rows":50}`
	req := httptest.NewRequest(http.MethodPost, "/api/v1/olap", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("x-provider", "aws")
	req.Header.Set("x-region", "us-east-1")
	req.Header.Set("x-cluster-id", "cluster-1")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	require.Equal(t, http.StatusOK, w.Code)
	assert.Equal(t, int64(50), engine.lastMaxRows)
	assert.NotContains(t, w.Body.String(), `"download_path"`)
}

func TestHandleSubmitExportResultIsPresignedURL(t *testing.T) {
	gin.SetMode(gin.TestMode)
	engine := &fakeEngine{statusSeq: []string{StatusSucceeded}}
	s := newTestService(engine, "us-east-1")
	r := gin.New()
	InitOlapApi(r.Group("/api/v1"), s)

	body := `{"sql":"SELECT 1","output_mode":"export"}`
	req := httptest.NewRequest(http.MethodPost, "/api/v1/olap", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("x-provider", "aws")
	req.Header.Set("x-region", "us-east-1")
	req.Header.Set("x-cluster-id", "cluster-1")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	require.Equal(t, http.StatusOK, w.Code)
	bodyStr := w.Body.String()
	assert.Contains(t, bodyStr, `"output_mode":"export"`)
	assert.Contains(t, bodyStr, `"result"`)
	assert.Contains(t, bodyStr, `"download_urls"`)
	assert.Contains(t, bodyStr, `https://example.com/exec-1.csv`)
	assert.NotContains(t, bodyStr, `"rows"`)
	assert.Equal(t, 0, engine.fetchCalled)
	assert.Equal(t, 1, engine.presignCalled)
}

func TestHandleDownload(t *testing.T) {
	gin.SetMode(gin.TestMode)
	engine := &fakeEngine{statusSeq: []string{StatusSucceeded}}
	s := newTestService(engine, "us-east-1")
	job := s.jobs.create(&Job{
		ID:          "q-1",
		ClusterID:   "cluster-1",
		EngineName:  "fake",
		ExecutionID: "exec-1",
		Status:      StatusSucceeded,
		Region:      "us-east-1",
		Provider:    "aws",
	})

	r := gin.New()
	InitOlapApi(r.Group("/api/v1"), s)
	req := httptest.NewRequest(http.MethodGet, "/api/v1/olap/"+job.ID+"/download?expires_seconds=600", nil)
	req.Header.Set("x-provider", "aws")
	req.Header.Set("x-region", "us-east-1")
	req.Header.Set("x-cluster-id", "cluster-1")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	require.Equal(t, http.StatusOK, w.Code)
	body := w.Body.String()
	assert.Contains(t, body, `"format":"csv"`)
	assert.Contains(t, body, `"download_urls"`)
	assert.Contains(t, body, `https://example.com/exec-1.csv`)
	assert.Equal(t, 1, engine.presignCalled)
}

func TestHandleResultsRejectsMaxRowsAboveAthenaLimit(t *testing.T) {
	gin.SetMode(gin.TestMode)
	engine := &fakeEngine{statusSeq: []string{StatusSucceeded}}
	s := newTestService(engine, "us-east-1")
	job := s.jobs.create(&Job{
		ID:          "q-2",
		ClusterID:   "cluster-1",
		EngineName:  "fake",
		ExecutionID: "exec-1",
		Status:      StatusSucceeded,
		Region:      "us-east-1",
		Provider:    "aws",
	})
	r := gin.New()
	InitOlapApi(r.Group("/api/v1"), s)
	req := httptest.NewRequest(http.MethodGet, "/api/v1/olap/"+job.ID+"/results?max_rows=1001", nil)
	req.Header.Set("x-provider", "aws")
	req.Header.Set("x-region", "us-east-1")
	req.Header.Set("x-cluster-id", "cluster-1")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	require.Equal(t, http.StatusBadRequest, w.Code)
}

func TestHandleListTables(t *testing.T) {
	gin.SetMode(gin.TestMode)
	s := newTestService(&fakeEngine{}, "us-east-1")
	r := gin.New()
	InitOlapApi(r.Group("/api/v1"), s)

	req := httptest.NewRequest(http.MethodGet, "/api/v1/olap/tables", nil)
	req.Header.Set("x-provider", "aws")
	req.Header.Set("x-region", "us-east-1")
	req.Header.Set("x-cluster-id", "cluster-1")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	require.Equal(t, http.StatusOK, w.Code)
	body := w.Body.String()
	assert.Contains(t, body, `"general_logs_us_east_1"`)
	assert.Contains(t, body, `"log_stats_us_east_1"`)
	assert.Contains(t, body, `"topsql_us_east_1"`)
	assert.Contains(t, body, `"sqlstatement_us_east_1"`)
	assert.Contains(t, body, `"slowlogs_us_east_1"`)
	assert.Contains(t, body, `"log-stats"`)
	assert.Contains(t, body, `"topsql"`)
	assert.Contains(t, body, `"sqlstatement"`)
	assert.Contains(t, body, `"slowlogs"`)
	assert.Contains(t, body, `"catalog_status":"MISSING"`)
	assert.Contains(t, body, `"data_status":"UNKNOWN"`)
}