package executor

import (
	"context"
	"database/sql"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/cql/parser"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/cql/translator"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/metrics"
)

type fakeDMLResult struct {
	rowsAffected int64
}

func (r fakeDMLResult) LastInsertId() (int64, error) {
	return 0, nil
}

func (r fakeDMLResult) RowsAffected() (int64, error) {
	return r.rowsAffected, nil
}

type fakeDMLExecutor struct {
	rowsAffected int64
}

func (e fakeDMLExecutor) ExecContext(context.Context, string, ...interface{}) (sql.Result, error) {
	return fakeDMLResult{rowsAffected: e.rowsAffected}, nil
}

func (e fakeDMLExecutor) QueryContext(context.Context, string, ...interface{}) (*sql.Rows, error) {
	return nil, nil
}

func TestCQLWriteRowsMetricSingleInsertUsesRequestedRows(t *testing.T) {
	metrics.SetClusterIDProviderForTest(nil)
	metrics.Init(metrics.Settings{Enabled: true}, nil)

	stmt := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (1, 'a'), (2, 'b')")
	insertStmt, ok := stmt.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected insert statement, got %T", stmt)
	}
	exec := &Executor{translator: translator.NewTranslator()}
	plan, err := exec.buildInsertPlan(context.Background(), "", "", insertStmt)
	if err != nil {
		t.Fatalf("buildInsertPlan failed: %v", err)
	}

	ctx := metrics.WithUsername(context.Background(), "ks.user")
	result, err := exec.executeDMLPlan(ctx, fakeDMLExecutor{rowsAffected: 4}, plan, nil)
	if err != nil {
		t.Fatalf("executeDMLPlan failed: %v", err)
	}
	metrics.ObserveCQLWriteRows(ctx, "insert", result.CQLWriteRows)

	assertCQLWriteRowsMetric(t, `proxy_cql_write_rows_total{ClusterId="-1",KeyspaceName="ks",operation_type="insert"} 2`)
}

func TestCQLWriteRowsMetricSingleUpdateUsesRequestedRows(t *testing.T) {
	metrics.SetClusterIDProviderForTest(nil)
	metrics.Init(metrics.Settings{Enabled: true}, nil)

	stmt := mustParseStatement(t, "UPDATE ks.tbl SET v = 'x' WHERE id = 1")
	updateStmt, ok := stmt.(*parser.UpdateStatement)
	if !ok {
		t.Fatalf("expected update statement, got %T", stmt)
	}
	exec := &Executor{translator: translator.NewTranslator()}
	dml := &dmlExecutor{Executor: exec}
	plan, err := dml.buildPlainUpdatePlan(context.Background(), "", "", updateStmt, []equalityCondition{
		{column: "id", value: &parser.IntegerLiteral{Value: 1}, op: "="},
	}, buildBindMarkerIndex(stmt), []string{"id"}, nil)
	if err != nil {
		t.Fatalf("buildPlainUpdatePlan failed: %v", err)
	}

	ctx := metrics.WithUsername(context.Background(), "ks.user")
	result, err := exec.executeDMLPlan(ctx, fakeDMLExecutor{rowsAffected: 0}, plan, nil)
	if err != nil {
		t.Fatalf("executeDMLPlan failed: %v", err)
	}
	metrics.ObserveCQLWriteRows(ctx, "update", result.CQLWriteRows)

	assertCQLWriteRowsMetric(t, `proxy_cql_write_rows_total{ClusterId="-1",KeyspaceName="ks",operation_type="update"} 1`)
}

func TestCQLWriteRowsMetricBatchUsesPreMergeRequestRows(t *testing.T) {
	metrics.SetClusterIDProviderForTest(nil)
	metrics.Init(metrics.Settings{Enabled: true}, nil)

	entries := []BatchEntry{
		{Statement: mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (1, 'a'), (2, 'b')")},
		{Statement: mustParseStatement(t, "UPDATE ks.tbl SET v = 'c' WHERE id = 1")},
		{Statement: mustParseStatement(t, "UPDATE ks.tbl SET v = 'd' WHERE id = 1")},
	}

	ctx := metrics.WithUsername(context.Background(), "ks.user")
	metrics.ObserveCQLWriteRows(ctx, "batch", cqlWriteRowsForBatchEntries(entries))

	assertCQLWriteRowsMetric(t, `proxy_cql_write_rows_total{ClusterId="-1",KeyspaceName="ks",operation_type="batch"} 4`)
}

func assertCQLWriteRowsMetric(t *testing.T, expected string) {
	t.Helper()
	h := metrics.Handler()
	if h == nil {
		t.Fatal("expected metrics handler when enabled")
	}
	rec := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/metrics", nil)
	h.ServeHTTP(rec, req)
	body := rec.Body.String()
	if !strings.Contains(body, expected) {
		t.Fatalf("expected metric line %q, got:\n%s", expected, body)
	}
}