package engine

import (
	"sort"
	"testing"
)

// tableRefsSorted returns a copy sorted by (database, table) for stable comparisons.
func tableRefsSorted(refs []TableReference) []TableReference {
	out := append([]TableReference(nil), refs...)
	sort.Slice(out, func(i, j int) bool {
		if out[i].Database != out[j].Database {
			return out[i].Database < out[j].Database
		}
		return out[i].Table < out[j].Table
	})
	return out
}

func assertTableRefsEqual(t *testing.T, got, want []TableReference) {
	t.Helper()
	gs, ws := tableRefsSorted(got), tableRefsSorted(want)
	if len(gs) != len(ws) {
		t.Fatalf("got %d refs (%v), want %d (%v)", len(gs), got, len(ws), want)
	}
	for i := range gs {
		if gs[i].Database != ws[i].Database || gs[i].Table != ws[i].Table {
			t.Errorf("ref[%d] = %+v, want %+v", i, gs[i], ws[i])
		}
	}
}

func TestIsStandardTable(t *testing.T) {
	tests := []struct {
		name     string
		table    string
		expected bool
	}{
		{"slow_query_logs", "slow_query_logs", true},
		{"statements", "statements", true},
		{"topsql", "topsql", true},
		{"topru", "topru", true},
		{"topsql_tidb", "topsql_tidb", true},
		{"topsql_tikv", "topsql_tikv", true},
		{"topsql_plan_meta", "topsql_plan_meta", true},
		{"topsql_sql_meta", "topsql_sql_meta", true},
		{"topsql_topru", "topsql_topru", true},
		{"log_stats", "log_stats", true},
		{"non-standard table", "other_table", false},
		{"with cluster suffix", "topsql_tidb_123456", false},
		{"slow_query with suffix", "slow_query_logs_123456", false},
		{"statements with suffix", "statements_123456", false},
		{"log_stats with suffix", "log_stats_123456", false},
		{"empty", "", false},
	}

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

func TestInferQueryType(t *testing.T) {
	tests := []struct {
		name     string
		table    string
		expected string
	}{
		{"slow_query_logs", "slow_query_logs", "SLOW_QUERY"},
		{"statements", "statements", "STATEMENTS"},
		{"topsql", "topsql", "TOP_SQL"},
		{"topru", "topru", "TOP_SQL"},
		{"topsql_tidb", "topsql_tidb", "TOP_SQL"},
		{"topsql_tikv", "topsql_tikv", "TOP_SQL"},
		{"topsql_plan_meta", "topsql_plan_meta", "TOP_SQL"},
		{"topsql_sql_meta", "topsql_sql_meta", "TOP_SQL"},
		{"topsql_topru", "topsql_topru", "TOP_SQL"},
		{"log_stats", "log_stats", "LOG_STATS"},
		{"unknown table", "other_table", ""},
		{"with cluster suffix", "topsql_tidb_123456", ""},
		{"log_stats with suffix", "log_stats_123456", ""},
		{"empty", "", ""},
	}

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

func TestUnsupportedStandardTableError(t *testing.T) {
	tests := []struct {
		name     string
		display  string
		expected string
	}{
		{
			name:     "query uses shared generic message",
			display:  ".slowquery",
			expected: `unsupported table: .slowquery, this API only supports standard tables (slow_query_logs, statements, topsql, topru, topsql_tidb, topsql_tikv, topsql_plan_meta, topsql_sql_meta, topsql_topru, log_stats)`,
		},
		{
			name:     "schema also uses shared generic message",
			display:  "slowquery",
			expected: `unsupported table: slowquery, this API only supports standard tables (slow_query_logs, statements, topsql, topru, topsql_tidb, topsql_tikv, topsql_plan_meta, topsql_sql_meta, topsql_topru, log_stats)`,
		},
		{
			name:     "non alias stays generic",
			display:  "foo",
			expected: `unsupported table: foo, this API only supports standard tables (slow_query_logs, statements, topsql, topru, topsql_tidb, topsql_tikv, topsql_plan_meta, topsql_sql_meta, topsql_topru, log_stats)`,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := UnsupportedStandardTableError(tt.display); got != tt.expected {
				t.Fatalf("UnsupportedStandardTableError(%q) = %q, want %q", tt.display, got, tt.expected)
			}
		})
	}
}

func TestInferCompatibleQueryType(t *testing.T) {
	tests := []struct {
		name    string
		refs    []TableReference
		want    string
		wantErr bool
	}{
		{
			name: "single statements family",
			refs: []TableReference{{Table: "statements"}},
			want: "STATEMENTS",
		},
		{
			name: "multiple topsql tables share one family",
			refs: []TableReference{{Table: "topsql_tidb"}, {Table: "topsql_sql_meta"}},
			want: "TOP_SQL",
		},
		{
			name: "existing non-statements mix keeps first-table behavior",
			refs: []TableReference{{Table: "slow_query_logs"}, {Table: "topsql"}},
			want: "SLOW_QUERY",
		},
		{
			name:    "statements and slow query are rejected",
			refs:    []TableReference{{Table: "statements"}, {Table: "slow_query_logs"}},
			wantErr: true,
		},
		{
			name:    "statements and topsql are rejected",
			refs:    []TableReference{{Table: "statements"}, {Table: "topsql"}},
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := InferCompatibleQueryType(tt.refs)
			if tt.wantErr {
				if err == nil {
					t.Fatalf("InferCompatibleQueryType() error = nil, want error")
				}
				return
			}
			if err != nil {
				t.Fatalf("InferCompatibleQueryType() error = %v", err)
			}
			if got != tt.want {
				t.Fatalf("InferCompatibleQueryType() = %q, want %q", got, tt.want)
			}
		})
	}
}

func TestValidateStatementsTableReferences(t *testing.T) {
	for _, refs := range [][]TableReference{
		{{Database: "clinic_data_proxy", Table: "statements"}},
		{{Database: "clinic-data-proxy", Table: "statements"}},
	} {
		if err := ValidateStatementsTableReferences(refs, "123"); err != nil {
			t.Fatalf("valid refs %v rejected: %v", refs, err)
		}
	}

	invalid := [][]TableReference{
		{{Table: "statements"}},
		{{Database: "CLINIC_DATA_PROXY", Table: "statements"}},
		{{Database: "other_db", Table: "statements"}},
		{{Database: "other_db", Table: "statements_123"}},
		{{Database: "clinic_data_proxy", Table: "statements_123_deadbeef"}},
	}
	for _, refs := range invalid {
		if err := ValidateStatementsTableReferences(refs, "123"); err == nil {
			t.Fatalf("caller-controlled Statements refs should be rejected: %v", refs)
		}
	}
}

func TestExtractStandardTables(t *testing.T) {
	refs := []TableReference{
		{Database: "db1", Table: "topsql_tidb"},
		{Database: "db1", Table: "topsql_sql_meta"},
		{Database: "db1", Table: "some_random_table"},
		{Database: "db1", Table: "slow_query_logs"},
		{Database: "db1", Table: "log_stats"},
	}

	got := ExtractStandardTables(refs)
	if len(got) != 4 {
		t.Fatalf("ExtractStandardTables returned %d tables, want 4", len(got))
	}
	expectedTables := []string{"topsql_tidb", "topsql_sql_meta", "slow_query_logs", "log_stats"}
	for i, ref := range got {
		if ref.Table != expectedTables[i] {
			t.Errorf("got table[%d] = %q, want %q", i, ref.Table, expectedTables[i])
		}
	}
}

func TestExtractTableReferences_TopSQL(t *testing.T) {
	tests := []struct {
		name     string
		sql      string
		expected []TableReference
	}{
		{
			name: "single topsql_tidb table",
			sql:  `SELECT * FROM "clinic_data_proxy"."topsql_tidb" WHERE "date" = '2025-03-24'`,
			expected: []TableReference{
				{Database: "clinic_data_proxy", Table: "topsql_tidb"},
			},
		},
		{
			name: "topsql JOIN with sql_meta",
			sql: `SELECT t.sql_digest, m.normalized_sql FROM "clinic_data_proxy"."topsql_tidb" t
				  JOIN "clinic_data_proxy"."topsql_sql_meta" m ON t.sql_digest = m.sql_digest`,
			expected: []TableReference{
				{Database: "clinic_data_proxy", Table: "topsql_tidb"},
				{Database: "clinic_data_proxy", Table: "topsql_sql_meta"},
			},
		},
		{
			name: "backtick format",
			sql:  "SELECT * FROM `clinic_data_proxy`.`topsql_plan_meta` WHERE `date` = '2025-03-24'",
			expected: []TableReference{
				{Database: "clinic_data_proxy", Table: "topsql_plan_meta"},
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := ExtractTableReferences(tt.sql)
			if len(got) != len(tt.expected) {
				t.Fatalf("got %d refs, want %d", len(got), len(tt.expected))
			}
			for i, ref := range got {
				if ref.Database != tt.expected[i].Database || ref.Table != tt.expected[i].Table {
					t.Errorf("ref[%d] = %v, want %v", i, ref, tt.expected[i])
				}
			}
		})
	}
}

func TestExtractTableReferences_CTEIgnoresSubqueryAliases(t *testing.T) {
	sql := "WITH `base` AS (SELECT `digest`,`Query` FROM `clinic_proxy_prod_ap_southeast_1`.`slow_query_logs_10748116073341772391` WHERE `date` IN ('20260405','20260406')), " +
		"`flagged` AS (SELECT `digest`,`Query` FROM `base`) " +
		"SELECT `digest`, COUNT(1) AS `hit_records` FROM `flagged` GROUP BY `digest`"
	got := ExtractTableReferences(sql)
	if len(got) != 1 {
		t.Fatalf("got %d refs (%v), want 1 physical table ref", len(got), got)
	}
	if got[0].Database != "clinic_proxy_prod_ap_southeast_1" || got[0].Table != "slow_query_logs_10748116073341772391" {
		t.Errorf("got ref = %+v, want slow_query_logs cluster table only", got[0])
	}
}

func TestExtractTableReferences_SingleTableName(t *testing.T) {
	tests := []struct {
		name     string
		sql      string
		expected []TableReference
	}{
		{
			name: "double-quoted single table",
			sql:  `select "sql_digest","normalized_sql" from "topsql_sql_meta_10733571517854741553" where "date"='2026-03-24' limit 3`,
			expected: []TableReference{
				{Database: "", Table: "topsql_sql_meta_10733571517854741553"},
			},
		},
		{
			name: "double-quoted JOIN without database",
			sql:  `select "t"."sql_digest","m"."normalized_sql" from "topsql_tidb_123" as "t" join "topsql_sql_meta_123" as "m" on "t"."sql_digest"="m"."sql_digest" where "t"."date"='2026-03-24'`,
			expected: []TableReference{
				{Database: "", Table: "topsql_tidb_123"},
				{Database: "", Table: "topsql_sql_meta_123"},
			},
		},
		{
			name: "unquoted single table",
			sql:  `SELECT sql_digest FROM topsql_tidb_123 WHERE date = '2026-03-24'`,
			expected: []TableReference{
				{Database: "", Table: "topsql_tidb_123"},
			},
		},
		{
			name: "mixed: two-part and single-part",
			sql:  `SELECT * FROM "db"."table1" JOIN "table2" ON "table1"."id" = "table2"."id"`,
			expected: []TableReference{
				{Database: "db", Table: "table1"},
				{Database: "", Table: "table2"},
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := ExtractTableReferences(tt.sql)
			if len(got) != len(tt.expected) {
				t.Fatalf("got %d refs (%v), want %d", len(got), got, len(tt.expected))
			}
			for i, ref := range got {
				if ref.Database != tt.expected[i].Database || ref.Table != tt.expected[i].Table {
					t.Errorf("ref[%d] = %v, want %v", i, ref, tt.expected[i])
				}
			}
		})
	}
}

func TestExtractTableReferences_SubqueriesAliasesNesting(t *testing.T) {
	tests := []struct {
		name     string
		sql      string
		expected []TableReference
	}{
		{
			name: "derived_table_in_FROM_ignores_alias",
			sql: "SELECT `x` FROM (SELECT `digest` AS `x` FROM `clinic_data_proxy`.`slow_query_logs_999` WHERE `date` = '20260101') AS `derived` " +
				"WHERE 1 = 1",
			expected: []TableReference{
				{Database: "clinic_data_proxy", Table: "slow_query_logs_999"},
			},
		},
		{
			name: "AS_alias_on_base_table_not_counted_as_table",
			sql: "SELECT `t`.`sql_digest` FROM `clinic_data_proxy`.`topsql_tidb_111` AS `t` " +
				"WHERE `t`.`date` = '20260101'",
			expected: []TableReference{
				{Database: "clinic_data_proxy", Table: "topsql_tidb_111"},
			},
		},
		{
			name: "JOIN_both_sides_with_AS_aliases",
			sql: "SELECT `a`.`sql_digest`, `b`.`normalized_sql` " +
				"FROM `clinic_data_proxy`.`topsql_tidb_222` AS `a` " +
				"JOIN `clinic_data_proxy`.`topsql_sql_meta_222` AS `b` ON `a`.`sql_digest` = `b`.`sql_digest` " +
				"WHERE `a`.`date` = '20260101'",
			expected: []TableReference{
				{Database: "clinic_data_proxy", Table: "topsql_sql_meta_222"},
				{Database: "clinic_data_proxy", Table: "topsql_tidb_222"},
			},
		},
		{
			name: "WHERE_IN_subquery_second_physical_table",
			sql: "SELECT * FROM `clinic_data_proxy`.`topsql_tidb_333` `t` " +
				"WHERE `t`.`sql_digest` IN (" +
				"  SELECT `sql_digest` FROM `clinic_data_proxy`.`topsql_sql_meta_333` WHERE `date` = '20260101'" +
				")",
			expected: []TableReference{
				{Database: "clinic_data_proxy", Table: "topsql_sql_meta_333"},
				{Database: "clinic_data_proxy", Table: "topsql_tidb_333"},
			},
		},
		{
			name: "EXISTS_subquery",
			sql: "SELECT * FROM `o11y`.`t_outer_444` `o` " +
				"WHERE EXISTS (" +
				"  SELECT 1 FROM `o11y`.`t_inner_444` `i` WHERE `i`.`k` = `o`.`k` AND `i`.`date` = '20260101'" +
				")",
			expected: []TableReference{
				{Database: "o11y", Table: "t_inner_444"},
				{Database: "o11y", Table: "t_outer_444"},
			},
		},
		{
			name: "nested_IN_subqueries_three_tables",
			sql: "SELECT * FROM `app`.`t1_555` WHERE `id` IN (" +
				"  SELECT `p` FROM `app`.`t2_555` WHERE `c` IN (" +
				"    SELECT `z` FROM `app`.`t3_555` WHERE `date` = '20260101'" +
				"  )" +
				")",
			expected: []TableReference{
				{Database: "app", Table: "t1_555"},
				{Database: "app", Table: "t2_555"},
				{Database: "app", Table: "t3_555"},
			},
		},
		{
			name: "scalar_subquery_in_SELECT_list",
			sql: "SELECT (SELECT COUNT(1) FROM `app`.`cnt_666` WHERE `date` = '20260101') AS `c`, `m`.`v` " +
				"FROM `app`.`main_666` AS `m` WHERE `m`.`date` = '20260101'",
			expected: []TableReference{
				{Database: "app", Table: "cnt_666"},
				{Database: "app", Table: "main_666"},
			},
		},
		{
			name: "derived_table_then_JOIN_inner_subquery",
			sql: "SELECT * FROM (" +
				"  SELECT `d` FROM `app`.`base_777` WHERE `date` = '20260101'" +
				") AS `x` " +
				"JOIN `app`.`side_777` AS `s` ON `x`.`d` = `s`.`d` " +
				"WHERE `s`.`id` IN (SELECT `id` FROM `app`.`lookup_777` WHERE `date` = '20260101')",
			expected: []TableReference{
				{Database: "app", Table: "base_777"},
				{Database: "app", Table: "lookup_777"},
				{Database: "app", Table: "side_777"},
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := ExtractTableReferences(tt.sql)
			assertTableRefsEqual(t, got, tt.expected)
		})
	}
}