package engine

import (
	"fmt"
	"regexp"
	"strings"
)

const SupportedStandardTables = "slow_query_logs, statements, topsql, topru, topsql_tidb, topsql_tikv, topsql_plan_meta, topsql_sql_meta, topsql_topru, log_stats"

// TableReference 表引用
type TableReference struct {
	Database string // 数据库名
	Table    string // 表名
}

// ExtractTableReferences 从 SQL 提取表引用
// 支持以下格式：
// - FROM "database"."table"  / JOIN "database"."table"
// - FROM database.table      / JOIN database.table
// - FROM `database`.`table`  / JOIN `database`.`table`
// - FROM "table"             / JOIN "table"  (单表名，database 为空)
// - FROM `table`             / JOIN `table`  (单表名，database 为空)
// - FROM table               / JOIN table    (单表名，database 为空)
//
// 优先使用 TiDB SQL parser 解析（正确识别 WITH 子句中的 CTE 名，避免把 FROM cte 当成物理表）；
// 解析失败时回退到正则（兼容 parser 不接受的方言片段）。
func ExtractTableReferences(sql string) []TableReference {
	refs, err := extractTableReferencesAST(sql)
	if err == nil {
		return refs
	}
	return extractTableReferencesRegexp(sql)
}

func extractTableReferencesRegexp(sql string) []TableReference {
	var refs []TableReference
	seen := make(map[string]bool) // 去重

	// 先匹配两段式（database.table），再匹配单段式（table only）
	// 两段式正则
	twoPartPatterns := []string{
		// 双引号格式
		`(?i)(?:FROM|JOIN)\s+"([a-zA-Z0-9_-]+)"\."([a-zA-Z0-9_-]+)"`,
		// 无引号格式
		`(?i)(?:FROM|JOIN)\s+([a-zA-Z0-9_-]+)\.([a-zA-Z0-9_-]+)`,
		// 反引号格式
		"(?i)(?:FROM|JOIN)\\s+`([a-zA-Z0-9_-]+)`.`([a-zA-Z0-9_-]+)`",
	}

	for _, pattern := range twoPartPatterns {
		re := regexp.MustCompile(pattern)
		matches := re.FindAllStringSubmatch(sql, -1)

		for _, match := range matches {
			if len(match) == 3 {
				database := match[1]
				table := match[2]

				key := database + "." + table
				if !seen[key] {
					refs = append(refs, TableReference{
						Database: database,
						Table:    table,
					})
					seen[key] = true
				}
			}
		}
	}

	// 单段式正则（仅表名，无 database 前缀）
	// 注意：需要排除已被两段式匹配的情况（通过 negative lookbehind 或事后去重）
	onePartPatterns := []string{
		// 双引号格式：FROM "table"（不跟 .）
		`(?i)(?:FROM|JOIN)\s+"([a-zA-Z0-9_-]+)"(?:\s|$|,)`,
		// 反引号格式：FROM `table`（不跟 .）
		"(?i)(?:FROM|JOIN)\\s+`([a-zA-Z0-9_-]+)`(?:\\s|$|,)",
		// 无引号格式：FROM table（不跟 .）
		`(?i)(?:FROM|JOIN)\s+([a-zA-Z0-9_][a-zA-Z0-9_-]*)(?:\s|$|,)`,
	}

	for _, pattern := range onePartPatterns {
		re := regexp.MustCompile(pattern)
		matches := re.FindAllStringSubmatch(sql, -1)

		for _, match := range matches {
			if len(match) == 2 {
				table := match[1]

				// 跳过 SQL 关键字
				upper := strings.ToUpper(table)
				if upper == "SELECT" || upper == "WHERE" || upper == "SET" ||
					upper == "ON" || upper == "AS" || upper == "INTO" ||
					upper == "VALUES" || upper == "GROUP" || upper == "ORDER" ||
					upper == "HAVING" || upper == "LIMIT" || upper == "UNION" ||
					upper == "LATERAL" || upper == "UNNEST" || upper == "TABLE" {
					continue
				}

				// 去重：用 "." + table 作为 key（database 为空）
				key := "." + table
				if !seen[key] {
					// 也检查是否已被两段式匹配过（任意 database 前缀）
					alreadyMatched := false
					for k := range seen {
						if strings.HasSuffix(k, "."+table) {
							alreadyMatched = true
							break
						}
					}
					if !alreadyMatched {
						refs = append(refs, TableReference{
							Database: "",
							Table:    table,
						})
						seen[key] = true
					}
				}
			}
		}
	}

	return refs
}

// ExtractPartitionFromDataSource 从 data_source 路径提取分区信息
// 例如：s3://bucket/path/date=20251231/ -> date=20251231
func ExtractPartitionFromDataSource(dataSource string) map[string]string {
	partitions := make(map[string]string)

	// 匹配 key=value 格式的分区
	re := regexp.MustCompile(`([a-zA-Z0-9_]+)=([a-zA-Z0-9_-]+)`)
	matches := re.FindAllStringSubmatch(dataSource, -1)

	for _, match := range matches {
		if len(match) == 3 {
			partitions[match[1]] = match[2]
		}
	}

	return partitions
}

// ExtractPartitionValuesFromSQL 从 SQL WHERE 条件中提取指定分区字段的值
// 例如：SQL = "WHERE \"date\" IN ('20251230', '20251231')", partitionColumn = "date"
// 返回：["20251230", "20251231"]
func ExtractPartitionValuesFromSQL(sql string, partitionColumn string) []string {
	var values []string
	seen := make(map[string]bool)

	// 模式1: WHERE "date" = '20251231'
	pattern1 := regexp.MustCompile(`(?i)` + regexp.QuoteMeta(`"`+partitionColumn+`"`) + `\s*=\s*'([^']+)'`)
	matches1 := pattern1.FindAllStringSubmatch(sql, -1)
	for _, match := range matches1 {
		if len(match) == 2 && !seen[match[1]] {
			values = append(values, match[1])
			seen[match[1]] = true
		}
	}

	// 模式2: WHERE "date" IN ('20251230', '20251231', ...)
	pattern2 := regexp.MustCompile(`(?i)` + regexp.QuoteMeta(`"`+partitionColumn+`"`) + `\s+IN\s*\(\s*([^)]+)\s*\)`)
	matches2 := pattern2.FindAllStringSubmatch(sql, -1)
	for _, match := range matches2 {
		if len(match) == 2 {
			// 提取括号内的所有值
			inValues := match[1]
			// 匹配所有单引号内的值
			valuePattern := regexp.MustCompile(`'([^']+)'`)
			valueMatches := valuePattern.FindAllStringSubmatch(inValues, -1)
			for _, vm := range valueMatches {
				if len(vm) == 2 && !seen[vm[1]] {
					values = append(values, vm[1])
					seen[vm[1]] = true
				}
			}
		}
	}

	// 模式3: 支持无引号的列名 (date = '...' 或 date IN (...))
	pattern3 := regexp.MustCompile(`(?i)\b` + partitionColumn + `\s*=\s*'([^']+)'`)
	matches3 := pattern3.FindAllStringSubmatch(sql, -1)
	for _, match := range matches3 {
		if len(match) == 2 && !seen[match[1]] {
			values = append(values, match[1])
			seen[match[1]] = true
		}
	}

	pattern4 := regexp.MustCompile(`(?i)\b` + partitionColumn + `\s+IN\s*\(\s*([^)]+)\s*\)`)
	matches4 := pattern4.FindAllStringSubmatch(sql, -1)
	for _, match := range matches4 {
		if len(match) == 2 {
			inValues := match[1]
			valuePattern := regexp.MustCompile(`'([^']+)'`)
			valueMatches := valuePattern.FindAllStringSubmatch(inValues, -1)
			for _, vm := range valueMatches {
				if len(vm) == 2 && !seen[vm[1]] {
					values = append(values, vm[1])
					seen[vm[1]] = true
				}
			}
		}
	}

	return values
}

// ValidatePartitionColumnsInSQL 验证 SQL 是否包含所有指定的分区字段的过滤条件
// 返回缺少的分区字段列表
func ValidatePartitionColumnsInSQL(sql string, partitionColumns []string) []string {
	var missingColumns []string

	for _, col := range partitionColumns {
		values := ExtractPartitionValuesFromSQL(sql, col)
		if len(values) == 0 {
			missingColumns = append(missingColumns, col)
		}
	}

	return missingColumns
}

// IsStandardTable 判断是否为标准表
// 标准表：slow_query_logs, statements, topsql, topsql_tidb, topsql_tikv, topsql_plan_meta, topsql_sql_meta
// 只支持精确匹配，不支持带 clusterId 后缀的形式
func IsStandardTable(tableName string) bool {
	switch tableName {
	case "slow_query_logs", "statements", "topsql", "topru", "topsql_tidb", "topsql_tikv", "topsql_plan_meta", "topsql_sql_meta", "topsql_topru", "log_stats":
		return true
	}
	return false
}

func UnsupportedStandardTableError(displayName string) string {
	return fmt.Sprintf("unsupported table: %s, this API only supports standard tables (%s)", displayName, SupportedStandardTables)
}

// InferQueryType 从表名推断查询类型
// 返回 SLOW_QUERY、STATEMENTS、TOP_SQL、LOG_STATS 或空字符串
func InferQueryType(tableName string) string {
	switch tableName {
	case "slow_query_logs":
		return "SLOW_QUERY"
	case "statements":
		return "STATEMENTS"
	case "topsql", "topru", "topsql_tidb", "topsql_tikv", "topsql_plan_meta", "topsql_sql_meta", "topsql_topru":
		return "TOP_SQL"
	case "log_stats":
		return "LOG_STATS"
	}
	return ""
}

// InferCompatibleQueryType preserves the existing first-table inference for
// established query families while isolating Statements onto its own data source.
func InferCompatibleQueryType(refs []TableReference) (string, error) {
	var firstType string
	var statementsType string
	var otherType string
	for _, ref := range refs {
		inferred := InferQueryType(ref.Table)
		if inferred == "" {
			continue
		}
		if firstType == "" {
			firstType = inferred
		}
		if inferred == "STATEMENTS" {
			statementsType = inferred
		} else if otherType == "" {
			otherType = inferred
		}
	}
	if statementsType != "" && otherType != "" {
		return "", fmt.Errorf("statements cannot be queried with another data-source family: %s", otherType)
	}
	return firstType, nil
}

// ValidateStatementsTableReferences prevents Statements queries from selecting
// a same-named physical table in a caller-controlled database.
func ValidateStatementsTableReferences(refs []TableReference, clusterID string) error {
	for _, ref := range refs {
		if !IsStatementsTableReference(ref.Table, clusterID) {
			continue
		}
		if ref.Table != "statements" {
			return fmt.Errorf("physical Statements table names are internal and cannot be queried directly")
		}
		database := strings.TrimSpace(ref.Database)
		if database != "clinic_data_proxy" && database != "clinic-data-proxy" {
			return fmt.Errorf("statements must use the clinic_data_proxy database, got %q", ref.Database)
		}
	}
	return nil
}

// IsStatementsTableReference treats every statements_* name as sensitive so
// generic endpoints cannot bypass the logical Statements authorization path.
func IsStatementsTableReference(table, clusterID string) bool {
	table = strings.ToLower(strings.TrimSpace(table))
	return table == "statements" || strings.HasPrefix(table, "statements_") ||
		LogicalStandardTableName(table, clusterID) == "statements"
}

// ExtractStandardTables 从表引用列表中提取标准表
func ExtractStandardTables(refs []TableReference) []TableReference {
	var standardTables []TableReference
	for _, ref := range refs {
		if IsStandardTable(ref.Table) {
			standardTables = append(standardTables, ref)
		}
	}
	return standardTables
}