opsxsql/backend/internal/handler/connection_tester.go

package handler

import (
	"context"
	"database/sql"
	"fmt"
	"strconv"
	"strings"
	"time"

	_ "github.com/go-sql-driver/mysql"
	_ "github.com/jackc/pgx/v5/stdlib"
	"github.com/redis/go-redis/v9"
	"sql-platform/internal/model"
)

// ConnectionResult holds test connection results
type ConnectionResult struct {
	Success     bool   `json:"success"`
	Message     string `json:"message"`
	Version     string `json:"version,omitempty"`
	ClusterMode bool   `json:"cluster_mode"`
}

// testDBConnection tests a real database connection based on db_type.
func testDBConnection(instance *model.Instance) (bool, string) {
	password, err := decryptPassword(instance.Password)
	if err != nil {
		return false, fmt.Sprintf("Password decryption failed: %v", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	switch instance.DBType {
	case "mysql":
		return testMySQL(ctx, instance.Host, instance.Port, instance.Username, password)
	case "postgresql", "postgres", "pg":
		dbname := ""
		if instance.Database != nil {
			dbname = *instance.Database
		}
		return testPostgreSQL(ctx, instance.Host, instance.Port, instance.Username, password, dbname)
	case "redis":
		ok, msg, _ := testRedis(ctx, instance.Host, instance.Port, password)
		return ok, msg
	default:
		return false, fmt.Sprintf("Unsupported database type: %s", instance.DBType)
	}
}

// testDBConnectionFull returns detailed result including version
func testDBConnectionFull(instance *model.Instance) *ConnectionResult {
	password, err := decryptPassword(instance.Password)
	if err != nil {
		return &ConnectionResult{Success: false, Message: fmt.Sprintf("Password decryption failed: %v", err)}
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	switch instance.DBType {
	case "mysql":
		ok, msg := testMySQL(ctx, instance.Host, instance.Port, instance.Username, password)
		return &ConnectionResult{Success: ok, Message: msg, Version: extractVersion(msg, "version: ")}
	case "postgresql", "postgres", "pg":
		dbname := ""
		if instance.Database != nil {
			dbname = *instance.Database
		}
		ok, msg := testPostgreSQL(ctx, instance.Host, instance.Port, instance.Username, password, dbname)
		return &ConnectionResult{Success: ok, Message: msg, Version: extractVersion(msg, "version: ")}
	case "redis":
		ok, msg, isCluster := testRedis(ctx, instance.Host, instance.Port, password)
		return &ConnectionResult{Success: ok, Message: msg, Version: extractVersion(msg, "version: "), ClusterMode: isCluster}
	default:
		return &ConnectionResult{Success: false, Message: fmt.Sprintf("Unsupported database type: %s", instance.DBType)}
	}
}

func extractVersion(msg, prefix string) string {
	idx := strings.Index(msg, prefix)
	if idx < 0 {
		return ""
	}
	ver := msg[idx+len(prefix):]
	// Cut at first comma, closing paren, or space (whichever comes first)
	for _, sep := range []string{")", ",", " "} {
		if end := strings.Index(ver, sep); end > 0 {
			ver = ver[:end]
			break
		}
	}
	return strings.TrimSpace(ver)
}

func testMySQL(ctx context.Context, host string, port int, user, pass string) (bool, string) {
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/?timeout=10s&parseTime=true", user, pass, host, port)
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return false, fmt.Sprintf("Connection failed: %v", err)
	}
	defer db.Close()

	if err := db.PingContext(ctx); err != nil {
		return false, fmt.Sprintf("Connection failed: %v", err)
	}

	var ver string
	db.QueryRowContext(ctx, "SELECT VERSION()").Scan(&ver)
	return true, fmt.Sprintf("MySQL connection successful (version: %s)", ver)
}

func testPostgreSQL(ctx context.Context, host string, port int, user, pass, dbname string) (bool, string) {
	if dbname == "" {
		dbname = "postgres"
	}
	dsn := fmt.Sprintf("postgres://%s:***@%s:%d/%s?sslmode=disable&connect_timeout=10",
		user, pass, host, port, dbname)
	db, err := sql.Open("pgx", dsn)
	if err != nil {
		return false, fmt.Sprintf("Connection failed: %v", err)
	}
	defer db.Close()

	if err := db.PingContext(ctx); err != nil {
		return false, fmt.Sprintf("Connection failed: %v", err)
	}

	var ver string
	db.QueryRowContext(ctx, "SELECT version()").Scan(&ver)
	return true, fmt.Sprintf("PostgreSQL connection successful (version: %s)", ver)
}

func testRedis(ctx context.Context, host string, port int, pass string) (bool, string, bool) {
	addr := fmt.Sprintf("%s:%d", host, port)
	rdb := redis.NewClient(&redis.Options{
		Addr:         addr,
		Password:     pass,
		DB:           0,
		DialTimeout:  10 * time.Second,
		ReadTimeout:  10 * time.Second,
	})
	defer rdb.Close()

	pong, err := rdb.Ping(ctx).Result()
	if err != nil {
		return false, fmt.Sprintf("Connection failed: %v", err), false
	}

	var version string
	var isCluster bool
	if info, err := rdb.Info(ctx, "server", "cluster").Result(); err == nil {
		for _, line := range strings.Split(info, "\n") {
			line = strings.TrimRight(line, "\r")
			if strings.HasPrefix(line, "redis_version:") {
				version = strings.TrimPrefix(line, "redis_version:")
			}
			if strings.HasPrefix(line, "cluster_enabled:") {
				val := strings.TrimPrefix(line, "cluster_enabled:")
				isCluster = strings.TrimSpace(val) == "1"
			}
		}
	}

	mode := "standalone"
	if isCluster {
		mode = "cluster"
	}

	if version != "" {
		return true, fmt.Sprintf("Redis connection successful (PING=%s, version: %s, mode: %s)", pong, version, mode), isCluster
	}
	return true, fmt.Sprintf("Redis connection successful (PING=%s, mode: %s)", pong, mode), isCluster
}

// getRedisMetrics fetches live Redis metrics for instance overview
func getRedisMetrics(instance *model.Instance) map[string]interface{} {
	password := instance.Password
	if strings.HasPrefix(password, "enc:") {
		decrypted, err := decryptPassword(strings.TrimPrefix(password, "enc:"))
		if err == nil {
			password = decrypted
		}
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	addr := fmt.Sprintf("%s:%d", instance.Host, instance.Port)
	isCluster := instance.ClusterMode != nil && *instance.ClusterMode

	metrics := map[string]interface{}{
		"connections":  0,
		"qps":          0,
		"latency":      0,
		"uptime":       "-",
		"version":      "",
		"cluster_mode": isCluster,
	}

	var info string
	if isCluster {
		rdb := redis.NewClusterClient(&redis.ClusterOptions{
			Addrs:       []string{addr},
			Password:    password,
			DialTimeout: 5 * time.Second,
			ReadTimeout: 5 * time.Second,
		})
		defer rdb.Close()

		_, err := rdb.Ping(ctx).Result()
		if err != nil {
			metrics["status"] = "offline"
			return metrics
		}
		metrics["status"] = "online"
		info, _ = rdb.Info(ctx, "server", "clients", "stats").Result()
		// Latency from PING
		start := time.Now()
		rdb.Ping(ctx)
		metrics["latency"] = time.Since(start).Milliseconds()
	} else {
		rdb := redis.NewClient(&redis.Options{
			Addr:         addr,
			Password:     password,
			DB:           0,
			DialTimeout:  5 * time.Second,
			ReadTimeout:  5 * time.Second,
		})
		defer rdb.Close()

		_, err := rdb.Ping(ctx).Result()
		if err != nil {
			metrics["status"] = "offline"
			return metrics
		}
		metrics["status"] = "online"
		info, _ = rdb.Info(ctx, "server", "clients", "stats").Result()
		// Latency from PING
		start := time.Now()
		rdb.Ping(ctx)
		metrics["latency"] = time.Since(start).Milliseconds()
	}

	if info == "" {
		return metrics
	}

	infoMap := make(map[string]string)
	for _, line := range strings.Split(info, "\n") {
		line = strings.TrimSpace(strings.TrimRight(line, "\r"))
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		parts := strings.SplitN(line, ":", 2)
		if len(parts) == 2 {
			infoMap[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1])
		}
	}

	if v, ok := infoMap["redis_version"]; ok {
		metrics["version"] = v
	}
	if v, ok := infoMap["connected_clients"]; ok {
		metrics["connections"] = v
	}
	if v, ok := infoMap["instantaneous_ops_per_sec"]; ok {
		metrics["qps"] = v
	}
	if v, ok := infoMap["uptime_in_days"]; ok {
		metrics["uptime"] = v + " 天"
	} else if v, ok := infoMap["uptime_in_seconds"]; ok {
		secs := 0
		fmt.Sscanf(v, "%d", &secs)
		if secs > 86400 {
			metrics["uptime"] = fmt.Sprintf("%d 天", secs/86400)
		} else if secs > 3600 {
			metrics["uptime"] = fmt.Sprintf("%d 小时", secs/3600)
		} else {
			metrics["uptime"] = fmt.Sprintf("%d 分钟", secs/60)
		}
	}

	return metrics
}

// getMySQLMetrics fetches live MySQL metrics (connections, QPS, latency, uptime, version)
func getMySQLMetrics(instance *model.Instance) map[string]interface{} {
	password := instance.Password
	if strings.HasPrefix(password, "enc:") {
		decrypted, err := decryptPassword(strings.TrimPrefix(password, "enc:"))
		if err == nil {
			password = decrypted
		}
	}

	metrics := map[string]interface{}{
		"status":      "offline",
		"connections": 0,
		"qps":         0,
		"latency":     0,
		"uptime":      "-",
		"version":     "",
	}
	if instance.Version != nil {
		metrics["version"] = *instance.Version
	}

	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/?timeout=5s&parseTime=true",
		instance.Username, password, instance.Host, instance.Port)
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return metrics
	}
	defer db.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// Latency from simple query
	start := time.Now()
	if err := db.PingContext(ctx); err != nil {
		return metrics
	}
	metrics["latency"] = time.Since(start).Milliseconds()
	metrics["status"] = "online"

	// Fetch STATUS variables
	statusVars := map[string]string{}
	rows, err := db.QueryContext(ctx, "SHOW GLOBAL STATUS WHERE Variable_name IN "+
		"('Threads_connected','Queries','Uptime','Connections','Slow_queries')")
	if err == nil {
		defer rows.Close()
		for rows.Next() {
			var name, value string
			rows.Scan(&name, &value)
			statusVars[name] = value
		}
	}

	if v, ok := statusVars["Threads_connected"]; ok {
		n, _ := strconv.Atoi(v)
		metrics["connections"] = n
	}
	if v, ok := statusVars["Uptime"]; ok {
		secs, _ := strconv.Atoi(v)
		if secs > 86400 {
			metrics["uptime"] = fmt.Sprintf("%d 天", secs/86400)
		} else if secs > 3600 {
			metrics["uptime"] = fmt.Sprintf("%d 小时", secs/3600)
		} else {
			metrics["uptime"] = fmt.Sprintf("%d 分钟", secs/60)
		}
	}
	// QPS = Queries / Uptime
	if q, ok := statusVars["Queries"]; ok {
		if u, ok2 := statusVars["Uptime"]; ok2 {
			qi, _ := strconv.ParseFloat(q, 64)
			ui, _ := strconv.ParseFloat(u, 64)
			if ui > 0 {
				metrics["qps"] = int(qi / ui)
			}
		}
	}

	// Version from DB if not already set
	if instance.Version == nil || *instance.Version == "" {
		var ver string
		if err := db.QueryRowContext(ctx, "SELECT VERSION()").Scan(&ver); err == nil {
			metrics["version"] = ver
		}
	}

	return metrics
}

// getPostgresMetrics fetches live PostgreSQL metrics
func getPostgresMetrics(instance *model.Instance) map[string]interface{} {
	password := instance.Password
	if strings.HasPrefix(password, "enc:") {
		decrypted, err := decryptPassword(strings.TrimPrefix(password, "enc:"))
		if err == nil {
			password = decrypted
		}
	}

	metrics := map[string]interface{}{
		"status":      "offline",
		"connections": 0,
		"qps":         0,
		"latency":     0,
		"uptime":      "-",
		"version":     "",
	}
	if instance.Version != nil {
		metrics["version"] = *instance.Version
	}

	dsn := fmt.Sprintf("postgres://%s:%s@%s:%d/postgres?sslmode=disable&connect_timeout=5",
		instance.Username, password, instance.Host, instance.Port)
	db, err := sql.Open("pgx", dsn)
	if err != nil {
		return metrics
	}
	defer db.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	start := time.Now()
	if err := db.PingContext(ctx); err != nil {
		return metrics
	}
	metrics["latency"] = time.Since(start).Milliseconds()
	metrics["status"] = "online"

	// Fetch pg_stat_activity counts
	var activeConns int
	db.QueryRowContext(ctx, "SELECT count(*) FROM pg_stat_activity WHERE state = 'active'").Scan(&activeConns)
	metrics["connections"] = activeConns

	// Uptime
	var uptimeSecs int
	if err := db.QueryRowContext(ctx, "SELECT extract(epoch from now() - pg_postmaster_start_time())::int").Scan(&uptimeSecs); err == nil {
		if uptimeSecs > 86400 {
			metrics["uptime"] = fmt.Sprintf("%d 天", uptimeSecs/86400)
		} else if uptimeSecs > 3600 {
			metrics["uptime"] = fmt.Sprintf("%d 小时", uptimeSecs/3600)
		} else {
			metrics["uptime"] = fmt.Sprintf("%d 分钟", uptimeSecs/60)
		}
	}

	// Version
	if instance.Version == nil || *instance.Version == "" {
		var ver string
		if err := db.QueryRowContext(ctx, "SELECT version()").Scan(&ver); err == nil {
			metrics["version"] = ver
		}
	}

	return metrics
}

func getDBConnection(instance *model.Instance, dbName string) (*sql.DB, error) {
	password, err := decryptPassword(instance.Password)
	if err != nil {
		return nil, fmt.Errorf("password decryption failed: %v", err)
	}

	switch instance.DBType {
	case "mysql":
		dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?timeout=10s&parseTime=true",
			instance.Username, password, instance.Host, instance.Port, dbName)
		return sql.Open("mysql", dsn)
	case "postgresql", "postgres", "pg":
		if dbName == "" {
			dbName = "postgres"
		}
		dsn := fmt.Sprintf("postgres://%s:***@%s:%d/%s?sslmode=disable&connect_timeout=10",
			instance.Username, password, instance.Host, instance.Port, dbName)
		return sql.Open("pgx", dsn)
	default:
		return nil, fmt.Errorf("unsupported database type: %s", instance.DBType)
	}
}

// listDatabases queries the actual databases from the instance
func listDatabases(instance *model.Instance) ([]string, error) {
	db, err := getDBConnection(instance, "")
	if err != nil {
		return nil, err
	}
	defer db.Close()

	switch instance.DBType {
	case "mysql":
		rows, err := db.Query("SHOW DATABASES")
		if err != nil {
			return nil, err
		}
		defer rows.Close()
		var databases []string
		for rows.Next() {
			var name string
			rows.Scan(&name)
			if name != "information_schema" && name != "performance_schema" &&
				name != "mysql" && name != "sys" {
				databases = append(databases, name)
			}
		}
		return databases, nil

	case "postgresql", "postgres", "pg":
		dbConn, err := getDBConnection(instance, "postgres")
		if err != nil {
			return nil, err
		}
		defer dbConn.Close()
		rows, err := dbConn.Query("SELECT datname FROM pg_database WHERE datistemplate = false")
		if err != nil {
			return nil, err
		}
		defer rows.Close()
		var databases []string
		for rows.Next() {
			var name string
			rows.Scan(&name)
			databases = append(databases, name)
		}
		return databases, nil

	default:
		return nil, fmt.Errorf("unsupported type: %s", instance.DBType)
	}
}

// listTables queries the actual tables from a database
func listTables(instance *model.Instance, dbName string) ([]map[string]interface{}, error) {
	db, err := getDBConnection(instance, dbName)
	if err != nil {
		return nil, err
	}
	defer db.Close()

	switch instance.DBType {
	case "mysql":
		rows, err := db.Query("SELECT TABLE_NAME, TABLE_TYPE, TABLE_ROWS, DATA_LENGTH FROM information_schema.TABLES WHERE TABLE_SCHEMA = ? ORDER BY TABLE_NAME", dbName)
		if err != nil {
			return nil, err
		}
		defer rows.Close()
		var tables []map[string]interface{}
		for rows.Next() {
			var name, tblType string
			var rows_ int
			var dataLen int64
			rows.Scan(&name, &tblType, &rows_, &dataLen)
			tables = append(tables, map[string]interface{}{
				"name": name, "type": tblType, "rows": rows_, "data_length": dataLen,
			})
		}
		return tables, nil

	case "postgresql", "postgres", "pg":
		rows, err := db.Query("SELECT tablename FROM pg_tables WHERE schemaname = 'public' ORDER BY tablename")
		if err != nil {
			return nil, err
		}
		defer rows.Close()
		var tables []map[string]interface{}
		for rows.Next() {
			var name string
			rows.Scan(&name)
			tables = append(tables, map[string]interface{}{
				"name": name, "type": "BASE TABLE",
			})
		}
		return tables, nil

	default:
		return nil, fmt.Errorf("unsupported type: %s", instance.DBType)
	}
}