package service

import (
	"testing"
)

func TestRemoveLeadingComments(t *testing.T) {
	tests := []struct {
		name     string
		input    string
		expected string
	}{
		{
			name:     "block comment before SELECT",
			input:    "/*app=bytebase*/SELECT * FROM table",
			expected: "SELECT * FROM table",
		},
		{
			name:     "line comment before SELECT",
			input:    "-- this is a comment\nSELECT * FROM table",
			expected: "SELECT * FROM table",
		},
		{
			name:     "multiple comments before SELECT",
			input:    "/*comment1*/ -- comment2\n/*comment3*/SELECT * FROM table",
			expected: "SELECT * FROM table",
		},
		{
			name:     "whitespace and comments",
			input:    "  \n  /*comment*/  \n  SELECT * FROM table",
			expected: "SELECT * FROM table",
		},
		{
			name:     "no comments",
			input:    "SELECT * FROM table WHERE id = 1",
			expected: "SELECT * FROM table WHERE id = 1",
		},
		{
			name:     "comment not at beginning",
			input:    "SELECT /*comment*/* FROM table",
			expected: "SELECT /*comment*/* FROM table",
		},
		{
			name:     "unclosed block comment",
			input:    "/* unclosed comment SELECT * FROM table",
			expected: "/* unclosed comment SELECT * FROM table",
		},
		{
			name:     "line comment only",
			input:    "-- only comment, no SQL",
			expected: "",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := removeLeadingComments(tt.input)
			if result != tt.expected {
				t.Errorf("removeLeadingComments() = %q, want %q", result, tt.expected)
			}
		})
	}
}

func TestShouldKeep(t *testing.T) {
	tests := []struct {
		name     string
		query    string
		expected bool
	}{
		{
			name:     "SELECT without comments",
			query:    "SELECT * FROM table",
			expected: true,
		},
		{
			name:     "SELECT with block comment",
			query:    "/*app=bytebase*/SELECT * FROM table",
			expected: true,
		},
		{
			name:     "SELECT with line comment",
			query:    "-- comment\nSELECT * FROM table",
			expected: true,
		},
		{
			name:     "CREATE statement",
			query:    "CREATE TABLE test",
			expected: false,
		},
		{
			name:     "CREATE with comment",
			query:    "/*comment*/CREATE TABLE test",
			expected: false,
		},
		{
			name:     "ANALYZE statement",
			query:    "ANALYZE TABLE test",
			expected: false,
		},
		{
			name:     "SELECT with force index",
			query:    "SELECT * FROM table force index(idx)",
			expected: false,
		},
		{
			name:     "SELECT with tiflash",
			query:    "SELECT * FROM table read_from_storage(tiflash)",
			expected: false,
		},
		{
			name:     "SELECT with tiflash hint",
			query:    "SELECT /*+ read_from_storage(tiflash[t1]) */ * FROM table",
			expected: true,
		},
		{
			name:     "INSERT statement",
			query:    "INSERT INTO table VALUES (1)",
			expected: false,
		},
		{
			name:     "UPDATE statement",
			query:    "UPDATE table SET col = 1",
			expected: false,
		},
		{
			name:     "DELETE statement",
			query:    "DELETE FROM table",
			expected: false,
		},
		{
			name:     "DROP statement",
			query:    "DROP TABLE test",
			expected: false,
		},
		{
			name:     "TRUNCATE statement",
			query:    "TRUNCATE TABLE test",
			expected: false,
		},
		{
			name:     "ALTER statement",
			query:    "ALTER TABLE test ADD COLUMN col INT",
			expected: false,
		},
		{
			name:     "RENAME statement",
			query:    "RENAME TABLE old TO new",
			expected: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := ShouldKeep(tt.query)
			if result != tt.expected {
				t.Errorf("ShouldKeep(%q) = %v, want %v", tt.query, result, tt.expected)
			}
		})
	}
}